Separate public and private ML-KEM allocations

dndfrink-ibm · t8m · commit b4fedba43caa · 2025-07-07T15:40:47.000+02:00
Previously, this change had grouped the public and private portions of the ML-KEM key structure into one allocation that was changed to use secure memory. There were concerns raised that there may be use cases where storage of many ML-KEM public keys may be necessary. Since the total secure memory size is configured by the user, reduce the footprint of secure memory usage to reduce the impact of these changes on users of these flows. Reviewed-by: Shane Lontis <shane.lontis@oracle.com> Reviewed-by: Viktor Dukhovni <viktor@openssl.org> (Merged from openssl#27625)
diff --git a/crypto/ml_dsa/ml_dsa_encoders.c b/crypto/ml_dsa/ml_dsa_encoders.c
@@ -747,7 +747,7 @@ int ossl_ml_dsa_sk_encode(ML_DSA_KEY *key)
 err:
     WPACKET_finish(&pkt);
     if (ret == 0)
-        OPENSSL_clear_free(enc, enc_len);
+        OPENSSL_secure_clear_free(enc, enc_len);
     return ret;
 }
 
@@ -770,7 +770,7 @@ int ossl_ml_dsa_sk_decode(ML_DSA_KEY *key, const uint8_t *in, size_t in_len)
     PACKET pkt;
 
     /* When loading from an explicit key, drop the seed. */
-    OPENSSL_clear_free(key->seed, ML_DSA_SEED_BYTES);
+    OPENSSL_secure_clear_free(key->seed, ML_DSA_SEED_BYTES);
     key->seed = NULL;
 
     /* Allow the key encoding to be already set to the provided pointer */
diff --git a/crypto/ml_dsa/ml_dsa_key.c b/crypto/ml_dsa/ml_dsa_key.c
@@ -55,17 +55,19 @@ int ossl_ml_dsa_set_prekey(ML_DSA_KEY *key, int flags_set, int flags_clr,
         memcpy(key->priv_encoding, sk, sk_len);
     }
 
-    if (seed != NULL
-        && (key->seed = OPENSSL_memdup(seed, seed_len)) == NULL)
-        goto end;
+    if (seed != NULL) {
+        if ((key->seed = OPENSSL_secure_malloc(seed_len)) == NULL)
+            goto end;
+        memcpy(key->seed, seed, seed_len);
+    }
     key->prov_flags |= flags_set;
     key->prov_flags &= ~flags_clr;
     ret = 1;
 
  end:
     if (!ret) {
-        OPENSSL_secure_free(key->priv_encoding);
-        OPENSSL_free(key->seed);
+        OPENSSL_secure_clear_free(key->priv_encoding, sk_len);
+        OPENSSL_secure_clear_free(key->seed, seed_len);
         key->priv_encoding = key->seed = NULL;
     }
     return ret;
@@ -152,12 +154,12 @@ void ossl_ml_dsa_key_reset(ML_DSA_KEY *key)
      * must not access after |s1|'s poly is freed.
      */
     if (key->s1.poly != NULL) {
-        vector_zero(&key->s1);
-        vector_zero(&key->s2);
-        vector_zero(&key->t0);
-        vector_secure_free(&key->s1);
-        key->s2.poly = NULL;
-        key->t0.poly = NULL;
+        const ML_DSA_PARAMS *params = key->params;
+        size_t k = params->k, l = params->l;
+
+        vector_secure_free(&key->s1, l + 2 * k);
+        vector_init(&key->s2, NULL, 0);
+        vector_init(&key->t0, NULL, 0);
     }
     /* The |t1| vector is public and allocated separately */
     vector_free(&key->t1);
@@ -168,7 +170,7 @@ void ossl_ml_dsa_key_reset(ML_DSA_KEY *key)
         OPENSSL_secure_clear_free(key->priv_encoding, key->params->sk_len);
     key->priv_encoding = NULL;
     if (key->seed != NULL)
-        OPENSSL_clear_free(key->seed, ML_DSA_SEED_BYTES);
+        OPENSSL_secure_clear_free(key->seed, ML_DSA_SEED_BYTES);
     key->seed = NULL;
 }
 
