Multipart sign/verify: dispatch functions, dispatch to ML_DSA

drivers/builtin/include/mbedtls/private/crypto_builtin_composites.h

@@ -176,6 +176,15 @@ typedef struct {
 #define MBEDTLS_VERIFY_SIGN_HASH_INTERRUPTIBLE_OPERATION_INIT { 0 }
 #endif
 
+/* There are currently no built-in functions for multipart
+ * sign/verify-message. However, we have some boilerplate code that references
+ * such hypothetical functions in the driver dispatch code, to serve as
+ * a demonstration of what kind of code is expected there. Have a macro
+ * definition in the code to reassure check_names.py that the macro name
+ * is not a mistake. */
+#if 0
+#define MBEDTLS_PSA_BUILTIN_SIGNATURE_MULTIPART
+#endif
 
 /* EC-JPAKE operation definitions */
 

drivers/pqcp/include/tf-psa-crypto/private/crypto_struct_pqcp.h

@@ -18,7 +18,34 @@
 
 #if defined(TF_PSA_CRYPTO_PQCP_MLDSA_ENABLED)
 
-/* No context structures yet. */
+#if defined(TF_PSA_CRYPTO_PQCP_OWN_SHAKE)
+/* This is a copy of the type mld_shake128ctx
+ * in drivers/mldsa-native/mldsa/src/fips202/fips202.h,
+ * which is not an installed header. */
+typedef struct {
+    uint64_t s[25];
+    unsigned int pos;
+} tf_psa_crypto_mldsa_shake256_struct_t;
+typedef union {
+    unsigned char bytes[sizeof(tf_psa_crypto_mldsa_shake256_struct_t)];
+    tf_psa_crypto_mldsa_shake256_struct_t structure;
+} tf_psa_crypto_mldsa_shake256_t;
+#else
+typedef struct psa_xof_operation_s tf_psa_crypto_mldsa_shake256_t;
+#endif
+
+typedef struct {
+    /* Depending on the library version and compilation options, some fields
+     * may be unused or repurposed. See psa_crypto_mldsa.c for actual usage
+     * details. */
+    uint8_t parameter_set;      /* 44, 65 or 87 */
+    uint8_t hedged;             /* 0=deterministic, 1=hedged */
+    uint8_t context_set;        /* boolean: has set_context been called? */
+    uint8_t reserved;           /* uses space anyway, reserver for future use */
+    size_t key_length;          /* size of key in bytes */
+    uint8_t *key;               /* heap pointer, owned by the driver */
+    tf_psa_crypto_mldsa_shake256_t shake;
+} tf_psa_crypto_mldsa_operation_t;
 
 #endif /* TF_PSA_CRYPTO_PQCP_MLDSA_ENABLED */
 

drivers/pqcp/src/psa_crypto_mldsa.c

@@ -12,6 +12,72 @@
 #include "psa_crypto_mldsa.h"
 #include "wrap_mldsa_native.h"
 #include <mbedtls/platform_util.h>
+#include <mbedtls/platform.h>
+
+#if defined(TF_PSA_CRYPTO_PQCP_OWN_SHAKE)
+#include "../mldsa-native/mldsa/src/fips202/fips202.h"
+
+/* The mldsa-native header defines the SHAKE context types, declares
+ * the functions that work on that type, and declares macros mld_xxx.
+ *
+ * We need to expose the context type in our public headers since it
+ * appears in multipart operation structure, but we don't want to
+ * expose the function declarations. (Maybe we will in the future, but
+ * at the time of writing, it would be a hassle because fips202.h
+ * references many headers of mldsa-native that we don't want to
+ * install.)
+ *
+ * Therefore we define a public type which has (at least) the same size
+ * and alignment requirements as context type from mldsa-native, but
+ * is a distinct type according to the C language. We call the
+ * mldsa-native SHAKE functions through a wrapper that puns the
+ * pointer type. We pun via a union to a byte array to make the compiler
+ * understand that this must be aliasing of memory.
+ *
+ * In this source file, we want to call the SHAKE256 operations through
+ * the same names as when using the PSA callbacks for SHAKE, i.e. the
+ * mld_xxx names. So we replace the mld_xxx macros by wrappers that
+ * do the necessary pointer punning.
+ */
+
+#undef mld_shake256_init
+#undef mld_shake256_absorb
+#undef mld_shake256_finalize
+#undef mld_shake256_squeeze
+#undef mld_shake256_release
+
+static inline void mld_shake256_init(tf_psa_crypto_mldsa_shake256_t *state)
+{
+    MLD_NAMESPACE(shake256_init)((mld_shake256ctx *) &state->bytes);
+}
+
+static inline void mld_shake256_absorb(tf_psa_crypto_mldsa_shake256_t *state,
+                                       const uint8_t *in, size_t inlen)
+{
+    MLD_NAMESPACE(shake256_absorb)((mld_shake256ctx *) &state->bytes,
+                                   in, inlen);
+}
+
+static inline void mld_shake256_finalize(tf_psa_crypto_mldsa_shake256_t *state)
+{
+    MLD_NAMESPACE(shake256_finalize)((mld_shake256ctx *) &state->bytes);
+}
+
+static inline void mld_shake256_squeeze(uint8_t *out, size_t outlen,
+                                        tf_psa_crypto_mldsa_shake256_t *state)
+{
+    MLD_NAMESPACE(shake256_squeeze)(out, outlen,
+                                    (mld_shake256ctx *) &state->bytes);
+}
+
+static inline void mld_shake256_release(tf_psa_crypto_mldsa_shake256_t *state)
+{
+    MLD_NAMESPACE(shake256_release)((mld_shake256ctx *) &state->bytes);
+}
+
+#else /* TF_PSA_CRYPTO_PQCP_OWN_SHAKE */
+#include "fips202_psa.h"
+#endif
 
 /* The size of an ML-DSA seed in bytes.
  * The PSA API uses the seed as the private key.
@@ -20,11 +86,24 @@
  */
 #define SEED_SIZE 32
 
+/* We want to expose size values in public headers, but we don't want to
+ * expose the header that defines macros for these values in mldsa-native.
+ * So we define our own macros in public headers, and check that the
+ * values match.
+ */
+#if MLDSA87_BYTES != PSA_MLDSA_SIGNATURE_SIZE(87)
+#error "PSA and mldsa-native disagree on the ML-DSA-87 signature size"
+#endif
+
 /* For now, hard-coded values for MLDSA-87 */
 #define TF_PSA_CRYPTO_MLDSA_EXPANDED_SECRET_MAX_SIZE MLDSA87_SECRETKEYBYTES
 #define TF_PSA_CRYPTO_MLDSA_PUBLIC_KEY_MAX_SIZE MLDSA87_PUBLICKEYBYTES
 #define TF_PSA_CRYPTO_MLDSA_SIGNATURE_MAX_SIZE MLDSA87_BYTES
 
+#if TF_PSA_CRYPTO_MLDSA_SIGNATURE_MAX_SIZE < PSA_MLDSA_SIGNATURE_MAX_SIZE
+#error "PSA and mldsa-native disagree on the maximum ML-DSA signature size"
+#endif
+
 static psa_status_t pqcp_to_psa_error(int ret)
 {
     /* At the time of writing, mldsa-native has very few documented error
@@ -193,4 +272,243 @@ psa_status_t tf_psa_crypto_mldsa_verify_message(
     }
 }
 
+static psa_status_t setup(
+    tf_psa_crypto_mldsa_operation_t *operation,
+    const psa_key_attributes_t *attributes,
+    psa_algorithm_t alg)
+{
+    memset(operation, 0, sizeof(*operation));
+
+    if (psa_get_key_bits(attributes) != 87) {
+        /* Other parameter sets are not supported yet. */
+        return PSA_ERROR_NOT_SUPPORTED;
+    }
+    operation->parameter_set = psa_get_key_bits(attributes);
+
+    switch (alg) {
+        case PSA_ALG_DETERMINISTIC_ML_DSA:
+            operation->hedged = 0;
+            break;
+        case PSA_ALG_ML_DSA:
+            operation->hedged = 1;
+            break;
+        default:
+            return PSA_ERROR_NOT_SUPPORTED;
+    }
+
+    return PSA_SUCCESS;
+}
+
+static void start_pure(tf_psa_crypto_mldsa_shake256_t *shake_ctx,
+                       const uint8_t *public_key, size_t public_key_length)
+{
+    /* Hash the public key */
+    mld_shake256_init(shake_ctx);
+    mld_shake256_absorb(shake_ctx, public_key, public_key_length);
+    mld_shake256_finalize(shake_ctx);
+    uint8_t tr[MLDSA_CRHBYTES];
+    mld_shake256_squeeze(tr, sizeof(tr), shake_ctx);
+    mld_shake256_release(shake_ctx);
+    mld_shake256_init(shake_ctx);
+    mld_shake256_absorb(shake_ctx, tr, sizeof(tr));
+
+    /* Hash the domain separation prefix */
+    tr[0] = 0;                  /* pure ML-DSA (1 for Hash-ML-DSA) */
+    tr[1] = 0;                  /* context length */
+    mld_shake256_absorb(shake_ctx, tr, 2);
+}
+
+psa_status_t tf_psa_crypto_mldsa_sign_setup(
+    tf_psa_crypto_mldsa_operation_t *operation,
+    const psa_key_attributes_t *attributes,
+    const uint8_t *key_buffer, size_t key_buffer_size,
+    psa_algorithm_t alg)
+{
+    psa_status_t status = setup(operation, attributes, alg);
+    if (status != PSA_SUCCESS) {
+        return status;
+    }
+
+    if (operation->hedged) {
+        /* not implemented yet */
+        return PSA_ERROR_NOT_SUPPORTED;
+    }
+
+    if (psa_get_key_type(attributes) != PSA_KEY_TYPE_ML_DSA_KEY_PAIR) {
+        return PSA_ERROR_INVALID_ARGUMENT;
+    }
+    if (key_buffer_size != SEED_SIZE) {
+        return PSA_ERROR_INVALID_ARGUMENT;
+    }
+
+    /* After this point, we may allocate memory, so we must go through
+     * cleanup. */
+    size_t public_key_length = MLDSA87_PUBLICKEYBYTES;
+    uint8_t *public_key = mbedtls_calloc(1, public_key_length);
+    if (public_key == NULL) {
+        status = PSA_ERROR_INSUFFICIENT_MEMORY;
+        goto cleanup;
+    }
+    operation->key = mbedtls_calloc(1, MLDSA87_SECRETKEYBYTES);
+    if (operation->key == NULL) {
+        status = PSA_ERROR_INSUFFICIENT_MEMORY;
+        goto cleanup;
+    }
+    operation->key_length = MLDSA87_SECRETKEYBYTES;
+
+    int ret = tf_psa_crypto_pqcp_mldsa87_keypair_internal(public_key,
+                                                          operation->key,
+                                                          key_buffer);
+    if (ret != 0) {
+        status = pqcp_to_psa_error(ret);
+        goto cleanup;
+    }
+
+    start_pure(&operation->shake, public_key, public_key_length);
+
+cleanup:
+    mbedtls_free(public_key);
+    if (status != PSA_SUCCESS) {
+        mbedtls_zeroize_and_free(operation->key, operation->key_length);
+        mld_shake256_release(&operation->shake);
+        mbedtls_platform_zeroize(operation, sizeof(*operation));
+    }
+    return status;
+}
+
+psa_status_t tf_psa_crypto_mldsa_verify_setup(
+    tf_psa_crypto_mldsa_operation_t *operation,
+    const psa_key_attributes_t *attributes,
+    const uint8_t *key_buffer, size_t key_buffer_size,
+    psa_algorithm_t alg)
+{
+    psa_status_t status = setup(operation, attributes, alg);
+    if (status != PSA_SUCCESS) {
+        return status;
+    }
+
+    if (psa_get_key_type(attributes) != PSA_KEY_TYPE_ML_DSA_PUBLIC_KEY) {
+        return PSA_ERROR_INVALID_ARGUMENT;
+    }
+    if (key_buffer_size != MLDSA87_PUBLICKEYBYTES) {
+        /* Technically setup() doesn't care about the public key size, only
+         * finish() will care. But it's easier for users to debug a wrong-key
+         * problem if we complain as soon as the problem is noticeable. */
+        return PSA_ERROR_INVALID_ARGUMENT;
+    }
+
+    start_pure(&operation->shake, key_buffer, key_buffer_size);
+
+    return PSA_SUCCESS;
+}
+
+psa_status_t tf_psa_crypto_mldsa_update(
+    tf_psa_crypto_mldsa_operation_t *operation,
+    const uint8_t *input, size_t input_length)
+{
+    mld_shake256_absorb(&operation->shake, input, input_length);
+    return PSA_SUCCESS;
+}
+
+psa_status_t tf_psa_crypto_mldsa_sign_finish(
+    tf_psa_crypto_mldsa_operation_t *operation,
+    const uint8_t *key_buffer, size_t key_buffer_size,
+    uint8_t *signature, size_t signature_size, size_t *signature_length)
+{
+    *signature_length = 0;
+
+    if (operation->parameter_set != 87) {
+        return PSA_ERROR_NOT_SUPPORTED;
+    }
+    if (signature_size < MLDSA87_BYTES) {
+        return PSA_ERROR_BUFFER_TOO_SMALL;
+    }
+
+    /* Rely on setup() having stored the expanded private key in the
+     * operation structure. This is a performance/memory trade-off:
+     * we could instead re-expand the private key from the seed
+     * in \p key_buffer here. */
+    if (operation->key_length != MLDSA87_SECRETKEYBYTES) {
+        return PSA_ERROR_CORRUPTION_DETECTED;
+    }
+    (void) key_buffer;
+    (void) key_buffer_size;
+
+    uint8_t mu[MLDSA_CRHBYTES];
+    mld_shake256_finalize(&operation->shake);
+    mld_shake256_squeeze(mu, sizeof(mu), &operation->shake);
+    mld_shake256_release(&operation->shake);
+
+    uint8_t rnd[MLDSA_RNDBYTES];
+    memset(rnd, 0, sizeof(rnd));
+
+    int ret = tf_psa_crypto_pqcp_mldsa87_signature_internal(
+        signature, signature_length,
+        mu, sizeof(mu),
+        NULL, 0, rnd,
+        operation->key, 1);
+
+    psa_status_t abort_status = tf_psa_crypto_mldsa_abort(operation);
+    if (abort_status != PSA_SUCCESS) {
+        return abort_status;
+    }
+    return pqcp_to_psa_error(ret);
+}
+
+psa_status_t tf_psa_crypto_mldsa_verify_finish(
+    tf_psa_crypto_mldsa_operation_t *operation,
+    const uint8_t *key_buffer, size_t key_buffer_size,
+    const uint8_t *signature, size_t signature_length)
+{
+    if (operation->parameter_set != 87) {
+        return PSA_ERROR_NOT_SUPPORTED;
+    }
+    if (key_buffer_size != MLDSA87_PUBLICKEYBYTES) {
+        return PSA_ERROR_INVALID_ARGUMENT;
+    }
+    if (signature_length != MLDSA87_BYTES) {
+        return PSA_ERROR_INVALID_SIGNATURE;
+    }
+
+    uint8_t mu[MLDSA_CRHBYTES];
+    mld_shake256_finalize(&operation->shake);
+    mld_shake256_squeeze(mu, sizeof(mu), &operation->shake);
+    mld_shake256_release(&operation->shake);
+
+    int ret = tf_psa_crypto_pqcp_mldsa87_verify_internal(
+        signature, signature_length,
+        mu, sizeof(mu),
+        NULL, 0,
+        key_buffer, 1);
+
+    psa_status_t abort_status = tf_psa_crypto_mldsa_abort(operation);
+    if (abort_status != PSA_SUCCESS) {
+        return abort_status;
+    }
+
+    if (ret == MLD_ERR_FAIL) {
+        return PSA_ERROR_INVALID_SIGNATURE;
+    } else {
+        return pqcp_to_psa_error(ret);
+    }
+}
+
+psa_status_t tf_psa_crypto_mldsa_abort(
+    tf_psa_crypto_mldsa_operation_t *operation)
+{
+    /* If operation->parameter_set is 0, we may have an operation object
+     * that's only partially initialized. This shouldn't happen, since
+     * the PSA crypto driver specification says that the core initialized
+     * driver contexts to all-bits-zero. But avoid calling free() in that
+     * case as an extra bit of robustness. Of course, if the operation
+     * object is completely uninitialized, there's no way to detect that.
+     */
+    if (operation->parameter_set != 0) {
+        mbedtls_zeroize_and_free(operation->key, operation->key_length);
+        mld_shake256_release(&operation->shake);
+    }
+    mbedtls_platform_zeroize(operation, sizeof(*operation));
+    return PSA_SUCCESS;
+}
+
 #endif /* MBEDTLS_PSA_CRYPTO_C && TF_PSA_CRYPTO_PQCP_MLDSA_ENABLED */