Add unit tests for FIPS202 APIs

Author: bremoran

mldsa/fips202/fips202.h

@@ -17,14 +17,18 @@
 #define SHA3_256_HASHBYTES 32
 #define SHA3_512_HASHBYTES 64
 
+#ifndef FIPS202_NAMESPACE
 #define FIPS202_NAMESPACE(s) mldsa_fips202_ref_##s
+#endif
 
+#define mld_shake128ctx FIPS202_NAMESPACE(shake128ctx)
 typedef struct
 {
   uint64_t s[MLD_KECCAK_LANES];
   unsigned int pos;
 } mld_shake128ctx;
 
+#define mld_shake256ctx FIPS202_NAMESPACE(shake256ctx)
 typedef struct
 {
   uint64_t s[MLD_KECCAK_LANES];

test/test_unit.c

@@ -14,10 +14,13 @@
 
 
 #define FIPS202_NAMESPACE(A) test_priv_##A
-#define MLD_FIPS202_FIPS202_H
 #define MLD_COMMON_H
+#define mld_memset memset
+#define mld_memcpy memcpy
+#define mld_zeroize(PTR, LEN) memset(PTR, 0, LEN)
 
 #include "../mldsa/fips202/keccakf1600.c"
+#include "../mldsa/fips202/fips202.c"
 
 #undef FIPS202_NAMESPACE
 #undef mld_keccakf1600_xor_bytes
@@ -26,13 +29,35 @@
 #undef mld_keccakf1600x4_extract_bytes
 #undef mld_keccakf1600_permute
 #undef mld_keccakf1600x4_permute
+#undef mld_shake256
+#undef mld_shake256_release
+#undef mld_shake256_squeeze
+#undef mld_shake256_finalize
+#undef mld_shake256_absorb
+#undef mld_shake256_init
+#undef mld_shake128_release
+#undef mld_shake128_squeeze
+#undef mld_shake128_finalize
+#undef mld_shake128_absorb
+#undef mld_shake128_init
+#undef mld_keccakf1600x4_permute
+#undef mld_keccakf1600_permute
+#undef mld_keccakf1600x4_xor_bytes
+#undef mld_keccakf1600x4_extract_bytes
+#undef mld_keccakf1600_xor_bytes
+#undef mld_keccakf1600_extract_bytes
+
+#undef mld_memset
+#undef mld_memcpy
+#undef mld_zeroize
 
 #undef MLD_FIPS202_KECCAKF1600_H
 #undef MLD_FIPS202_FIPS202_H
 #undef MLD_COMMON_H
 
 /* Under test */
 #include "../mldsa/fips202/keccakf1600.h"
+#include "../mldsa/fips202/fips202.h"
 
 #define CHECK(x)                                              \
   do                                                          \
@@ -49,7 +74,7 @@
 #define STATE_BYTES (MLD_KECCAK_LANES * sizeof(uint64_t))
 
 #if defined(MLD_USE_FIPS202_X1_XOR_NATIVE) || defined(MLD_USE_FIPS202_X4_XOR_NATIVE)
-/* Bit interleave/deinterleave helpers as specified */
+/* Bit interleave/deinterleave helpers */
 static uint64_t test_priv_bit_interleave(uint64_t in)
 {
   uint64_t e = in & 0x5555555555555555ULL;
@@ -155,6 +180,141 @@ static void unpack_x4_native(uint64_t ch_interleaved_out[MLD_KECCAK_WAY][MLD_KEC
 }
 #endif /* MLD_USE_FIPS202_X4_NATIVE && MLD_USE_FIPS202_X4_XOR_NATIVE */
 
+static int shake128_ctx_equals(const test_priv_shake128ctx *ref,
+                               const mld_shake128ctx *dut)
+{
+  if (ref->pos != dut->pos) return 0;
+#if defined(MLD_USE_FIPS202_X1_XOR_NATIVE)
+  uint64_t norm[MLD_KECCAK_LANES];
+  deinterleave_state(norm, dut->s);
+  return memcmp(ref->s, norm, sizeof norm) == 0;
+#else
+  return memcmp(ref->s, dut->s, sizeof ref->s) == 0;
+#endif
+}
+
+static int shake256_ctx_equals(const test_priv_shake256ctx *ref,
+                               const mld_shake256ctx *dut)
+{
+  if (ref->pos != dut->pos) return 0;
+#if defined(MLD_USE_FIPS202_X1_XOR_NATIVE)
+  uint64_t norm[MLD_KECCAK_LANES];
+  deinterleave_state(norm, dut->s);
+  return memcmp(ref->s, norm, sizeof norm) == 0;
+#else
+  return memcmp(ref->s, dut->s, sizeof ref->s) == 0;
+#endif
+}
+
+static int test_shake128_api(void)
+{
+  int fails = 0;
+  static const size_t absorb_lens[] = {0, 1, SHAKE128_RATE - 1, SHAKE128_RATE,
+                                       SHAKE128_RATE + 7, SHAKE128_RATE * 2 + 5};
+  static const size_t squeeze_lens[] = {0, 1, 32, SHAKE128_RATE,
+                                        SHAKE128_RATE + 11, SHAKE128_RATE * 3};
+  uint8_t chunk[SHAKE128_RATE * 3];
+  uint8_t out_ref[SHAKE128_RATE * 4];
+  uint8_t out_dut[SHAKE128_RATE * 4];
+  test_priv_shake128ctx ref_ctx;
+  mld_shake128ctx dut_ctx;
+
+  test_priv_shake128_init(&ref_ctx);
+  mld_shake128_init(&dut_ctx);
+  CHECK(shake128_ctx_equals(&ref_ctx, &dut_ctx));
+
+  for (size_t i = 0; i < sizeof(absorb_lens) / sizeof(absorb_lens[0]); i++) {
+    size_t len = absorb_lens[i];
+    if (len > 0) randombytes(chunk, len);
+    test_priv_shake128_absorb(&ref_ctx, chunk, len);
+    mld_shake128_absorb(&dut_ctx, chunk, len);
+    CHECK(shake128_ctx_equals(&ref_ctx, &dut_ctx));
+  }
+
+  test_priv_shake128_finalize(&ref_ctx);
+  mld_shake128_finalize(&dut_ctx);
+  CHECK(shake128_ctx_equals(&ref_ctx, &dut_ctx));
+
+  for (size_t i = 0; i < sizeof(squeeze_lens) / sizeof(squeeze_lens[0]); i++) {
+    size_t len = squeeze_lens[i];
+    memset(out_ref, 0, len);
+    memset(out_dut, 0, len);
+    test_priv_shake128_squeeze(out_ref, len, &ref_ctx);
+    mld_shake128_squeeze(out_dut, len, &dut_ctx);
+    CHECK(memcmp(out_ref, out_dut, len) == 0);
+    CHECK(shake128_ctx_equals(&ref_ctx, &dut_ctx));
+  }
+
+  return fails;
+}
+
+static int test_shake256_api(void)
+{
+  int fails = 0;
+  static const size_t absorb_lens[] = {0, 1, SHAKE256_RATE - 1, SHAKE256_RATE,
+                                       SHAKE256_RATE + 9, SHAKE256_RATE * 2 + 3};
+  static const size_t squeeze_lens[] = {0, 1, 48, SHAKE256_RATE,
+                                        SHAKE256_RATE + 13, SHAKE256_RATE * 3};
+  uint8_t chunk[SHAKE256_RATE * 3];
+  uint8_t out_ref[SHAKE256_RATE * 4];
+  uint8_t out_dut[SHAKE256_RATE * 4];
+  test_priv_shake256ctx ref_ctx;
+  mld_shake256ctx dut_ctx;
+
+  test_priv_shake256_init(&ref_ctx);
+  mld_shake256_init(&dut_ctx);
+  CHECK(shake256_ctx_equals(&ref_ctx, &dut_ctx));
+
+  for (size_t i = 0; i < sizeof(absorb_lens) / sizeof(absorb_lens[0]); i++) {
+    size_t len = absorb_lens[i];
+    if (len > 0) randombytes(chunk, len);
+    test_priv_shake256_absorb(&ref_ctx, chunk, len);
+    mld_shake256_absorb(&dut_ctx, chunk, len);
+    CHECK(shake256_ctx_equals(&ref_ctx, &dut_ctx));
+  }
+
+  test_priv_shake256_finalize(&ref_ctx);
+  mld_shake256_finalize(&dut_ctx);
+  CHECK(shake256_ctx_equals(&ref_ctx, &dut_ctx));
+
+  for (size_t i = 0; i < sizeof(squeeze_lens) / sizeof(squeeze_lens[0]); i++) {
+    size_t len = squeeze_lens[i];
+    memset(out_ref, 0, len);
+    memset(out_dut, 0, len);
+    test_priv_shake256_squeeze(out_ref, len, &ref_ctx);
+    mld_shake256_squeeze(out_dut, len, &dut_ctx);
+    CHECK(memcmp(out_ref, out_dut, len) == 0);
+    CHECK(shake256_ctx_equals(&ref_ctx, &dut_ctx));
+  }
+
+  return fails;
+}
+
+static int test_shake256_function(void)
+{
+  int fails = 0;
+  static const size_t in_lens[] = {0, 1, 50, 200};
+  static const size_t out_lens[] = {0, 1, 64, 256};
+  uint8_t inbuf[200];
+  uint8_t out_ref[256];
+  uint8_t out_dut[256];
+
+  for (size_t i = 0; i < sizeof(in_lens) / sizeof(in_lens[0]); i++) {
+    size_t inlen = in_lens[i];
+    if (inlen > 0) randombytes(inbuf, inlen);
+    for (size_t j = 0; j < sizeof(out_lens) / sizeof(out_lens[0]); j++) {
+      size_t outlen = out_lens[j];
+      memset(out_ref, 0, outlen);
+      memset(out_dut, 0, outlen);
+      test_priv_shake256(out_ref, outlen, inbuf, inlen);
+      mld_shake256(out_dut, outlen, inbuf, inlen);
+      CHECK(memcmp(out_ref, out_dut, outlen) == 0);
+    }
+  }
+
+  return fails;
+}
+
 #if defined(MLD_USE_FIPS202_X4_NATIVE)
 static int test_x4_xor_bytes(void)
 {
@@ -421,6 +581,10 @@ int main(void)
 
   randombytes_reset();
 
+  fails += test_shake128_api();
+  fails += test_shake256_api();
+  fails += test_shake256_function();
+
 #if defined(MLD_USE_FIPS202_X1_NATIVE)
   fails += test_xor_bytes();
   fails += test_extract_bytes();