From 48975c38e02e0ccf57e2d9dcee6999442b59c1d5 Mon Sep 17 00:00:00 2001
From: Juan Lang <juan.lang@gmail.com>
Date: Tue, 20 Oct 2009 15:22:21 -0700
Subject: [PATCH] rsaenh: Store key when algid is a CALG_RSA_* algid, too.

---
 dlls/rsaenh/rsaenh.c       | 16 ++++++----------
 dlls/rsaenh/tests/rsaenh.c | 31 +++++++++++++++++++++++++++++++
 2 files changed, 37 insertions(+), 10 deletions(-)

diff --git a/dlls/rsaenh/rsaenh.c b/dlls/rsaenh/rsaenh.c
index de3125cc6a8..1b23391affe 100644
--- a/dlls/rsaenh/rsaenh.c
+++ b/dlls/rsaenh/rsaenh.c
@@ -3059,11 +3059,9 @@ BOOL WINAPI RSAENH_CPGenKey(HCRYPTPROV hProv, ALG_ID Algid, DWORD dwFlags, HCRYP
             if (pCryptKey) { 
                 new_key_impl(pCryptKey->aiAlgid, &pCryptKey->context, pCryptKey->dwKeyLen);
                 setup_key(pCryptKey);
-                if (Algid == AT_SIGNATURE) {
-                    RSAENH_CPDestroyKey(hProv, pKeyContainer->hSignatureKeyPair);
-                    copy_handle(&handle_table, *phKey, RSAENH_MAGIC_KEY,
-                                &pKeyContainer->hSignatureKeyPair);
-                }
+                RSAENH_CPDestroyKey(hProv, pKeyContainer->hSignatureKeyPair);
+                copy_handle(&handle_table, *phKey, RSAENH_MAGIC_KEY,
+                            &pKeyContainer->hSignatureKeyPair);
             }
             break;
 
@@ -3073,11 +3071,9 @@ BOOL WINAPI RSAENH_CPGenKey(HCRYPTPROV hProv, ALG_ID Algid, DWORD dwFlags, HCRYP
             if (pCryptKey) { 
                 new_key_impl(pCryptKey->aiAlgid, &pCryptKey->context, pCryptKey->dwKeyLen);
                 setup_key(pCryptKey);
-                if (Algid == AT_KEYEXCHANGE) {
-                    RSAENH_CPDestroyKey(hProv, pKeyContainer->hKeyExchangeKeyPair);
-                    copy_handle(&handle_table, *phKey, RSAENH_MAGIC_KEY,
-                                &pKeyContainer->hKeyExchangeKeyPair);
-                }
+                RSAENH_CPDestroyKey(hProv, pKeyContainer->hKeyExchangeKeyPair);
+                copy_handle(&handle_table, *phKey, RSAENH_MAGIC_KEY,
+                            &pKeyContainer->hKeyExchangeKeyPair);
             }
             break;
             
diff --git a/dlls/rsaenh/tests/rsaenh.c b/dlls/rsaenh/tests/rsaenh.c
index 8afa931b8e4..bfe2b5a6a54 100644
--- a/dlls/rsaenh/tests/rsaenh.c
+++ b/dlls/rsaenh/tests/rsaenh.c
@@ -2086,6 +2086,37 @@ static void test_null_provider(void)
     CryptAcquireContext(&prov, szContainer, NULL, PROV_RSA_FULL,
      CRYPT_DELETEKEYSET);
 
+    /* Test for being able to get a key generated with CALG_RSA_SIGN. */
+    result = CryptAcquireContext(&prov, szContainer, NULL, PROV_RSA_FULL,
+     CRYPT_NEWKEYSET);
+    ok(result, "CryptAcquireContext failed: %08x\n", GetLastError());
+    result = CryptGenKey(prov, CALG_RSA_SIGN, 0, &key);
+    ok(result, "CryptGenKey with CALG_RSA_SIGN failed with error %08x\n", GetLastError());
+    result = CryptGetUserKey(prov, AT_KEYEXCHANGE, &key);
+    ok(!result, "expected CryptGetUserKey to fail\n");
+    result = CryptGetUserKey(prov, AT_SIGNATURE, &key);
+    ok(result, "CryptGetUserKey with AT_SIGNATURE failed: %08x\n", GetLastError());
+    CryptDestroyKey(key);
+    CryptReleaseContext(prov, 0);
+
+    CryptAcquireContext(&prov, szContainer, NULL, PROV_RSA_FULL,
+     CRYPT_DELETEKEYSET);
+
+    /* Test for being able to get a key generated with CALG_RSA_KEYX. */
+    result = CryptAcquireContext(&prov, szContainer, NULL, PROV_RSA_FULL,
+     CRYPT_NEWKEYSET);
+    ok(result, "CryptAcquireContext failed: %08x\n", GetLastError());
+    result = CryptGenKey(prov, CALG_RSA_KEYX, 0, &key);
+    ok(result, "CryptGenKey with CALG_RSA_KEYX failed with error %08x\n", GetLastError());
+    result = CryptGetUserKey(prov, AT_KEYEXCHANGE, &key);
+    ok(result, "CryptGetUserKey with AT_KEYEXCHANGE failed: %08x\n", GetLastError());
+    result = CryptGetUserKey(prov, AT_SIGNATURE, &key);
+    ok(!result, "expected CryptGetUserKey to fail\n");
+    CryptDestroyKey(key);
+    CryptReleaseContext(prov, 0);
+
+    CryptAcquireContext(&prov, szContainer, NULL, PROV_RSA_FULL,
+     CRYPT_DELETEKEYSET);
 
     /* test for the bug in accessing the user key in a container
      */