From 7805496f0d5ab5cbfceee502a1c602e18adee4a3 Mon Sep 17 00:00:00 2001
From: Colton Willey <colton@wolfssl.com>
Date: Wed, 4 Sep 2024 11:24:10 -0700
Subject: [PATCH] Update PKCS11 examples to set RNG for RSA operations

---
 pkcs11/pkcs11_rsa.c  | 16 +++++++++-------
 pkcs11/pkcs11_test.c | 32 ++++++++++++++++++++++++--------
 2 files changed, 33 insertions(+), 15 deletions(-)

diff --git a/pkcs11/pkcs11_rsa.c b/pkcs11/pkcs11_rsa.c
index 4c1f8cbc..6074b91c 100644
--- a/pkcs11/pkcs11_rsa.c
+++ b/pkcs11/pkcs11_rsa.c
@@ -187,6 +187,7 @@ static const unsigned char client_keypub_der_2048[] =
         0x03, 0x01, 0x00, 0x01
 };
 static const int sizeof_client_keypub_der_2048 = sizeof(client_keypub_der_2048);
+WC_RNG rng;
 
 static int decode_private_key(RsaKey* key, int devId)
 {
@@ -245,7 +246,7 @@ static int rsa_sign_verify(int devId)
     if (ret == 0) {
         fprintf(stderr, "Signing\n");
         sigSz = ret = wc_RsaSSL_Sign(hash, hashSz, sig, (int)sigSz, &priv,
-                                     NULL);
+                                     &rng);
         if (ret < 0)
             fprintf(stderr, "Failed to sign: %d\n", ret);
         else
@@ -282,18 +283,12 @@ static int rsa_sign_verify_pss(int devId)
     word32 hashSz, ptSz, sigSz;
     RsaKey priv;
     RsaKey pub;
-    WC_RNG rng;
 
     memset(hash, 9, sizeof(hash));
     hashSz = sizeof(hash);
     sigSz = sizeof(sig);
     ptSz = sizeof(pt);
 
-    ret = wc_InitRng(&rng);
-    if (ret != 0) {
-        fprintf(stderr, "Failed to initialize RNG: %d\n", ret);
-    }
-
     ret = decode_private_key(&priv, devId);
     if (ret == 0) {
         fprintf(stderr, "PSS Signing\n");
@@ -375,6 +370,12 @@ int main(int argc, char* argv[])
                 fprintf(stderr, "Failed to register PKCS#11 token\n");
                 ret = 2;
             }
+            if (ret == 0) {
+                ret = wc_InitRng(&rng);
+                if (ret != 0) {
+                    fprintf(stderr, "Failed to initialize RNG: %d\n", ret);
+                }
+            }
             if (ret == 0) {
             #ifndef NO_RSA
                 ret = rsa_sign_verify(devId);
@@ -394,6 +395,7 @@ int main(int argc, char* argv[])
         wc_Pkcs11_Finalize(&dev);
     }
 
+    wc_FreeRng(&rng);
     wolfCrypt_Cleanup();
 
     if (ret == 0)
diff --git a/pkcs11/pkcs11_test.c b/pkcs11/pkcs11_test.c
index 0e656cd1..2e87a24d 100644
--- a/pkcs11/pkcs11_test.c
+++ b/pkcs11/pkcs11_test.c
@@ -429,6 +429,10 @@ int rsaenc_test(RsaKey* key)
     outSz = sizeof(out);
     decSz = sizeof(dec);
 
+#ifdef WC_RSA_BLINDING
+    ret = wc_RsaSetRNG(key, &rng);
+#endif
+
     if (ret == 0) {
         outSz = ret = wc_RsaPublicEncrypt_ex(plain, plainSz, out, (int)outSz,
             key, &rng, WC_RSA_PKCSV15_PAD, WC_HASH_TYPE_NONE, WC_MGF1NONE, NULL,
@@ -460,27 +464,39 @@ int rsaenc_test(RsaKey* key)
 int rsasig_test(RsaKey* key)
 {
     int    ret = 0;
-    byte   plain[128], out[2048/8];
-    word32 plainSz, outSz;
+    byte   plain[128], sig[2048/8], pt[2048/8];
+    word32 plainSz, sigSz, ptSz;
 
     memset(plain, 9, sizeof(plain));
     plainSz = sizeof(plain);
-    outSz = sizeof(out);
+    sigSz = sizeof(sig);
+    ptSz = sizeof(pt);
 
     if (ret == 0) {
-        outSz = ret = wc_RsaSSL_Sign(plain, plainSz, out, (int)outSz, key,
-            NULL);
+        sigSz = ret = wc_RsaSSL_Sign(plain, plainSz, sig, (int)sigSz, key,
+            &rng);
         if (ret < 0)
             fprintf(stderr, "Failed to sign: %d\n", ret);
         else
             ret = 0;
     }
     if (ret == 0) {
-        ret = wc_RsaSSL_Verify(out, outSz, plain, (int)plainSz, key);
+        ret = wc_RsaSSL_Verify(sig, sigSz, pt, (int)ptSz, key);
         if (ret < 0)
             fprintf(stderr, "Failed to verify: %d\n", ret);
-        else
-            ret = 0;
+
+        if (ret != plainSz) {
+            fprintf(stderr, "Failed to verify: %d\n", ret);
+            ret = -1;
+        }
+        if (ret > 0) {
+            if (XMEMCMP(plain, pt, ret) != 0) {
+                ret = -1;
+            }
+            else {
+                ret = 0;
+            }
+        }
     }
 
     return ret;