Fix ML-DSA missing public key guards and memory allocation, add ECC pubkey derivation logs

Author: Emma Stensland

doc/dox_comments/header_files/wc_mldsa.h

@@ -257,6 +257,27 @@ int wc_MlDsaKey_MakeKey(wc_MlDsaKey* key, WC_RNG* rng);
 */
 int wc_MlDsaKey_MakeKeyFromSeed(wc_MlDsaKey* key, const byte* seed);
 
+/*!
+    \ingroup ML_DSA
+
+    \brief Derives and caches the public key (t1 and the public seed)
+    for a wc_MlDsaKey that already has its private key set, e.g. after
+    wc_MlDsaKey_ImportPrivRaw() or decoding a private-key-only DER. A
+    no-op returning 0 if the public key is already set.
+
+    \return 0 on success, or if the public key was already set.
+    \return BAD_FUNC_ARG if key is NULL, key->params is NULL, or the
+    private key is not set.
+    \return MEMORY_E on allocation failure.
+    \return Other negative when an error occurs.
+
+    \param [in,out] key Pointer to a wc_MlDsaKey with the private key set.
+
+    \sa wc_MlDsaKey_ImportPrivRaw
+    \sa wc_MlDsaKey_MakeKey
+*/
+int wc_MlDsaKey_MakePublicKey(wc_MlDsaKey* key);
+
 /*!
     \ingroup ML_DSA
 

tests/api/test_asn.c

@@ -2481,3 +2481,223 @@ int test_wc_AsnFeatureCoverage(void)
 #endif /* !NO_ASN && HAVE_ECC && USE_CERT_BUFFERS_256 && !HAVE_FIPS */
     return EXPECT_RESULT();
 }