@@ -222,14 +224,16 @@ ML_DSA_KEY *ossl_ml_dsa_key_dup(const ML_DSA_KEY *src, int selection)
                         vector_copy(&ret->t0, &src->t0);
                     }
                     ret->priv_encoding = OPENSSL_secure_malloc(src->params->sk_len);
-                    if (!ret->priv_encoding)
+                    if (ret->priv_encoding == NULL)
                         goto err;
                     memcpy(ret->priv_encoding, src->priv_encoding, src->params->sk_len);
                 }
-                if (src->seed != NULL
-                    && (ret->seed = OPENSSL_memdup(src->seed,
-                                                   ML_DSA_SEED_BYTES)) == NULL)
-                    goto err;
+                if (src->seed != NULL) {
+                    ret->seed = OPENSSL_secure_malloc(ML_DSA_SEED_BYTES);
+                    if (ret->seed == NULL)
+                        goto err;
+                    memcpy(ret->seed, src->seed, ML_DSA_SEED_BYTES);
+                }
             }
         }
         EVP_MD_up_ref(src->shake128_md);
@@ -465,10 +469,10 @@ int ossl_ml_dsa_generate_key(ML_DSA_KEY *out)
     int ret;
 
     if (out->seed == NULL) {
-        if ((out->seed = OPENSSL_malloc(seed_len)) == NULL)
+        if ((out->seed = OPENSSL_secure_malloc(seed_len)) == NULL)
             return 0;
         if (RAND_priv_bytes_ex(out->libctx, out->seed, seed_len, 0) <= 0) {
-            OPENSSL_free(out->seed);
+            OPENSSL_secure_free(out->seed);
             out->seed = NULL;
             return 0;
         }
diff --git a/crypto/ml_dsa/ml_dsa_vector.h b/crypto/ml_dsa/ml_dsa_vector.h
@@ -59,9 +59,9 @@ void vector_free(VECTOR *v)
 }
 
 static ossl_inline ossl_unused
-void vector_secure_free(VECTOR *v)
+void vector_secure_free(VECTOR *v, size_t rank)
 {
-    OPENSSL_secure_clear_free(v->poly, v->num_poly * sizeof(POLY));
+    OPENSSL_secure_clear_free(v->poly, rank * sizeof(POLY));
     v->poly = NULL;
     v->num_poly = 0;
 }
diff --git a/crypto/ml_kem/ml_kem.c b/crypto/ml_kem/ml_kem.c
@@ -94,27 +94,31 @@ typedef struct ossl_ml_kem_scalar_st {
 } scalar;
 
 /* Key material allocation layout */
-#define DECLARE_ML_KEM_KEYDATA(name, rank, private_sz) \
+#define DECLARE_ML_KEM_PUBKEYDATA(name, rank) \
     struct name##_alloc { \
         /* Public vector |t| */ \
         scalar tbuf[(rank)]; \
         /* Pre-computed matrix |m| (FIPS 203 |A| transpose) */ \
-        scalar mbuf[(rank)*(rank)] \
-        /* optional private key data */ \
-        private_sz \
+        scalar mbuf[(rank)*(rank)]; \
+    }
+
+#define DECLARE_ML_KEM_PRVKEYDATA(name, rank) \
+    struct name##_alloc { \
+        scalar sbuf[rank]; \
+        uint8_t zbuf[2 * ML_KEM_RANDOM_BYTES]; \
     }
 
 /* Declare variant-specific public and private storage */
 #define DECLARE_ML_KEM_VARIANT_KEYDATA(bits) \
