diff --git a/sdk/software/examples/mlkem_test/main.c b/sdk/software/examples/mlkem_test/main.c index 45951ef..0c4ff11 100644 --- a/sdk/software/examples/mlkem_test/main.c +++ b/sdk/software/examples/mlkem_test/main.c @@ -49,9 +49,19 @@ unsigned long CORE_CLOCKS_PER_SEC = 33000000L; #define MLKEM_DK_WINDOW (MLKEM_BASE + 0x2000) /* W preload dk, R read dk */ #define MLKEM_CT_WINDOW (MLKEM_BASE + 0x3000) /* W preload c, R read ct */ +/* PRNG sub-window in the shared ML-KEM/FFT/PRNG AXI slave. */ +#define PRNG_BASE (MLKEM_BASE + 0xD000) +#define PRNG_CTRL (PRNG_BASE + 0x000) /* W bit0=mix seed bit1=clr count */ +#define PRNG_STATUS (PRNG_BASE + 0x004) /* R bit0=ready */ +#define PRNG_DATA (PRNG_BASE + 0x008) /* R next pseudo-random word */ +#define PRNG_COUNT (PRNG_BASE + 0x00C) /* R data read count */ +#define PRNG_SEED (PRNG_BASE + 0x010) /* RW software seed mix value */ + #define CTRL_START 0x1 #define STATUS_BUSY 0x1 #define STATUS_DONE 0x2 +#define PRNG_READY 0x1 +#define PRNG_CTRL_MIX 0x1 #define OP_KEYGEN 0 #define OP_ENCAPS 1 @@ -132,6 +142,44 @@ static void read_block32(uint32_t base, uint8_t *bytes) } } +/* Filled from the PRNG endpoint at runtime. */ +static uint8_t SEED_D[32]; +static uint8_t SEED_Z[32]; +static uint8_t MSG_M[32]; + +static uint32_t prng_read_word(void) +{ + while ((REG_RD(PRNG_STATUS) & PRNG_READY) == 0) + ; + return REG_RD(PRNG_DATA); +} + +static void prng_mix_seed(uint32_t seed) +{ + REG_WR(PRNG_SEED, seed); + REG_WR(PRNG_CTRL, PRNG_CTRL_MIX); +} + +static void prng_fill(uint8_t *bytes, int n) +{ + int i; + for (i = 0; i < n; i += 4) { + uint32_t v = prng_read_word(); + bytes[i + 0] = (uint8_t)(v & 0xFF); + if (i + 1 < n) bytes[i + 1] = (uint8_t)((v >> 8) & 0xFF); + if (i + 2 < n) bytes[i + 2] = (uint8_t)((v >> 16) & 0xFF); + if (i + 3 < n) bytes[i + 3] = (uint8_t)((v >> 24) & 0xFF); + } +} + +static void generate_prng_inputs(void) +{ + prng_mix_seed((uint32_t)get_clock_count()); + prng_fill(SEED_D, sizeof(SEED_D)); + prng_fill(SEED_Z, sizeof(SEED_Z)); + prng_fill(MSG_M, sizeof(MSG_M)); +} + /* Stream `n` bytes into a preload port. Each PRE_DAT write auto-increments * the preload address; the bus wait-states each write until the cross-domain * byte handshake completes, so a plain loop is correct. */ @@ -198,20 +246,6 @@ static void copy_from_uc(uint8_t *dst, volatile uint8_t *src, int n) dst[i] = src[i]; } -/* Arbitrary but fixed test seeds / message (32 bytes each). */ -static uint8_t SEED_D[32] = { - 0x00,0x01,0x02,0x03,0x04,0x05,0x06,0x07,0x08,0x09,0x0A,0x0B,0x0C,0x0D,0x0E,0x0F, - 0x10,0x11,0x12,0x13,0x14,0x15,0x16,0x17,0x18,0x19,0x1A,0x1B,0x1C,0x1D,0x1E,0x1F -}; -static uint8_t SEED_Z[32] = { - 0x20,0x21,0x22,0x23,0x24,0x25,0x26,0x27,0x28,0x29,0x2A,0x2B,0x2C,0x2D,0x2E,0x2F, - 0x30,0x31,0x32,0x33,0x34,0x35,0x36,0x37,0x38,0x39,0x3A,0x3B,0x3C,0x3D,0x3E,0x3F -}; -static uint8_t MSG_M[32] = { - 0xA0,0xA1,0xA2,0xA3,0xA4,0xA5,0xA6,0xA7,0xA8,0xA9,0xAA,0xAB,0xAC,0xAD,0xAE,0xAF, - 0xB0,0xB1,0xB2,0xB3,0xB4,0xB5,0xB6,0xB7,0xB8,0xB9,0xBA,0xBB,0xBC,0xBD,0xBE,0xBF -}; - typedef struct { uint8_t ek[EK_BYTES] __attribute__((aligned(64))); uint8_t dk[DK_BYTES] __attribute__((aligned(64))); @@ -343,6 +377,7 @@ int main(void) unsigned long cpu_cycles, dma_cycles; printf("MLKEM_START\n"); + generate_prng_inputs(); cpu_cycles = run_roundtrip_cpu(SEED_D, SEED_Z, MSG_M, &cpu_res); dma_cycles = run_roundtrip_dma(SEED_D, SEED_Z, MSG_M, &dma_res);