+
+#if defined(USE_WOLFSSL_MEMORY) && !defined(WOLFSSL_NO_MALLOC) && \
+    !defined(WOLFSSL_STATIC_MEMORY)
+/* Rather than failing the first allocation seen after arming, count
+ * allocations and fail only the Nth (ecc_oom_fail_at). This lets the test
+ * target the allocation specific to public key derivation instead of
+ * whichever unrelated allocation (ASN parsing, big-int growth, etc.)
+ * happens to run first -- that varies by math library/build config. */
+static int ecc_oom_failed = 0;
+static int ecc_oom_inject = 0;
+static int ecc_oom_count = 0;
+static int ecc_oom_fail_at = 0;
+
+/* Custom malloc for testing OOM.
+ * Returns allocated pointer or NULL on failure. */
+#ifdef WOLFSSL_DEBUG_MEMORY
+static void* ecc_oom_malloc_cb(size_t size, const char* func,
+    unsigned int line)
+{
+    (void)func;
+    (void)line;
+#else
+static void* ecc_oom_malloc_cb(size_t size)
+{
+#endif
+    if (ecc_oom_inject) {
+        ecc_oom_count++;
+        if (!ecc_oom_failed &&
+                (ecc_oom_fail_at != 0) && (ecc_oom_count == ecc_oom_fail_at)) {
+            ecc_oom_failed = 1;
+            return NULL;
+        }
+    }
+    return malloc(size);
+}
+
+/* Custom free for testing OOM. */
+#ifdef WOLFSSL_DEBUG_MEMORY
+static void ecc_oom_free_cb(void* ptr, const char* func, unsigned int line)
+{
+    (void)func;
+    (void)line;
+#else
+static void ecc_oom_free_cb(void* ptr)
+{
+#endif
+    free(ptr);
+}
+
+/* Custom realloc for testing OOM.
+ * Returns reallocated pointer or NULL on failure. */
+#ifdef WOLFSSL_DEBUG_MEMORY
+static void* ecc_oom_realloc_cb(void* ptr, size_t size, const char* func,
+    unsigned int line)
+{
+    (void)func;
+    (void)line;
+#else
+static void* ecc_oom_realloc_cb(void* ptr, size_t size)
+{
+#endif
+    return realloc(ptr, size);
+}
+#endif /* USE_WOLFSSL_MEMORY */
+
+/* wc_EccPrivateKeyDecode should derive and cache the public point (best
+ * effort) when it decodes a SEC1 private key whose optional public point
+ * was omitted, so the key comes out fully usable. */
+int test_wc_EccPrivateKeyDecode_derive_pub(void)
+{
+    EXPECT_DECLS;
+#if !defined(NO_ASN) && defined(HAVE_ECC) && !defined(NO_ECC_MAKE_PUB) && \
+    defined(USE_CERT_BUFFERS_256) && !defined(HAVE_FIPS) && \
+    !defined(WOLFSSL_ATECC508A) && !defined(WOLFSSL_ATECC608A) && \
+    !defined(WOLFSSL_MICROCHIP_TA100) && !defined(WOLFSSL_CRYPTOCELL) && \
+    !defined(WOLFSSL_SILABS_SE_ACCEL) && !defined(WOLFSSL_KCAPI_ECC) && \
+    !defined(WOLFSSL_QNX_CAAM) && !defined(WOLFSSL_IMXRT1170_CAAM)
+    ecc_key fullKey;
+    ecc_key privOnlyKey;
+    WC_RNG  rng;
+    word32  idx;
+    byte    privOnlyDer[256];
+    int     privOnlyDerSz;
+    byte    fullPub[256];
+    word32  fullPubSz = sizeof(fullPub);
+    byte    derivedPub[256];
+    word32  derivedPubSz = sizeof(derivedPub);
+
+    XMEMSET(&fullKey, 0, sizeof(fullKey));
+    XMEMSET(&privOnlyKey, 0, sizeof(privOnlyKey));
+
+    ExpectIntEQ(wc_InitRng(&rng), 0);
+
+    ExpectIntEQ(wc_ecc_init(&fullKey), 0);
+    idx = 0;
+    ExpectIntEQ(wc_EccPrivateKeyDecode(ecc_clikey_der_256, &idx, &fullKey,
+        sizeof_ecc_clikey_der_256), 0);
+    ExpectIntEQ(fullKey.type, ECC_PRIVATEKEY);
+    PRIVATE_KEY_UNLOCK();
+    ExpectIntEQ(wc_ecc_export_x963(&fullKey, fullPub, &fullPubSz), 0);
+    PRIVATE_KEY_LOCK();
+
+    /* Re-encode as a private-key-only SEC1 DER (no public point). */
+    ExpectIntGT(privOnlyDerSz = wc_EccPrivateKeyToDer(&fullKey, privOnlyDer,
+        sizeof(privOnlyDer)), 0);
+
+#ifdef ECC_TIMING_RESISTANT
+    /* Without an RNG set, decode must not derive unblinded: the key is
+     * left as ECC_PRIVATEKEY_ONLY, same as before auto-derivation existed. */
+    ExpectIntEQ(wc_ecc_init(&privOnlyKey), 0);
+    idx = 0;
+    ExpectIntEQ(wc_EccPrivateKeyDecode(privOnlyDer, &idx, &privOnlyKey,
+        (word32)privOnlyDerSz), 0);
+    ExpectIntEQ(privOnlyKey.type, ECC_PRIVATEKEY_ONLY);
+    wc_ecc_free(&privOnlyKey);
+#endif
+
+    /* Opting into blinding (as every other blinded ECC op in this library
+     * requires) makes decode derive and cache the public point. */
+    ExpectIntEQ(wc_ecc_init(&privOnlyKey), 0);
+    ExpectIntEQ(wc_ecc_set_rng(&privOnlyKey, &rng), 0);
+    idx = 0;
+    ExpectIntEQ(wc_EccPrivateKeyDecode(privOnlyDer, &idx, &privOnlyKey,
+        (word32)privOnlyDerSz), 0);
+
+    /* The public point should have been derived automatically, making the
+     * key fully usable rather than left as ECC_PRIVATEKEY_ONLY. */
+    ExpectIntEQ(privOnlyKey.type, ECC_PRIVATEKEY);
+    PRIVATE_KEY_UNLOCK();
+    ExpectIntEQ(wc_ecc_export_x963(&privOnlyKey, derivedPub, &derivedPubSz),
+        0);
+    PRIVATE_KEY_LOCK();
+    ExpectIntEQ(derivedPubSz, fullPubSz);
+    ExpectBufEQ(derivedPub, fullPub, fullPubSz);
+
+    wc_ecc_free(&privOnlyKey);
+    wc_ecc_free(&fullKey);
+
+#if defined(USE_WOLFSSL_MEMORY) && !defined(WOLFSSL_NO_MALLOC) && \
+    !defined(WOLFSSL_STATIC_MEMORY)
+    {
+        wolfSSL_Malloc_cb prevMalloc = NULL;
+        wolfSSL_Free_cb prevFree = NULL;
+        wolfSSL_Realloc_cb prevRealloc = NULL;
+        int allocatorsSet = 0;
+        int baseAllocCount = 0;
+        int totalAllocCount = 0;
+
+        ExpectIntEQ(wolfSSL_GetAllocators(&prevMalloc, &prevFree, &prevRealloc),
+            0);
+        ExpectIntEQ(wolfSSL_SetAllocators(ecc_oom_malloc_cb, ecc_oom_free_cb,
+            ecc_oom_realloc_cb), 0);
+        if (EXPECT_SUCCESS()) {
+            allocatorsSet = 1;
+        }
+
+        /* Count (without failing) allocations for a decode that never
+         * derives -- ecc_clikey_der_256 already has its public point, so
+         * wc_ecc_derive_pub_best_effort() is a no-op. This is the baseline
+         * allocation count to skip past below. */
+        ecc_oom_count = 0;
+        ecc_oom_fail_at = 0;
+        ecc_oom_inject = 1;
+        ExpectIntEQ(wc_ecc_init(&fullKey), 0);
+        idx = 0;
+        ExpectIntEQ(wc_EccPrivateKeyDecode(ecc_clikey_der_256, &idx, &fullKey,
+            sizeof_ecc_clikey_der_256), 0);
+        ecc_oom_inject = 0;
+        wc_ecc_free(&fullKey);
+        baseAllocCount = ecc_oom_count;
+
+        /* Count allocations for a decode that does derive, for comparison. */
+        ecc_oom_count = 0;
+        ecc_oom_inject = 1;
+        ExpectIntEQ(wc_ecc_init(&privOnlyKey), 0);
+        ExpectIntEQ(wc_ecc_set_rng(&privOnlyKey, &rng), 0);
+        idx = 0;
+        ExpectIntEQ(wc_EccPrivateKeyDecode(privOnlyDer, &idx, &privOnlyKey,
+            (word32)privOnlyDerSz), 0);
+        ecc_oom_inject = 0;
+        wc_ecc_free(&privOnlyKey);
+        totalAllocCount = ecc_oom_count;
+
+        /* Only derivation allocates more than the baseline -- if it
+         * allocates nothing (e.g. stack-only math build) there is nothing
+         * to fault-inject into. */
+        if (totalAllocCount > baseAllocCount) {
+            /* Test that a derivation failure leaves the key as
+             * ECC_PRIVATEKEY_ONLY but does not fail the overall decode. */
+            ecc_oom_failed = 0;
+            ecc_oom_count = 0;
+            ecc_oom_fail_at = baseAllocCount + 1;
+
+            ExpectIntEQ(wc_ecc_init(&privOnlyKey), 0);
+            ExpectIntEQ(wc_ecc_set_rng(&privOnlyKey, &rng), 0);
+
+            idx = 0;
+            ecc_oom_inject = 1;
+            /* Decode should succeed, but leave the key as
+             * ECC_PRIVATEKEY_ONLY. */
+            ExpectIntEQ(wc_EccPrivateKeyDecode(privOnlyDer, &idx, &privOnlyKey,
+                (word32)privOnlyDerSz), 0);
+            ExpectIntEQ(ecc_oom_failed, 1);
+            ExpectIntEQ(privOnlyKey.type, ECC_PRIVATEKEY_ONLY);
+
+            wc_ecc_free(&privOnlyKey);
+            ecc_oom_inject = 0;
+        }
+
+        if (allocatorsSet) {
+            (void)wolfSSL_SetAllocators(prevMalloc, prevFree, prevRealloc);
+        }
+    }
+#endif /* USE_WOLFSSL_MEMORY */
+
+    wc_FreeRng(&rng);
+#endif /* !NO_ASN && HAVE_ECC && !NO_ECC_MAKE_PUB && USE_CERT_BUFFERS_256 &&
+        * !HAVE_FIPS */
+    return EXPECT_RESULT();
+}

tests/api/test_asn.h

@@ -44,6 +44,7 @@ int test_ToTraditional_ex_negative(void);
 int test_ToTraditional_ex_mldsa_bad_params(void);
 int test_wc_AsnDecisionCoverage(void);
 int test_wc_AsnFeatureCoverage(void);
+int test_wc_EccPrivateKeyDecode_derive_pub(void);
 
 #define TEST_ASN_DECLS                                              \
     TEST_DECL_GROUP("asn", test_SetAsymKeyDer),                     \
@@ -65,6 +66,7 @@ int test_wc_AsnFeatureCoverage(void);
     TEST_DECL_GROUP("asn", test_ToTraditional_ex_negative),         \
     TEST_DECL_GROUP("asn", test_ToTraditional_ex_mldsa_bad_params), \
     TEST_DECL_GROUP("asn", test_wc_AsnDecisionCoverage),           \
-    TEST_DECL_GROUP("asn", test_wc_AsnFeatureCoverage)
+    TEST_DECL_GROUP("asn", test_wc_AsnFeatureCoverage),            \
+    TEST_DECL_GROUP("asn", test_wc_EccPrivateKeyDecode_derive_pub)
 
 #endif /* WOLFCRYPT_TEST_ASN_H */