From 1d8687e45f60bc1f8d2f3d8342e9b577792ddecc Mon Sep 17 00:00:00 2001
From: David Garske <david@wolfssl.com>
Date: Mon, 13 Nov 2023 14:15:14 -0800
Subject: [PATCH] Add TPM crypto callback support for RSA key generation.

---
 src/tpm2_cryptocb.c | 33 +++++++++++++++++++++++++--------
 src/tpm2_wrap.c     |  8 ++++----
 wolftpm/tpm2_wrap.h |  9 +++++++++
 3 files changed, 38 insertions(+), 12 deletions(-)

diff --git a/src/tpm2_cryptocb.c b/src/tpm2_cryptocb.c
index 8daad7fe..bda28d79 100644
--- a/src/tpm2_cryptocb.c
+++ b/src/tpm2_cryptocb.c
@@ -78,14 +78,31 @@ int wolfTPM2_CryptoDevCb(int devId, wc_CryptoInfo* info, void* ctx)
     #ifndef NO_RSA
         /* RSA */
         if (info->pk.type == WC_PK_TYPE_RSA_KEYGEN) {
-            /* TODO: Add crypto callback RSA keygen support */
-            #if 0
-            RsaKey* key;
-            int     size;
-            long    e;
-            WC_RNG* rng;
-            #endif
-            rc = exit_rc;
+        #ifdef WOLFSSL_KEY_GEN
+            if (tlsCtx->rsaKey != NULL && tlsCtx->storageKey != NULL) {
+                /* create a new RSA key */
+                TPMT_PUBLIC publicTemplate;
+                XMEMSET(&publicTemplate, 0, sizeof(publicTemplate));
+                rc = GetKeyTemplateRSA(&publicTemplate, WOLFTPM2_WRAP_DIGEST,
+                    (TPMA_OBJECT_sensitiveDataOrigin | TPMA_OBJECT_userWithAuth |
+                     TPMA_OBJECT_decrypt | TPMA_OBJECT_noDA),
+                    info->pk.rsakg.size,
+                    info->pk.rsakg.e,
+                    TPM_ALG_NULL /* TPM_ALG_RSASSA */,
+                    WOLFTPM2_WRAP_DIGEST);
+                if (rc == 0) {
+                    rc = wolfTPM2_CreateAndLoadKey(tlsCtx->dev, tlsCtx->rsaKey,
+                        &tlsCtx->storageKey->handle, &publicTemplate, NULL, 0);
+                }
+                if (rc == 0) {
+                    /* export public portion of new key to wolf RsaKey struct */
+                    rc = wolfTPM2_RsaKey_TpmToWolf(tlsCtx->dev, tlsCtx->rsaKey,
+                        info->pk.rsakg.key);
+                }
+            }
+            else
+        #endif
+                rc = exit_rc;
         }
         else if (info->pk.type == WC_PK_TYPE_RSA) {
             switch (info->pk.rsa.type) {
diff --git a/src/tpm2_wrap.c b/src/tpm2_wrap.c
index 2f4f464b..c19dfede 100644
--- a/src/tpm2_wrap.c
+++ b/src/tpm2_wrap.c
@@ -5280,8 +5280,8 @@ int wolfTPM2_ChangePlatformAuth(WOLFTPM2_DEV* dev, WOLFTPM2_SESSION* session)
 /* --- BEGIN Utility Functions -- */
 /******************************************************************************/
 
-static int GetKeyTemplateRSA(TPMT_PUBLIC* publicTemplate,
-    TPM_ALG_ID nameAlg, TPMA_OBJECT objectAttributes, int keyBits, int exponent,
+int GetKeyTemplateRSA(TPMT_PUBLIC* publicTemplate,
+    TPM_ALG_ID nameAlg, TPMA_OBJECT objectAttributes, int keyBits, long exponent,
     TPM_ALG_ID sigScheme, TPM_ALG_ID sigHash)
 {
     if (publicTemplate == NULL)
@@ -5293,7 +5293,7 @@ static int GetKeyTemplateRSA(TPMT_PUBLIC* publicTemplate,
     publicTemplate->nameAlg = nameAlg;
     publicTemplate->objectAttributes = objectAttributes;
     publicTemplate->parameters.rsaDetail.keyBits = keyBits;
-    publicTemplate->parameters.rsaDetail.exponent = exponent;
+    publicTemplate->parameters.rsaDetail.exponent = (UINT32)exponent;
     publicTemplate->parameters.rsaDetail.scheme.scheme = sigScheme;
     publicTemplate->parameters.rsaDetail.scheme.details.anySig.hashAlg = sigHash;
     /* For fixedParent or (decrypt and restricted) enable symmetric */
@@ -5311,7 +5311,7 @@ static int GetKeyTemplateRSA(TPMT_PUBLIC* publicTemplate,
     return TPM_RC_SUCCESS;
 }
 
-static int GetKeyTemplateECC(TPMT_PUBLIC* publicTemplate,
+int GetKeyTemplateECC(TPMT_PUBLIC* publicTemplate,
     TPM_ALG_ID nameAlg, TPMA_OBJECT objectAttributes, TPM_ECC_CURVE curve,
     TPM_ALG_ID sigScheme, TPM_ALG_ID sigHash)
 {
diff --git a/wolftpm/tpm2_wrap.h b/wolftpm/tpm2_wrap.h
index a43274f6..4394acb4 100644
--- a/wolftpm/tpm2_wrap.h
+++ b/wolftpm/tpm2_wrap.h
@@ -3387,6 +3387,15 @@ WOLFTPM_API int wolfTPM2_PolicyAuthorizeMake(TPM_ALG_ID pcrAlg,
     const TPM2B_PUBLIC* pub, byte* digest, word32* digestSz,
     const byte* policyRef, word32 policyRefSz);
 
+
+/* Internal API's */
+WOLFTPM_LOCAL int GetKeyTemplateRSA(TPMT_PUBLIC* publicTemplate,
+    TPM_ALG_ID nameAlg, TPMA_OBJECT objectAttributes, int keyBits, long exponent,
+    TPM_ALG_ID sigScheme, TPM_ALG_ID sigHash);
+WOLFTPM_LOCAL int GetKeyTemplateECC(TPMT_PUBLIC* publicTemplate,
+    TPM_ALG_ID nameAlg, TPMA_OBJECT objectAttributes, TPM_ECC_CURVE curve,
+    TPM_ALG_ID sigScheme, TPM_ALG_ID sigHash);
+
 #ifdef __cplusplus
     }  /* extern "C" */
 #endif