fix keytools public key der export to use ml-dsa level passed as env var

wolfSSL · Dec 16, 2024 · 0d18b25 · 0d18b25
1 parent 9454deb
commit 0d18b25
Showing 1 changed file with 32 additions and 13 deletions.
diff --git a/tools/keytools/keygen.c b/tools/keytools/keygen.c
@@ -1088,21 +1088,38 @@ static void keygen_ml_dsa(const char *priv_fname, uint32_t id_mask)
 
     if (exportPubKey) {
         if (saveAsDer) {
+            uint8_t*  pubDer;
+            size_t    pubDerSz;
+            int       pubOutLen;
+            const int WITH_ALG_SPKI = 1;
+
+            /* Size the buffer based on the ML DSA level */
+            switch (ml_dsa_level) {
+                case WC_ML_DSA_44:
+                    pubDerSz = ML_DSA_LEVEL2_PUB_KEY_DER_SIZE;
+                    break;
+                case WC_ML_DSA_65:
+                    pubDerSz = ML_DSA_LEVEL3_PUB_KEY_DER_SIZE;
+                    break;
+                case WC_ML_DSA_87:
+                    pubDerSz = ML_DSA_LEVEL5_PUB_KEY_DER_SIZE;
+                    break;
+                default:
+                    fprintf(stderr, "Error: Unsupported ML DSA level\n");
+                    exit(1);
+                    break;
+            }
+            pubDer = malloc(pubDerSz);
+            if (pubDer == NULL) {
+                fprintf(stderr,
+                        "Error: Failed to allocate memory for DER export\n");
+                exit(1);
+            }
+
             /* Export public key in DER format */
-            uint8_t pubDer[
-            #if ML_DSA_LEVEL == 2
-                ML_DSA_LEVEL2_PUB_KEY_DER_SIZE
-            #elif ML_DSA_LEVEL == 3
-                ML_DSA_LEVEL3_PUB_KEY_DER_SIZE
-            #elif ML_DSA_LEVEL == 5
-                ML_DSA_LEVEL5_PUB_KEY_DER_SIZE
-            #endif
-            ];
-            int     pubOutLen;
 
-            const int WITH_ALG_SPKI = 1;
-            pubOutLen = wc_Dilithium_PublicKeyToDer(
-                &key, pubDer, sizeof(pubDer), WITH_ALG_SPKI);
+            pubOutLen = wc_Dilithium_PublicKeyToDer(&key, pubDer, pubDerSz,
+                                                    WITH_ALG_SPKI);
             if (pubOutLen < 0) {
                 fprintf(stderr, "Unable to export public key to DER, ret=%d\n",
                         pubOutLen);
@@ -1113,6 +1130,8 @@ static void keygen_ml_dsa(const char *priv_fname, uint32_t id_mask)
                 fprintf(stderr, "Unable to export public key to file\n");
                 exit(1);
             }
+
+            free(pubDer);
         }
         else {
             /* Export public key in raw format */