New API generating ML-DSA public key from private key

tests/api/test_mldsa.c

@@ -726,6 +726,115 @@ int test_wc_dilithium_make_key(void)
     return EXPECT_RESULT();
 }
 
+int test_wc_dilithium_pub_from_priv(void)
+{
+    EXPECT_DECLS;
+#if defined(HAVE_DILITHIUM) && defined(WOLFSSL_WC_DILITHIUM) && \
+    defined(WOLFSSL_DILITHIUM_PRIVATE_KEY) && defined(WOLFSSL_DILITHIUM_PUBLIC_KEY)
+    dilithium_key* key = NULL;
+    dilithium_key* importKey = NULL;
+    WC_RNG rng;
+    byte* privKey = NULL;
+    word32 privKeyLen = DILITHIUM_MAX_KEY_SIZE;
+    byte* pubKey = NULL;
+    word32 pubKeyLen = DILITHIUM_MAX_PUB_KEY_SIZE;
+    byte* origPub = NULL;
+    word32 origPubLen = DILITHIUM_MAX_PUB_KEY_SIZE;
+    int ret;
+
+    key = (dilithium_key*)XMALLOC(sizeof(*key), NULL, DYNAMIC_TYPE_TMP_BUFFER);
+    ExpectNotNull(key);
+    importKey = (dilithium_key*)XMALLOC(sizeof(*importKey), NULL,
+        DYNAMIC_TYPE_TMP_BUFFER);
+    ExpectNotNull(importKey);
+    privKey = (byte*)XMALLOC(DILITHIUM_MAX_KEY_SIZE, NULL,
+        DYNAMIC_TYPE_TMP_BUFFER);
+    ExpectNotNull(privKey);
+    pubKey = (byte*)XMALLOC(DILITHIUM_MAX_PUB_KEY_SIZE, NULL,
+        DYNAMIC_TYPE_TMP_BUFFER);
+    ExpectNotNull(pubKey);
+    origPub = (byte*)XMALLOC(DILITHIUM_MAX_PUB_KEY_SIZE, NULL,
+        DYNAMIC_TYPE_TMP_BUFFER);
+    ExpectNotNull(origPub);
+
+    if (key != NULL) XMEMSET(key, 0, sizeof(*key));
+    if (importKey != NULL) XMEMSET(importKey, 0, sizeof(*importKey));
+    XMEMSET(&rng, 0, sizeof(WC_RNG));
+
+    ExpectIntEQ(wc_InitRng(&rng), 0);
+    ExpectIntEQ(wc_dilithium_init(key), 0);
+
+#ifndef WOLFSSL_NO_ML_DSA_44
+    ExpectIntEQ(wc_dilithium_set_level(key, WC_ML_DSA_44), 0);
+#elif !defined(WOLFSSL_NO_ML_DSA_65)
+    ExpectIntEQ(wc_dilithium_set_level(key, WC_ML_DSA_65), 0);
+#else
+    ExpectIntEQ(wc_dilithium_set_level(key, WC_ML_DSA_87), 0);
+#endif
+
+#if !defined(WOLFSSL_DILITHIUM_NO_MAKE_KEY)
+    /* Generate key then export private to simulate having only priv key */
+    ExpectIntEQ(wc_dilithium_make_key(key, &rng), 0);
+    ExpectIntEQ(wc_dilithium_export_private(key, privKey, &privKeyLen), 0);
+    /* Export original public key for later comparison. */
+    origPubLen = DILITHIUM_MAX_PUB_KEY_SIZE;
+    ExpectIntEQ(wc_dilithium_export_public(key, origPub, &origPubLen), 0);
+    ExpectIntGT(origPubLen, 0);
+#else
+    /* Use built-in bench private keys when make_key not available */
+#ifndef WOLFSSL_NO_ML_DSA_44
+    XMEMCPY(privKey, bench_dilithium_level2_key, sizeof_bench_dilithium_level2_key);
+    privKeyLen = sizeof_bench_dilithium_level2_key;
+#elif !defined(WOLFSSL_NO_ML_DSA_65)
+    XMEMCPY(privKey, bench_dilithium_level3_key, sizeof_bench_dilithium_level3_key);
+    privKeyLen = sizeof_bench_dilithium_level3_key;
+#else
+    XMEMCPY(privKey, bench_dilithium_level5_key, sizeof_bench_dilithium_level5_key);
+    privKeyLen = sizeof_bench_dilithium_level5_key;
+#endif
+#endif
+
+    ExpectIntEQ(wc_dilithium_init(importKey), 0);
+    /* Ensure importKey has the same security level set as key so import
+     * functions that validate level do not fail. */
+#ifndef WOLFSSL_NO_ML_DSA_44
+    ExpectIntEQ(wc_dilithium_set_level(importKey, WC_ML_DSA_44), 0);
+#elif !defined(WOLFSSL_NO_ML_DSA_65)
+    ExpectIntEQ(wc_dilithium_set_level(importKey, WC_ML_DSA_65), 0);
+#else
+    ExpectIntEQ(wc_dilithium_set_level(importKey, WC_ML_DSA_87), 0);
+#endif
+    ExpectIntEQ(wc_dilithium_import_private(privKey, privKeyLen, importKey), 0);
+
+    /* At this point importKey should only have private key; derive public */
+    ret = wc_dilithium_pub_from_priv(importKey);
+    ExpectIntEQ(ret, 0);
+
+    pubKeyLen = DILITHIUM_MAX_PUB_KEY_SIZE;
+    ExpectIntEQ(wc_dilithium_export_public(importKey, pubKey, &pubKeyLen), 0);
+    ExpectIntGT(pubKeyLen, 0);
+
+    /* If we generated the original key, compare its public key to the one
+     * derived from the imported private key. */
+#if !defined(WOLFSSL_DILITHIUM_NO_MAKE_KEY)
+    ExpectIntEQ(origPubLen, pubKeyLen);
+    ExpectIntEQ(XMEMCMP(origPub, pubKey, pubKeyLen), 0);
+#endif
+
+    wc_dilithium_free(importKey);
+    wc_dilithium_free(key);
+    wc_FreeRng(&rng);
+
+    XFREE(pubKey, NULL, DYNAMIC_TYPE_TMP_BUFFER);
+    XFREE(origPub, NULL, DYNAMIC_TYPE_TMP_BUFFER);
+    XFREE(privKey, NULL, DYNAMIC_TYPE_TMP_BUFFER);
+    XFREE(importKey, NULL, DYNAMIC_TYPE_TMP_BUFFER);
+    XFREE(key, NULL, DYNAMIC_TYPE_TMP_BUFFER);
+#endif /* HAVE_DILITHIUM && WOLFSSL_WC_DILITHIUM && PRIVATE+PUBLIC */
+
+    return EXPECT_RESULT();
+}
+
 int test_wc_dilithium_sign(void)
 {
     EXPECT_DECLS;

tests/api/test_mldsa.h

@@ -26,6 +26,7 @@
 
 int test_wc_dilithium(void);
 int test_wc_dilithium_make_key(void);
+int test_wc_dilithium_pub_from_priv(void);
 int test_wc_dilithium_sign(void);
 int test_wc_dilithium_verify(void);
 int test_wc_dilithium_sign_vfy(void);
@@ -43,6 +44,7 @@ int test_mldsa_pkcs12(void);
 #define TEST_MLDSA_DECLS                                                       \
     TEST_DECL_GROUP("mldsa", test_wc_dilithium),                               \
     TEST_DECL_GROUP("mldsa", test_wc_dilithium_make_key),                      \
+    TEST_DECL_GROUP("mldsa", test_wc_dilithium_pub_from_priv),                 \
     TEST_DECL_GROUP("mldsa", test_wc_dilithium_sign),                          \
     TEST_DECL_GROUP("mldsa", test_wc_dilithium_verify),                        \
     TEST_DECL_GROUP("mldsa", test_wc_dilithium_sign_vfy),                      \

wolfcrypt/src/dilithium.c

@@ -7980,6 +7980,145 @@ static int dilithium_make_key_from_seed(dilithium_key* key, const byte* seed)
 #endif
 }
 
