Preserve in-place message on signing failure

mldsa/mldsa_native.h

@@ -444,6 +444,9 @@ int MLD_API_NAMESPACE(signature_extmu)(
  *                     MLD_CONFIG_CONTEXT_PARAMETER is defined; type set by
  *                     MLD_CONFIG_CONTEXT_PARAMETER_TYPE.
  *
+ * On failure, *smlen is set to 0. No partial signed-message output is
+ * preserved, but if sm == m, the input message bytes remain unchanged.
+ *
  * @retval 0                               Success.
  * @retval MLD_ERR_OUT_OF_MEMORY           MLD_CONFIG_CUSTOM_ALLOC_FREE was
  *                                         used and an allocation via

mldsa/src/sign.c

@@ -1151,8 +1151,17 @@ int mld_sign(uint8_t *sm, size_t *smlen, const uint8_t *m, size_t mlen,
              MLD_CONFIG_CONTEXT_PARAMETER_TYPE context)
 {
   int ret;
+  uint8_t sig[MLDSA_CRYPTO_BYTES];
+  size_t siglen;
   size_t i;
 
+  ret = mld_sign_signature(sig, &siglen, m, mlen, ctx, ctxlen, sk, context);
+  if (ret != 0)
+  {
+    *smlen = 0;
+    goto cleanup;
+  }
+
   for (i = 0; i < mlen; ++i)
   __loop__(
     assigns(i, object_whole(sm))
@@ -1162,12 +1171,12 @@ int mld_sign(uint8_t *sm, size_t *smlen, const uint8_t *m, size_t mlen,
   {
     sm[MLDSA_CRYPTO_BYTES + mlen - 1 - i] = m[mlen - 1 - i];
   }
-  ret = mld_sign_signature(sm, smlen, sm + MLDSA_CRYPTO_BYTES, mlen, ctx,
-                           ctxlen, sk, context);
-  if (ret == 0)
-  {
-    *smlen += mlen;
-  }
+
+  mld_memcpy(sm, sig, MLDSA_CRYPTO_BYTES);
+  *smlen = siglen + mlen;
+
+cleanup:
+  mld_zeroize(sig, sizeof(sig));
   return ret;
 }
 #endif /* !MLD_CONFIG_NO_RANDOMIZED_API */

mldsa/src/sign.h

@@ -259,6 +259,9 @@ __contract__(
  *                     MLD_CONFIG_CONTEXT_PARAMETER is defined; type set by
  *                     MLD_CONFIG_CONTEXT_PARAMETER_TYPE.
  *
+ * On failure, *smlen is set to 0. No partial signed-message output is
+ * preserved, but if sm == m, the input message bytes remain unchanged.
+ *
  * @retval 0                               Success.
  * @retval MLD_ERR_OUT_OF_MEMORY           MLD_CONFIG_CUSTOM_ALLOC_FREE was
  *                                         used and an allocation via

test/src/test_mldsa.c

@@ -134,6 +134,73 @@ static int test_sign_unaligned(void)
   return test_sign_core(pk + 1, sk + 1, sm + 1, m + 1, m2 + 1, ctx + 1);
 }
 
+static int test_sign_inplace_success(void)
+{
+  uint8_t pk[CRYPTO_PUBLICKEYBYTES];
+  uint8_t sk[CRYPTO_SECRETKEYBYTES];
+  uint8_t sm[MLEN + CRYPTO_BYTES];
+  uint8_t expected[MLEN];
+  uint8_t m2[MLEN + CRYPTO_BYTES];
+  uint8_t ctx[CTXLEN];
+  size_t smlen;
+  size_t mlen;
+  int rc;
+
+  CHECK(crypto_sign_keypair(pk, sk) == 0);
+  CHECK(randombytes(ctx, sizeof(ctx)) == 0);
+  MLD_CT_TESTING_SECRET(ctx, sizeof(ctx));
+  CHECK(randombytes(sm, MLEN) == 0);
+  MLD_CT_TESTING_SECRET(sm, MLEN);
+  memcpy(expected, sm, MLEN);
+  MLD_CT_TESTING_SECRET(expected, sizeof(expected));
+
+  CHECK_SIGN_RC(crypto_sign(sm, &smlen, sm, MLEN, ctx, sizeof(ctx), sk));
+  rc = crypto_sign_open(m2, &mlen, sm, smlen, ctx, sizeof(ctx), pk);
+
+  MLD_CT_TESTING_DECLASSIFY(&rc, sizeof(rc));
+  MLD_CT_TESTING_DECLASSIFY(expected, sizeof(expected));
+  MLD_CT_TESTING_DECLASSIFY(m2, sizeof(m2));
+
+  CHECK(rc == 0);
+  CHECK(smlen == MLEN + CRYPTO_BYTES);
+  CHECK(mlen == MLEN);
+  CHECK(memcmp(expected, m2, MLEN) == 0);
+
+  return 0;
+}
+
+static int test_sign_inplace_context_failure_preserves_message(void)
+{
+  uint8_t pk[CRYPTO_PUBLICKEYBYTES];
+  uint8_t sk[CRYPTO_SECRETKEYBYTES];
+  uint8_t sm[MLEN + CRYPTO_BYTES];
+  uint8_t expected[MLEN];
+  uint8_t ctx[256];
+  size_t smlen = MLEN + CRYPTO_BYTES;
+  int rc;
+
+  CHECK(crypto_sign_keypair(pk, sk) == 0);
+  CHECK(randombytes(ctx, sizeof(ctx)) == 0);
+  MLD_CT_TESTING_SECRET(ctx, sizeof(ctx));
+  CHECK(randombytes(sm, MLEN) == 0);
+  MLD_CT_TESTING_SECRET(sm, MLEN);
+  memcpy(expected, sm, MLEN);
+  MLD_CT_TESTING_SECRET(expected, sizeof(expected));
+
+  rc = crypto_sign(sm, &smlen, sm, MLEN, ctx, sizeof(ctx), sk);
+
+  MLD_CT_TESTING_DECLASSIFY(&rc, sizeof(rc));
+  MLD_CT_TESTING_DECLASSIFY(&smlen, sizeof(smlen));
+  MLD_CT_TESTING_DECLASSIFY(sm, MLEN);
+  MLD_CT_TESTING_DECLASSIFY(expected, sizeof(expected));
+
+  CHECK(rc == MLD_ERR_FAIL);
+  CHECK(smlen == 0);
+  CHECK(memcmp(sm, expected, MLEN) == 0);
+
+  return 0;
+}
+
 static int test_sign_extmu(void)
 {
   uint8_t pk[CRYPTO_PUBLICKEYBYTES];
@@ -509,6 +576,8 @@ int main(void)
     !defined(MLD_CONFIG_NO_VERIFY_API)
     r |= test_sign();
     r |= test_sign_unaligned();
+    r |= test_sign_inplace_success();
+    r |= test_sign_inplace_context_failure_preserves_message();
     r |= test_wrong_pk();
     r |= test_wrong_sig();
     r |= test_wrong_ctx();

test/src/test_rng_fail.c

@@ -158,6 +158,46 @@ static int test_sign_combined_rng_failure(void)
   return 0;
 }
 
+static int test_sign_combined_inplace_rng_failure_preserves_message(void)
+{
+  uint8_t sm[CRYPTO_BYTES + TEST_VECTOR_MSG_LEN];
+  uint8_t expected[TEST_VECTOR_MSG_LEN];
+  size_t smlen = CRYPTO_BYTES + TEST_VECTOR_MSG_LEN;
+  int rc;
+
+  memcpy(sm, TEST_VECTOR_MSG, TEST_VECTOR_MSG_LEN);
+  memcpy(expected, TEST_VECTOR_MSG, TEST_VECTOR_MSG_LEN);
+
+  reset_all();
+  randombytes_fail_on_counter = 0;
+  rc = crypto_sign(sm, &smlen, sm, TEST_VECTOR_MSG_LEN,
+                   (const uint8_t *)TEST_VECTOR_CTX, TEST_VECTOR_CTX_LEN,
+                   test_vector_sk);
+  if (rc != MLD_ERR_RNG_FAIL)
+  {
+    fprintf(stderr,
+            "ERROR: crypto_sign returned %d on in-place RNG failure "
+            "(expected %d)\n",
+            rc, MLD_ERR_RNG_FAIL);
+    return 1;
+  }
+  if (smlen != 0)
+  {
+    fprintf(stderr,
+            "ERROR: crypto_sign returned smlen=%zu on in-place RNG failure\n",
+            smlen);
+    return 1;
+  }
+  if (memcmp(sm, expected, TEST_VECTOR_MSG_LEN) != 0)
+  {
+    fprintf(stderr,
+            "ERROR: crypto_sign changed the in-place message on RNG failure\n");
+    return 1;
+  }
+
+  return 0;
+}
+
 static int test_signature_extmu_rng_failure(void)
 {
   uint8_t sig[CRYPTO_BYTES];
@@ -249,6 +289,7 @@ int main(void)
 #if !defined(MLD_CONFIG_NO_SIGN_API)
   r |= test_sign_rng_failure();
   r |= test_sign_combined_rng_failure();
+  r |= test_sign_combined_inplace_rng_failure_preserves_message();
   r |= test_signature_extmu_rng_failure();
   r |= test_signature_pre_hash_shake256_rng_failure();
 #endif /* !MLD_CONFIG_NO_SIGN_API */