[PATCH 7/7] bcrypt: Check output size early in key_asymmetric_encrypt() for RSA.

[Paul Gofman]

From: Paul Gofman

---
 dlls/bcrypt/gnutls.c       |  7 +++++++
 dlls/bcrypt/tests/bcrypt.c | 33 +++++++++++++++++++++++++++++----
 2 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/dlls/bcrypt/gnutls.c b/dlls/bcrypt/gnutls.c
index c8c901f3460..78f30dc7c83 100644
--- a/dlls/bcrypt/gnutls.c
+++ b/dlls/bcrypt/gnutls.c
@@ -2777,6 +2777,13 @@ static NTSTATUS key_asymmetric_encrypt( void *args )

if (!key_data(params->key)->a.pubkey) return STATUS_INVALID_HANDLE;

+    if (params->key->alg_id == ALG_ID_RSA
+        && (!params->output || len_from_bitlen( params->key->u.a.bitlen ) > params->output_len))
+    {
+        *params->ret_len = len_from_bitlen( params->key->u.a.bitlen );
+        return !params->output ? STATUS_SUCCESS : STATUS_BUFFER_TOO_SMALL;
+    }
+
     if (params->key->alg_id == ALG_ID_RSA && params->flags & BCRYPT_PAD_OAEP)
     {
         BCRYPT_OAEP_PADDING_INFO *pad = params->padding;
diff --git a/dlls/bcrypt/tests/bcrypt.c b/dlls/bcrypt/tests/bcrypt.c
index 3ff33a1c177..9a0fdf1f6b9 100644
--- a/dlls/bcrypt/tests/bcrypt.c
+++ b/dlls/bcrypt/tests/bcrypt.c
@@ -2614,6 +2614,12 @@ static void test_rsa_encrypt(void)
     /*   No padding    */
     todo_wine {
     memset(input_no_padding, 0, sizeof(input_no_padding));
+
+    encrypted_size = 0;
+    ret = BCryptEncrypt(key, input_no_padding, sizeof(input_no_padding), NULL, NULL, 0, NULL, 0, &encrypted_size, BCRYPT_PAD_NONE);
+    ok(ret == STATUS_SUCCESS, "got %lx\n", ret);
+    ok(encrypted_size == 64, "got size of %ld\n", encrypted_size);
+
     strcpy((char *)input_no_padding, "Hello World");
     encrypted_size = 0;
     ret = BCryptEncrypt(key, input_no_padding, sizeof(input_no_padding), NULL, NULL, 0, NULL, 0, &encrypted_size, BCRYPT_PAD_NONE);
@@ -2709,26 +2715,45 @@ static void test_rsa_encrypt(void)
     ret = BCryptFinalizeKeyPair(key, 0);
     ok(ret == STATUS_SUCCESS, "got %#lx\n", ret);

-    todo_wine {
     encrypted_size = 0;
     ret = BCryptEncrypt(key, input, sizeof(input), &oaep_pad, NULL, 0, NULL, 0, &encrypted_size, BCRYPT_PAD_OAEP);
     ok(ret == STATUS_SUCCESS, "got %lx\n", ret);
     ok(encrypted_size == 80, "got size of %ld\n", encrypted_size);

+    encrypted_size = 0;
+    ret = BCryptEncrypt(key, input, sizeof(input), NULL, NULL, 0, NULL, 0, &encrypted_size, BCRYPT_PAD_OAEP);
+    ok(ret == STATUS_SUCCESS, "got %lx\n", ret);
+    ok(encrypted_size == 80, "got size of %ld\n", encrypted_size);
+
     encrypted_a = realloc(encrypted_a, encrypted_size);
     memset(encrypted_a, 0, encrypted_size);
     encrypted_b = realloc(encrypted_b, encrypted_size);
     memset(encrypted_b, 0, encrypted_size);

-    ret = BCryptEncrypt(key, input, sizeof(input), &oaep_pad, NULL, 0, encrypted_a, encrypted_size, &encrypted_size, BCRYPT_PAD_OAEP);
+    ret = BCryptEncrypt(key, input, sizeof(input), NULL, NULL, 0, NULL, encrypted_size, &encrypted_size, BCRYPT_PAD_OAEP);
     ok(ret == STATUS_SUCCESS, "got %lx\n", ret);
     ok(encrypted_size == 80, "got size of %ld\n", encrypted_size);

+    encrypted_size = 0;
+    ret = BCryptEncrypt(key, input, sizeof(input), NULL, NULL, 0, encrypted_a, 0, &encrypted_size, BCRYPT_PAD_OAEP);
+    ok(ret == STATUS_BUFFER_TOO_SMALL, "got %lx\n", ret);
+    ok(encrypted_size == 80, "got size of %ld\n", encrypted_size);
+
+    ret = BCryptEncrypt(key, input, sizeof(input), &oaep_pad, NULL, 0, encrypted_a, encrypted_size, &encrypted_size, BCRYPT_PAD_OAEP);
+    todo_wine ok(ret == STATUS_SUCCESS, "got %lx\n", ret);
+    ok(encrypted_size == 80, "got size of %ld\n", encrypted_size);
+
     ret = BCryptEncrypt(key, input, sizeof(input), &oaep_pad, NULL, 0, encrypted_b, encrypted_size, &encrypted_size, BCRYPT_PAD_OAEP);
-    ok(ret == STATUS_SUCCESS, "got %lx\n", ret);
+    todo_wine ok(ret == STATUS_SUCCESS, "got %lx\n", ret);
     ok(encrypted_size == 80, "got size of %ld\n", encrypted_size);
-    ok(memcmp(encrypted_a, encrypted_b, encrypted_size), "Both outputs are the same\n");
+    todo_wine ok(memcmp(encrypted_a, encrypted_b, encrypted_size), "Both outputs are the same\n");
+
+    decrypted_size = 0;
+    memset(decrypted, 0, sizeof(decrypted));
+    ret = BCryptDecrypt(key, encrypted_a, encrypted_size, NULL, NULL, 0, NULL, 0, &decrypted_size, BCRYPT_PAD_OAEP);
+    ok(ret == STATUS_INVALID_PARAMETER, "got %lx\n", ret);

+    todo_wine {
     decrypted_size = 0;
     memset(decrypted, 0, sizeof(decrypted));
     ret = BCryptDecrypt(key, encrypted_a, encrypted_size, &oaep_pad, NULL, 0, NULL, 0, &decrypted_size, BCRYPT_PAD_OAEP);