-    DECLARE_ML_KEM_KEYDATA(pubkey_##bits, ML_KEM_##bits##_RANK,;); \
-    DECLARE_ML_KEM_KEYDATA(prvkey_##bits, ML_KEM_##bits##_RANK,;\
-        scalar sbuf[ML_KEM_##bits##_RANK]; \
-        uint8_t zbuf[2 * ML_KEM_RANDOM_BYTES];)
+    DECLARE_ML_KEM_PUBKEYDATA(pubkey_##bits, ML_KEM_##bits##_RANK); \
+    DECLARE_ML_KEM_PRVKEYDATA(prvkey_##bits, ML_KEM_##bits##_RANK)
+
 DECLARE_ML_KEM_VARIANT_KEYDATA(512);
 DECLARE_ML_KEM_VARIANT_KEYDATA(768);
 DECLARE_ML_KEM_VARIANT_KEYDATA(1024);
 #undef DECLARE_ML_KEM_VARIANT_KEYDATA
-#undef DECLARE_ML_KEM_KEYDATA
+#undef DECLARE_ML_KEM_PUBKEYDATA
+#undef DECLARE_ML_KEM_PRVKEYDATA
 
 typedef __owur
 int (*CBD_FUNC)(scalar *out, uint8_t in[ML_KEM_RANDOM_BYTES + 1],
@@ -1534,26 +1538,35 @@ int decap(uint8_t secret[ML_KEM_SHARED_SECRET_BYTES],
 /*
  * After allocating storage for public or private key data, update the key
  * component pointers to reference that storage.
+ *
+ * The caller should only store private data in `priv` *after* a successful
+ * (non-zero) return from this function.
  */
 static __owur
-int add_storage(scalar *p, int private, ML_KEM_KEY *key)
+int add_storage(scalar *pub, scalar *priv, int private, ML_KEM_KEY *key)
 {
     int rank = key->vinfo->rank;
 
-    if (p == NULL)
+    if (pub == NULL || (private && priv == NULL)) {
+        /*
+         * One of these could be allocated correctly. It is legal to call free with a NULL
+         * pointer, so always attempt to free both allocations here
+         */
+        OPENSSL_free(pub);
+        OPENSSL_secure_free(priv);
         return 0;
+    }
 
     /*
-     * We're adding key material, the seed buffer will now hold |rho| and
-     * |pkhash|.
+     * We're adding key material, set up rho and pkhash to point to the rho_pkhash buffer
      */
-    memset(key->seedbuf, 0, sizeof(key->seedbuf));
-    key->rho = key->seedbuf;
-    key->pkhash = key->seedbuf + ML_KEM_RANDOM_BYTES;
+    memset(key->rho_pkhash, 0, sizeof(key->rho_pkhash));
+    key->rho = key->rho_pkhash;
+    key->pkhash = key->rho_pkhash + ML_KEM_RANDOM_BYTES;
     key->d = key->z = NULL;
 
     /* A public key needs space for |t| and |m| */
-    key->m = (key->t = p) + rank;
+    key->m = (key->t = pub) + rank;
 
     /*
      * A private key also needs space for |s| and |z|.
@@ -1563,7 +1576,7 @@ int add_storage(scalar *p, int private, ML_KEM_KEY *key)
      * non-NULL |d| pointer.
      */
     if (private)
-        key->z = (uint8_t *)(rank + (key->s = key->m + rank * rank));
+        key->z = (uint8_t *)(rank + (key->s = priv));
     return 1;
 }
 
@@ -1574,27 +1587,29 @@ int add_storage(scalar *p, int private, ML_KEM_KEY *key)
 void
 ossl_ml_kem_key_reset(ML_KEM_KEY *key)
 {
-    if (key->t == NULL)
-        return;
+    /*
+     * seedbuf can be allocated and contain |z| and |d| if the key is
+     * being created from a private key encoding.  Similarly a pending
+     * serialised (encoded) private key may be queued up to load.
+     * Clear and free that data now.
+     */
+    if (key->seedbuf != NULL)
+        OPENSSL_secure_clear_free(key->seedbuf, ML_KEM_SEED_BYTES);
+    if (ossl_ml_kem_have_dkenc(key))
+        OPENSSL_secure_clear_free(key->encoded_dk, key->vinfo->prvkey_bytes);
+
     /*-
      * Cleanse any sensitive data:
      * - The private vector |s| is immediately followed by the FO failure
      *   secret |z|, and seed |d|, we can cleanse all three in one call.
-     *
-     * - Otherwise, when key->d is set, cleanse the stashed seed.
-     *
-     * If the memory has been allocated with secure memory, it will be cleared
-     * before being free'd under the OPENSSL_secure_free call.
      */
-    if (ossl_ml_kem_have_prvkey(key)) {
-        if (!CRYPTO_secure_allocated(key->t))
-            OPENSSL_cleanse(key->s, key->vinfo->rank * sizeof(scalar) + 2 * ML_KEM_RANDOM_BYTES);
-        OPENSSL_secure_free(key->t);
-    } else {
+    if (key->t != NULL) {
+        if (ossl_ml_kem_have_prvkey(key))
+            OPENSSL_secure_clear_free(key->s, key->vinfo->prvalloc);
         OPENSSL_free(key->t);
     }
-
-    key->d = key->z = (uint8_t *)(key->s = key->m = key->t = NULL);
+    key->d = key->z = key->seedbuf = key->encoded_dk =
+        (uint8_t *)(key->s = key->m = key->t = NULL);
 }
 
 /*
@@ -1640,7 +1655,7 @@ ML_KEM_KEY *ossl_ml_kem_key_new(OSSL_LIB_CTX *libctx, const char *properties,
     key->shake256_md = EVP_MD_fetch(libctx, "SHAKE256", properties);
     key->sha3_256_md = EVP_MD_fetch(libctx, "SHA3-256", properties);
     key->sha3_512_md = EVP_MD_fetch(libctx, "SHA3-512", properties);
-    key->d = key->z = key->rho = key->pkhash = key->encoded_dk = NULL;
+    key->d = key->z = key->rho = key->pkhash = key->encoded_dk = key->seedbuf = NULL;
     key->s = key->m = key->t = NULL;
 
     if (key->shake128_md != NULL
@@ -1660,7 +1675,8 @@ ML_KEM_KEY *ossl_ml_kem_key_dup(const ML_KEM_KEY *key, int selection)
 {
     int ok = 0;
     ML_KEM_KEY *ret;
-    void *tmp;
+    void *tmp_pub;
+    void *tmp_priv;
 
     /*
      * Partially decoded keys, not yet imported or loaded, should never be
@@ -1669,9 +1685,11 @@ ML_KEM_KEY *ossl_ml_kem_key_dup(const ML_KEM_KEY *key, int selection)
     if (ossl_ml_kem_decoded_key(key))
         return NULL;
 
-    if (key == NULL
-        || (ret = OPENSSL_memdup(key, sizeof(*key))) == NULL)
+    if (key == NULL)
+        return NULL;
+    else if ((ret = OPENSSL_memdup(key, sizeof(*key))) == NULL)
         return NULL;
+
     ret->d = ret->z = ret->rho = ret->pkhash = NULL;
     ret->s = ret->m = ret->t = NULL;
 
@@ -1686,16 +1704,19 @@ ML_KEM_KEY *ossl_ml_kem_key_dup(const ML_KEM_KEY *key, int selection)
         ok = 1;
         break;
     case OSSL_KEYMGMT_SELECT_PUBLIC_KEY:
-        ok = add_storage(OPENSSL_memdup(key->t, key->vinfo->puballoc), 0, ret);
-        ret->rho = ret->seedbuf;
-        ret->pkhash = ret->rho + ML_KEM_RANDOM_BYTES;
+        ok = add_storage(OPENSSL_memdup(key->t, key->vinfo->puballoc), NULL, 0, ret);
         break;
     case OSSL_KEYMGMT_SELECT_PRIVATE_KEY:
-        tmp = OPENSSL_secure_malloc(key->vinfo->prvalloc);
-        if (tmp == NULL)
+        tmp_pub = OPENSSL_memdup(key->t, key->vinfo->puballoc);
+        if (tmp_pub == NULL)
+            break;
+        tmp_priv = OPENSSL_secure_malloc(key->vinfo->prvalloc);
+        if (tmp_priv == NULL) {
+            OPENSSL_free(tmp_pub);
             break;
-        memcpy(tmp, key->t, key->vinfo->prvalloc);
-        ok = add_storage(tmp, 1, ret);
+        }
+        if ((ok = add_storage(tmp_pub, tmp_priv, 1, ret)) != 0)
+            memcpy(tmp_priv, key->s, key->vinfo->prvalloc);
         /* Duplicated keys retain |d|, if available */
         if (key->d != NULL)
             ret->d = ret->z + ML_KEM_RANDOM_BYTES;
@@ -1725,11 +1746,6 @@ void ossl_ml_kem_key_free(ML_KEM_KEY *key)
     EVP_MD_free(key->sha3_256_md);
     EVP_MD_free(key->sha3_512_md);
 
-    if (ossl_ml_kem_decoded_key(key)) {
-        OPENSSL_cleanse(key->seedbuf, sizeof(key->seedbuf));
-        if (ossl_ml_kem_have_dkenc(key))
-            OPENSSL_secure_clear_free(key->encoded_dk, key->vinfo->prvkey_bytes);
-    }
     ossl_ml_kem_key_reset(key);
     OPENSSL_free(key);
 }
@@ -1783,10 +1799,13 @@ ML_KEM_KEY *ossl_ml_kem_set_seed(const uint8_t *seed, size_t seedlen, ML_KEM_KEY
         || ossl_ml_kem_have_seed(key)
         || seedlen != ML_KEM_SEED_BYTES)
         return NULL;
-    /*
-     * With no public or private key material on hand, we can use the seed
-     * buffer for |z| and |d|, in that order.
-     */
+
+    if (key->seedbuf == NULL) {
+        key->seedbuf = OPENSSL_secure_malloc(seedlen);
+        if (key->seedbuf == NULL)
+            return NULL;
+    }
+
     key->z = key->seedbuf;
     key->d = key->z + ML_KEM_RANDOM_BYTES;
     memcpy(key->d, seed, ML_KEM_RANDOM_BYTES);
@@ -1813,7 +1832,7 @@ int ossl_ml_kem_parse_public_key(const uint8_t *in, size_t len, ML_KEM_KEY *key)
         || (mdctx = EVP_MD_CTX_new()) == NULL)
         return 0;
 
-    if (add_storage(OPENSSL_malloc(vinfo->puballoc), 0, key))
+    if (add_storage(OPENSSL_malloc(vinfo->puballoc), NULL, 0, key))
         ret = parse_pubkey(in, mdctx, key);
 
     if (!ret)
@@ -1841,7 +1860,8 @@ int ossl_ml_kem_parse_private_key(const uint8_t *in, size_t len,
         || (mdctx = EVP_MD_CTX_new()) == NULL)
         return 0;
 
-    if (add_storage(OPENSSL_secure_malloc(vinfo->prvalloc), 1, key))
+    if (add_storage(OPENSSL_malloc(vinfo->puballoc),
+                    OPENSSL_secure_malloc(vinfo->prvalloc), 1, key))
         ret = parse_prvkey(in, mdctx, key);
 
     if (!ret)
@@ -1870,10 +1890,10 @@ int ossl_ml_kem_genkey(uint8_t *pubenc, size_t publen, ML_KEM_KEY *key)
     if (pubenc != NULL && publen != vinfo->pubkey_bytes)
         return 0;
 
-    if (ossl_ml_kem_have_seed(key)) {
+    if (key->seedbuf != NULL) {
         if (!ossl_ml_kem_encode_seed(seed, sizeof(seed), key))
             return 0;
-        key->d = key->z = NULL;
+        ossl_ml_kem_key_reset(key);
     } else if (RAND_priv_bytes_ex(key->libctx, seed, sizeof(seed),
                                   key->vinfo->secbits) <= 0) {
         return 0;
@@ -1888,7 +1908,8 @@ int ossl_ml_kem_genkey(uint8_t *pubenc, size_t publen, ML_KEM_KEY *key)
      */
     CONSTTIME_SECRET(seed, ML_KEM_SEED_BYTES);
 
-    if (add_storage(OPENSSL_secure_malloc(vinfo->prvalloc), 1, key))
+    if (add_storage(OPENSSL_malloc(vinfo->puballoc),
+                    OPENSSL_secure_malloc(vinfo->prvalloc), 1, key))
         ret = genkey(seed, mdctx, pubenc, key);
     OPENSSL_cleanse(seed, sizeof(seed));
 
diff --git a/include/crypto/ml_kem.h b/include/crypto/ml_kem.h

Original file line number	Diff line number	Diff line change
`@@ -59,9 +59,9 @@ void vector_free(VECTOR *v)`
`59`	`59`	`}`
`60`	`60`
`61`	`61`	`static ossl_inline ossl_unused`
`62`		`-void vector_secure_free(VECTOR *v)`
	`62`	`+void vector_secure_free(VECTOR *v, size_t rank)`
`63`	`63`	`{`
`64`		`- OPENSSL_secure_clear_free(v->poly, v->num_poly * sizeof(POLY));`
	`64`	`+ OPENSSL_secure_clear_free(v->poly, rank * sizeof(POLY));`
`65`	`65`	`v->poly = NULL;`
`66`	`66`	`v->num_poly = 0;`
`67`	`67`	`}`