Fix ML-KEM ARM64 NEON ciphertext comparison reduction

ZD#21457 (30)

Author: mattia-moffa

tests/api/test_mlkem.c

@@ -3950,3 +3950,67 @@ int test_wc_mlkem_decapsulate_pubonly_fails(void)
     return EXPECT_RESULT();
 } /* END test_wc_mlkem_decapsulate_pubonly_fails */
 
+/* Verify that the FO re-encryption check catches ciphertext tampering
+ * at various byte offsets and falls back to implicit rejection. */
+int test_wc_mlkem_decap_fo_reject(void)
+{
+    EXPECT_DECLS;
+#if !defined(HAVE_FIPS) || FIPS_VERSION3_GE(7,0,0)
+#if defined(WOLFSSL_HAVE_MLKEM) && defined(WOLFSSL_WC_MLKEM) && \
+    !defined(WOLFSSL_NO_ML_KEM) && !defined(WOLFSSL_MLKEM_NO_DECAPSULATE) && \
+    !defined(WOLFSSL_MLKEM_NO_ENCAPSULATE) && \
+    !defined(WOLFSSL_MLKEM_NO_MAKE_KEY)
+    MlKemKey* key = NULL;
+    WC_RNG rng;
+    byte ct[WC_ML_KEM_MAX_CIPHER_TEXT_SIZE];
+    byte ctTampered[WC_ML_KEM_MAX_CIPHER_TEXT_SIZE];
+    byte ss[WC_ML_KEM_SS_SZ];
+    byte ssDec[WC_ML_KEM_SS_SZ];
+    byte ssTampered[WC_ML_KEM_SS_SZ];
+    word32 ctLen = 0;
+
+    key = (MlKemKey*)XMALLOC(sizeof(*key), NULL, DYNAMIC_TYPE_TMP_BUFFER);
+    ExpectNotNull(key);
+
+    XMEMSET(&rng, 0, sizeof(rng));
+    ExpectIntEQ(wc_InitRng(&rng), 0);
+
+#ifndef WOLFSSL_NO_ML_KEM_768
+    ExpectIntEQ(wc_MlKemKey_Init(key, WC_ML_KEM_768, NULL, INVALID_DEVID), 0);
+#elif !defined(WOLFSSL_NO_ML_KEM_512)
+    ExpectIntEQ(wc_MlKemKey_Init(key, WC_ML_KEM_512, NULL, INVALID_DEVID), 0);
+#else
+    ExpectIntEQ(wc_MlKemKey_Init(key, WC_ML_KEM_1024, NULL, INVALID_DEVID), 0);
+#endif
+
+    ExpectIntEQ(wc_MlKemKey_CipherTextSize(key, &ctLen), 0);
+    ExpectIntEQ(wc_MlKemKey_MakeKey(key, &rng), 0);
+    ExpectIntEQ(wc_MlKemKey_Encapsulate(key, ct, ss, &rng), 0);
+
+    /* Untampered ciphertext recovers the original ss. */
+    XMEMSET(ssDec, 0, sizeof(ssDec));
+    ExpectIntEQ(wc_MlKemKey_Decapsulate(key, ssDec, ct, ctLen), 0);
+    ExpectIntEQ(XMEMCMP(ssDec, ss, WC_ML_KEM_SS_SZ), 0);
+
+    /* Tamper at byte 32: implicit rejection must fire. */
+    XMEMCPY(ctTampered, ct, ctLen);
+    ctTampered[32] ^= 0x01;
+    XMEMSET(ssTampered, 0, sizeof(ssTampered));
+    ExpectIntEQ(wc_MlKemKey_Decapsulate(key, ssTampered, ctTampered, ctLen), 0);
+    ExpectIntNE(XMEMCMP(ssTampered, ss, WC_ML_KEM_SS_SZ), 0);
+
+    /* Tamper at byte 0: also must be rejected. */
+    XMEMCPY(ctTampered, ct, ctLen);
+    ctTampered[0] ^= 0x01;
+    XMEMSET(ssTampered, 0, sizeof(ssTampered));
+    ExpectIntEQ(wc_MlKemKey_Decapsulate(key, ssTampered, ctTampered, ctLen), 0);
+    ExpectIntNE(XMEMCMP(ssTampered, ss, WC_ML_KEM_SS_SZ), 0);
+
+    DoExpectIntEQ(wc_FreeRng(&rng), 0);
+    wc_MlKemKey_Free(key);
+    XFREE(key, NULL, DYNAMIC_TYPE_TMP_BUFFER);
+#endif
+#endif
+    return EXPECT_RESULT();
+} /* END test_wc_mlkem_decap_fo_reject */
+

tests/api/test_mlkem.h

@@ -28,11 +28,13 @@ int test_wc_mlkem_make_key_kats(void);
 int test_wc_mlkem_encapsulate_kats(void);
 int test_wc_mlkem_decapsulate_kats(void);
 int test_wc_mlkem_decapsulate_pubonly_fails(void);
+int test_wc_mlkem_decap_fo_reject(void);
 
-#define TEST_MLKEM_DECLS                                      \
-    TEST_DECL_GROUP("mlkem", test_wc_mlkem_make_key_kats),    \
-    TEST_DECL_GROUP("mlkem", test_wc_mlkem_encapsulate_kats), \
-    TEST_DECL_GROUP("mlkem", test_wc_mlkem_decapsulate_kats), \
-    TEST_DECL_GROUP("mlkem", test_wc_mlkem_decapsulate_pubonly_fails)
+#define TEST_MLKEM_DECLS                                                \
+    TEST_DECL_GROUP("mlkem", test_wc_mlkem_make_key_kats),              \
+    TEST_DECL_GROUP("mlkem", test_wc_mlkem_encapsulate_kats),           \
+    TEST_DECL_GROUP("mlkem", test_wc_mlkem_decapsulate_kats),           \
+    TEST_DECL_GROUP("mlkem", test_wc_mlkem_decapsulate_pubonly_fails),  \
+    TEST_DECL_GROUP("mlkem", test_wc_mlkem_decap_fo_reject)
 
 #endif /* WOLFCRYPT_TEST_MLKEM_H */

wolfcrypt/src/port/arm/armv8-mlkem-asm.S

@@ -8927,7 +8927,7 @@ L_mlkem_aarch64_cmp_neon_done:
 	orr	v8.16b, v8.16b, v9.16b
 	orr	v10.16b, v10.16b, v11.16b
 	orr	v8.16b, v8.16b, v10.16b
-	ins	v9.b[0], v8.b[1]
+	ext	v9.16b, v8.16b, v8.16b, #8
 	orr	v8.16b, v8.16b, v9.16b
 	mov	x0, v8.d[0]
 	subs	x0, x0, xzr

wolfcrypt/src/port/arm/armv8-mlkem-asm_c.c

@@ -8404,7 +8404,7 @@ int mlkem_cmp_neon(const byte* a, const byte* b, int sz)
         "orr	v8.16b, v8.16b, v9.16b\n\t"
         "orr	v10.16b, v10.16b, v11.16b\n\t"
         "orr	v8.16b, v8.16b, v10.16b\n\t"
-        "ins	v9.b[0], v8.b[1]\n\t"
+        "ext	v9.16b, v8.16b, v8.16b, #8\n\t"
         "orr	v8.16b, v8.16b, v9.16b\n\t"
         "mov	x0, v8.d[0]\n\t"
         "subs	x0, x0, xzr\n\t"