+static int dilithium_pub_from_priv(dilithium_key* key)
+{
+    int ret = 0;
+    const wc_dilithium_params* params = key->params;
+    const byte* pub_seed = key->k;
+    const byte* s1p = pub_seed + DILITHIUM_PUB_SEED_SZ + DILITHIUM_K_SZ + DILITHIUM_TR_SZ;
+    const byte* s2p = s1p + params->s1EncSz;
+    sword32* a  = NULL;
+    sword32* s1 = NULL;
+    sword32* s2 = NULL;
+    sword32* t  = NULL;
+    byte* t0 = NULL;
+    byte* t1 = key->p + DILITHIUM_PUB_SEED_SZ;
+
+    /* Allocate and create cached values. */
+#ifndef WC_DILITHIUM_CACHE_MATRIX_A
+    a = (sword32*)XMALLOC(params->aSz, key->heap,
+        DYNAMIC_TYPE_DILITHIUM);
+    if (a == NULL) {
+        ret = MEMORY_E;
+    }
+    else {
+        XMEMSET(a, 0, params->aSz);
+    }
+
+    if (ret == 0) {
+        ret = dilithium_expand_a(&key->shake, pub_seed, params->k, params->l,
+            a, key->heap);
+    }
+#else
+    if (ret == 0) {
+        if (key->a == NULL) {
+            ret = BAD_STATE_E;
+        }
+        else {
+            a = key->a;
+        }
+    }
+#endif
+#ifndef WC_DILITHIUM_CACHE_PRIV_VECTORS
+    if (ret == 0) {
+        s1 = (sword32*)XMALLOC(params->s1Sz + params->s2Sz, key->heap,
+            DYNAMIC_TYPE_DILITHIUM);
+        if (s1 == NULL) {
+            ret = MEMORY_E;
+           }
+        else {
+            s2 = s1 + params->s1Sz / sizeof(*s1);
+
+            dilithium_vec_decode_eta_bits(s1p, params->eta, s1, params->l);
+            dilithium_vec_decode_eta_bits(s2p, params->eta, s2, params->k);
+
+            dilithium_vec_ntt_small(s1, params->l);
+        }
+    }
+#else
+    if (ret == 0) {
+        if (key->s1 == NULL || key->s2 == NULL) {
+            ret = BAD_STATE_E;
+        }
+        else {
+            s1 = key->s1;
+            s2 = key->s2;
+        }
+    }
+#endif
+
+    if (ret == 0) {
+        t0 = (byte*)XMALLOC(params->s2Sz, key->heap, DYNAMIC_TYPE_DILITHIUM);
+        if (t0 == NULL) {
+            ret = MEMORY_E;
+        }
+    }
+
+    if (ret == 0) {
+        t = (sword32*)XMALLOC(params->s2Sz, key->heap, DYNAMIC_TYPE_DILITHIUM);
+        if (t == NULL) {
+            ret = MEMORY_E;
+        }
+    }
+
+    /* cal t and get t0, t1 */
+    if (ret == 0) {
+        /* Copy public seed into public key. */
+        XMEMCPY(key->p, pub_seed, DILITHIUM_PUB_SEED_SZ);
+
+        /* t <- NTT-1(A_circum o NTT(s1)) + s2 */
+        dilithium_matrix_mul(t, a, s1, params->k, params->l);
+        dilithium_vec_invntt_full(t, params->k);
+        dilithium_vec_add(t, s2, params->k);
+        /* NTT s2 */
+        dilithium_vec_ntt_small(s2, params->k);
+
+        /* Make positive for decomposing. */
+        dilithium_vec_make_pos(t, params->k);
+        /* Decompose t in t0 and t1 and encode into public and private key. */
+        dilithium_vec_encode_t0_t1(t, params->k, t0, t1);
+    }
+
+#ifndef WC_DILITHIUM_CACHE_MATRIX_A
+    XMEMSET(a, 0, params->aSz);
+    XFREE(a, key->heap, DYNAMIC_TYPE_DILITHIUM);
+#endif
+#ifndef WC_DILITHIUM_CACHE_PRIV_VECTORS
+    XMEMSET(s1, 0, params->s1Sz + params->s2Sz);
+    XFREE(s1, key->heap, DYNAMIC_TYPE_DILITHIUM);
+#endif
+    XMEMSET(t0, 0, params->s2Sz);
+    XMEMSET(t, 0, params->s2Sz);
+    XFREE(t0, key->heap, DYNAMIC_TYPE_DILITHIUM);
+    XFREE(t, key->heap, DYNAMIC_TYPE_DILITHIUM);
+
+    if (ret == 0) {
+#ifdef WC_DILITHIUM_CACHE_PUB_VECTORS
+    #ifndef WC_DILITHIUM_FIXED_ARRAY
+        /* Allocate t1 if required. */
+        if (key->t1 == NULL) {
+            key->t1 = (sword32*)XMALLOC(params->s2Sz, key->heap,
+                DYNAMIC_TYPE_DILITHIUM);
+            if (key->t1 == NULL) {
+                ret = MEMORY_E;
+            }
+            else {
+                XMEMSET(key->t1, 0, key->params->s2Sz);
+            }
+        }
+    #endif
+    }
+    if (ret == 0) {
+        /* Compute t1 from public key data. */
+        dilithium_make_pub_vec(key, key->t1);
+#endif
+        /* Public key is set. */
+        key->pubKeySet = 1;
+    }
+
+    return ret;
+}
+
 /* Make a key from a random seed.
  *
  * FIPS 204. 5.1: Algorithm 1 ML-DSA.KeyGen()
@@ -10097,6 +10236,25 @@ int wc_dilithium_make_key_from_seed(dilithium_key* key, const byte* seed)
 
     return ret;
 }
+
+int wc_dilithium_pub_from_priv(dilithium_key* key)
+{
+    int ret = 0;
+
+    if (key == NULL) {
+        ret = BAD_FUNC_ARG;
+    }
+
+    if (ret == 0) {
+#ifdef WOLFSSL_WC_DILITHIUM
+        ret = dilithium_pub_from_priv(key);
+#elif defined(HAVE_LIBOQS)
+        ret = NOT_COMPILED_IN;
+#endif
+    }
+
+    return ret;
+}
 #endif
 
 #ifndef WOLFSSL_DILITHIUM_NO_SIGN

wolfssl/wolfcrypt/dilithium.h

@@ -806,6 +806,8 @@ WOLFSSL_API
 int wc_dilithium_make_key(dilithium_key* key, WC_RNG* rng);
 WOLFSSL_API
 int wc_dilithium_make_key_from_seed(dilithium_key* key, const byte* seed);
+WOLFSSL_API
+int wc_dilithium_pub_from_priv(dilithium_key* key);
 
 WOLFSSL_API
 int wc_dilithium_sign_msg(const byte* msg, word32 msgLen, byte* sig,