/* * mlkem_hello - one ML-KEM hardware accelerator communication demo. * * Alice generates an ML-KEM keypair, Bob encapsulates a shared secret to * Alice's public key, and Alice decapsulates Bob's ciphertext. The demo uses * fixed KAT count=0 inputs so it is deterministic and does not require a PRNG. */ #include #include #include unsigned long UART_BASE = 0xbf000000; unsigned long CONFREG_TIMER_BASE = 0xbf20f100; unsigned long CONFREG_CLOCKS_PER_SEC = 50000000L; unsigned long CORE_CLOCKS_PER_SEC = 33000000L; #define MLKEM_BASE 0xbf400000 #define MLKEM_CTRL (MLKEM_BASE + 0x000) #define MLKEM_STATUS (MLKEM_BASE + 0x004) #define MLKEM_CONFIG (MLKEM_BASE + 0x008) #define MLKEM_D (MLKEM_BASE + 0x010) #define MLKEM_Z (MLKEM_BASE + 0x030) #define MLKEM_MSG (MLKEM_BASE + 0x050) #define MLKEM_SS (MLKEM_BASE + 0x070) #define MLKEM_PRE_SEL (MLKEM_BASE + 0x100) #define MLKEM_PRE_ADR (MLKEM_BASE + 0x104) #define MLKEM_PRE_DAT (MLKEM_BASE + 0x108) #define MLKEM_RDB_SEL (MLKEM_BASE + 0x10C) #define MLKEM_RDB_ADR (MLKEM_BASE + 0x110) #define MLKEM_RDB_DAT (MLKEM_BASE + 0x114) #define CTRL_START 0x1 #define CTRL_SOFT_RESET 0x2 #define STATUS_DONE 0x2 #define OP_KEYGEN 0 #define OP_ENCAPS 1 #define OP_DECAPS 2 #define PRE_EK 0 #define PRE_DK 1 #define PRE_C 2 #define RDB_EK 0 #define RDB_DK 2 #define RDB_CT 3 #define MLKEM_K 3 #define EK_BYTES (384 * MLKEM_K + 32) #define DK_BYTES (768 * MLKEM_K + 96) #define CT_BYTES (32 * 10 * MLKEM_K + 32 * 4) #define SYM_BYTES 32 #define APP_MSG "hello_world" #define APP_MSG_BYTES (sizeof(APP_MSG) - 1) #define REG_WR(addr, val) (*(volatile uint32_t *)(addr) = (uint32_t)(val)) #define REG_RD(addr) (*(volatile uint32_t *)(addr)) static uint8_t alice_ek[EK_BYTES] __attribute__((aligned(64))); static uint8_t alice_dk[DK_BYTES] __attribute__((aligned(64))); static uint8_t bob_ct[CT_BYTES] __attribute__((aligned(64))); static uint8_t bob_ss[SYM_BYTES]; static uint8_t alice_ss[SYM_BYTES]; static uint8_t app_cipher[APP_MSG_BYTES]; static uint8_t app_plain_out[APP_MSG_BYTES + 1]; /* Deterministic ML-KEM KAT count=0 inputs. */ static const uint8_t SEED_D[SYM_BYTES] = { 0x6d, 0xbb, 0xc4, 0x37, 0x51, 0x36, 0xdf, 0x3b, 0x07, 0xf7, 0xc7, 0x0e, 0x63, 0x9e, 0x22, 0x3e, 0x17, 0x7e, 0x7f, 0xd5, 0x3b, 0x16, 0x1b, 0x3f, 0x4d, 0x57, 0x79, 0x17, 0x94, 0xf1, 0x26, 0x24 }; static const uint8_t SEED_Z[SYM_BYTES] = { 0xf6, 0x96, 0x48, 0x40, 0x48, 0xec, 0x21, 0xf9, 0x6c, 0xf5, 0x0a, 0x56, 0xd0, 0x75, 0x9c, 0x44, 0x8f, 0x37, 0x79, 0x75, 0x2f, 0x03, 0x83, 0xd3, 0x74, 0x49, 0x69, 0x06, 0x94, 0xcf, 0x7a, 0x68 }; static const uint8_t MSG_M[SYM_BYTES] = { 0x20, 0xa7, 0xb7, 0xe1, 0x0f, 0x70, 0x49, 0x6c, 0xc3, 0x82, 0x20, 0xb9, 0x44, 0xde, 0xf6, 0x99, 0xbf, 0x14, 0xd1, 0x4e, 0x55, 0xcf, 0x4c, 0x90, 0xa1, 0x2c, 0x1b, 0x33, 0xfc, 0x80, 0xff, 0xff }; static void print_hex(const char *label, const uint8_t *buf, int n) { int i; printf("%s", label); for (i = 0; i < n; i++) printf("%02x", buf[i]); printf("\n"); } static void xor_with_shared_secret(uint8_t *dst, const uint8_t *src, int n, const uint8_t *ss) { int i; for (i = 0; i < n; i++) dst[i] = src[i] ^ ss[i % SYM_BYTES]; } static void write_block32(uint32_t base, const uint8_t *bytes) { int w; for (w = 0; w < 8; w++) { uint32_t v = (uint32_t)bytes[w * 4 + 0] | ((uint32_t)bytes[w * 4 + 1] << 8) | ((uint32_t)bytes[w * 4 + 2] << 16) | ((uint32_t)bytes[w * 4 + 3] << 24); REG_WR(base + w * 4, v); } } static void read_block32(uint32_t base, uint8_t *bytes) { int w; for (w = 0; w < 8; w++) { uint32_t v = REG_RD(base + w * 4); bytes[w * 4 + 0] = (uint8_t)(v & 0xFF); bytes[w * 4 + 1] = (uint8_t)((v >> 8) & 0xFF); bytes[w * 4 + 2] = (uint8_t)((v >> 16) & 0xFF); bytes[w * 4 + 3] = (uint8_t)((v >> 24) & 0xFF); } } static void clear_done(void) { REG_WR(MLKEM_STATUS, STATUS_DONE); } static void wait_done(void) { while ((REG_RD(MLKEM_STATUS) & STATUS_DONE) == 0) ; } static void mlkem_soft_reset(void) { REG_WR(MLKEM_CTRL, CTRL_SOFT_RESET); clear_done(); } static void run_op(uint32_t op) { REG_WR(MLKEM_CONFIG, ((op & 0x3) << 4) | (MLKEM_K & 0x7)); clear_done(); REG_WR(MLKEM_CTRL, CTRL_START); wait_done(); } static void preload(uint32_t sel, const uint8_t *buf, int n) { int i; REG_WR(MLKEM_PRE_SEL, sel); REG_WR(MLKEM_PRE_ADR, 0); for (i = 0; i < n; i++) REG_WR(MLKEM_PRE_DAT, buf[i]); } static void readback(uint32_t sel, uint8_t *buf, int n) { int i; REG_WR(MLKEM_RDB_SEL, sel); for (i = 0; i < n; i++) { REG_WR(MLKEM_RDB_ADR, i); buf[i] = (uint8_t)(REG_RD(MLKEM_RDB_DAT) & 0xFF); } } int main(void) { int ok; printf("MLKEM_HELLO_START\n"); printf("MLKEM_HELLO_LEVEL=ML-KEM-768 k=3\n"); printf("MLKEM_HELLO_INPUTS=STATIC_KAT_COUNT_0\n"); memset(alice_ek, 0, sizeof(alice_ek)); memset(alice_dk, 0, sizeof(alice_dk)); memset(bob_ct, 0, sizeof(bob_ct)); memset(bob_ss, 0, sizeof(bob_ss)); memset(alice_ss, 0, sizeof(alice_ss)); memset(app_cipher, 0, sizeof(app_cipher)); memset(app_plain_out, 0, sizeof(app_plain_out)); mlkem_soft_reset(); printf("\n[1] Alice: KeyGen(d,z) -> public key ek, secret key dk\n"); write_block32(MLKEM_D, SEED_D); write_block32(MLKEM_Z, SEED_Z); run_op(OP_KEYGEN); readback(RDB_EK, alice_ek, EK_BYTES); readback(RDB_DK, alice_dk, DK_BYTES); print_hex("ALICE_PUBLIC_KEY_EK32=", alice_ek, 32); printf("\n[2] Bob: Encaps(ek,m) -> ciphertext ct, shared secret ss_bob\n"); preload(PRE_EK, alice_ek, EK_BYTES); write_block32(MLKEM_MSG, MSG_M); run_op(OP_ENCAPS); read_block32(MLKEM_SS, bob_ss); readback(RDB_CT, bob_ct, CT_BYTES); print_hex("BOB_CIPHERTEXT_CT32=", bob_ct, 32); print_hex("BOB_SHARED_SECRET=", bob_ss, 32); printf("\n[3] Bob: Encrypt application plaintext with ss_bob\n"); printf("APP_PLAINTEXT=\"%s\"\n", APP_MSG); xor_with_shared_secret(app_cipher, (const uint8_t *)APP_MSG, APP_MSG_BYTES, bob_ss); print_hex("APP_CIPHERTEXT=", app_cipher, APP_MSG_BYTES); printf("\n[4] Alice: Decaps(dk,ct) -> shared secret ss_alice\n"); preload(PRE_DK, alice_dk, DK_BYTES); preload(PRE_C, bob_ct, CT_BYTES); run_op(OP_DECAPS); read_block32(MLKEM_SS, alice_ss); print_hex("ALICE_SHARED_SECRET=", alice_ss, 32); printf("\n[5] Alice: Decrypt application ciphertext with ss_alice\n"); xor_with_shared_secret(app_plain_out, app_cipher, APP_MSG_BYTES, alice_ss); app_plain_out[APP_MSG_BYTES] = '\0'; printf("APP_DECRYPTED=\"%s\"\n", app_plain_out); ok = (memcmp(bob_ss, alice_ss, SYM_BYTES) == 0) && (memcmp(app_plain_out, APP_MSG, APP_MSG_BYTES) == 0); printf("\n[6] Verification\n"); printf("SHARED_SECRET_MATCH=%s\n", ok ? "YES" : "NO"); printf("PLAINTEXT_RECOVERED=%s\n", ok ? "YES" : "NO"); printf("MLKEM_HELLO_RESULT=%s\n", ok ? "PASS" : "FAIL"); printf("MLKEM_HELLO_DONE\n"); return ok ? 0 : 1; }