fallensigh/mlkem-sync
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(enc): Encaps E0 - op_i/msg_i/ek-load scaffold + H(ek)+G(m||H(ek))

Browse Source 
Extend mlkem_top with a runtime op_i select (0=KeyGen, 1=Encaps) and the
first Encaps stages, reusing the shared keccak_core and the ST_H multi-block
SHA3-256 machinery:
  ST_ENC_H: H(ek) over preloaded ek_bram (same FSM as KeyGen ST_H)
  ST_ENC_G: (K,r) = G(m||H(ek)) via new 64-byte single-block SHA3-512

- sha3_top_shared: add mode=2'b11 = SHA3-512 over a full 512-bit message
  (g512_pad). Standalone tb_sha3_g512 confirms it byte-exact.
- mlkem_top: new ports op_i, msg_i, ek_in_{we,addr,byte} (ek preload), ss_o,
  dbg_ct_*, dbg_r_o/dbg_hek_o. st widened 4->5 bits; ST_ENC_* states added.
  Renamed message port to msg_i to avoid collision with ST_M counter m_i.
- TB tb_mlkem_enc_katK + gen_encaps_vectors.py (per-byte ek/m/ct/ss vectors).

Verified ss==KAT.ss for K=2/3/4, cases 0-2 (all PASS). KeyGen unaffected
(K=2 c0 still ek==pk, dk==sk byte-exact).
This commit is contained in:
FallenSigh
2026-06-29 01:00:47 +08:00
parent c4669480d1
commit 0a8b3dae69
43 changed files with 21956 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

155
sync_rtl/top/mlkem_top.v
View File

@@ -28,12 +28,25 @@ module mlkem_top #(
input clk,
input rst_n,
input [2:0] k_i, // RUNTIME ML-KEM param: 2=512, 3=768, 4=1024
input op_i, // 0 = KeyGen, 1 = Encaps (captured at start_i)
input [255:0] d_i, // KeyGen seed d (byte 0 in d_i[7:0])
input [255:0] z_i, // implicit-rejection seed z
input [255:0] msg_i, // Encaps message m (byte 0 in msg_i[7:0])
input start_i, // pulse to begin KeyGen
output busy_o, // high while running
output done_o, // pulse when ek/dk ready
// Encaps ek input: stream ek bytes (=pk) into ek_bram before start_i with
// op_i=1. ek_in_we writes byte ek_in_byte at address ek_in_addr.
input ek_in_we, // write one ek byte (only honored in ST_IDLE for Encaps preload)
input [10:0] ek_in_addr, // ek byte address 0..ek_bytes-1
input [7:0] ek_in_byte, // ek byte value
// Encaps shared secret output (= K), valid at done_o
output [255:0] ss_o,
// Encaps ciphertext readback tap: ct byte by index 0..ct_bytes-1
input [10:0] dbg_ct_idx_i,
output [7:0] dbg_ct_o,
// Debug readback tap: read one stored coefficient by (poly slot, index).
// Lets stage TBs verify intermediates without wide buses.
input [3:0] dbg_slot_i, // poly slot (see localparams below)
@@ -49,7 +62,10 @@ module mlkem_top #(
output [7:0] dbg_dk_o,
// Debug taps for hash outputs
output [255:0] dbg_rho_o,
output [255:0] dbg_sigma_o
output [255:0] dbg_sigma_o,
// Encaps debug taps: r (G high half) and H(ek)
output [255:0] dbg_r_o,
output [255:0] dbg_hek_o
);
localparam Q = `Q; // 3329
@@ -70,6 +86,27 @@ module mlkem_top #(
wire [3:0] h_nblk_rt = (k_r == 3'd2) ? 4'd6 : (k_r == 3'd3) ? 4'd9 : 4'd12;
wire [11:0] h_last_rt = {6'b0, h_nblk_rt} * 12'd136 - 12'd1; // final padded byte index
// ---- Encaps runtime params ----
reg op_r; // 0=KeyGen 1=Encaps (captured at start)
reg [255:0] m_r; // Encaps message m (captured at start)
reg [255:0] ss_r; // Encaps shared secret K (= G output low half)
reg [255:0] r_r; // Encaps PRF seed r (= G output high half)
// FIPS 203: eta2 = 2 for all parameter sets.
wire [1:0] eta2_rt = 2'd2;
// Compression params: (du,dv) = (10,4) for k=2/3, (11,5) for k=4.
wire [4:0] du_rt = (k_r == 3'd4) ? 5'd11 : 5'd10;
wire [4:0] dv_rt = (k_r == 3'd4) ? 5'd5 : 5'd4;
// Ciphertext byte sizes: c1 = 32*du*k, c2 = 32*dv, ct = c1+c2.
wire [11:0] c1_bytes_rt = 12'd32 * {7'b0, du_rt} * {9'b0, k_r}; // 640/960/1408
wire [11:0] c2_bytes_rt = 12'd32 * {7'b0, dv_rt}; // 128/128/160
wire [11:0] ct_bytes_rt = c1_bytes_rt + c2_bytes_rt; // 768/1088/1568
assign ss_o = ss_r;
// ct readback tap (ct_bram added in E5/E7); tied off until then.
/* verilator lint_off UNUSEDSIGNAL */
wire [10:0] dbg_ct_idx_unused = dbg_ct_idx_i;
/* verilator lint_on UNUSEDSIGNAL */
assign dbg_ct_o = 8'd0;
// ================================================================
// Polynomial storage, sized for KMAX (worst case). Runtime k uses a
// sub-range. Slot layout (each slot = 256 coeffs):
@@ -203,10 +240,17 @@ module mlkem_top #(
reg [10:0] ek_wa, dkp_wa;
reg [7:0] ek_wd, dkp_wd;
// ek BRAM write port: KeyGen ST_E drives ek_we/ek_wa/ek_wd; Encaps preloads
// ek from the external ek_in_* port (TB streams ek=pk before start_i). The
// two never overlap (preload happens in ST_IDLE before an Encaps run).
wire ek_we_mux = ek_in_we ? 1'b1 : ek_we;
wire [10:0] ek_wa_mux = ek_in_we ? ek_in_addr : ek_wa;
wire [7:0] ek_wd_mux = ek_in_we ? ek_in_byte : ek_wd;
sd_bram #(.W(8), .D(2048), .A(11)) u_ek_bram (
.clk(clk),
.rd_addr(ek_rd_addr), .rd_data(ek_rd_data),
.wr_en(ek_we), .wr_addr(ek_wa), .wr_data(ek_wd)
.wr_en(ek_we_mux), .wr_addr(ek_wa_mux), .wr_data(ek_wd_mux)
);
sd_bram #(.W(8), .D(2048), .A(11)) u_dkp_bram (
.clk(clk),
@@ -236,17 +280,29 @@ module mlkem_top #(
// ================================================================
// Top-level FSM (built incrementally). Stage 2a: G only.
// ================================================================
localparam ST_IDLE = 4'd0;
localparam ST_G = 4'd1; // run G(d||K), capture rho/sigma
localparam ST_A = 4'd2; // generate A_hat[i][j] via SampleNTT
localparam ST_C = 4'd3; // generate s[i],e[i] via CBD
localparam ST_N = 4'd4; // forward NTT of s[i],e[i] in place
localparam ST_M = 4'd5; // matrix accumulate t_hat = e_hat + sum A o s_hat
localparam ST_E = 4'd6; // byteEncode12 -> ek/dk BRAM
localparam ST_H = 4'd7; // H(ek) via multi-block SHA3-256
localparam ST_DONE = 4'd15;
localparam ST_IDLE = 5'd0;
localparam ST_G = 5'd1; // run G(d||K), capture rho/sigma
localparam ST_A = 5'd2; // generate A_hat[i][j] via SampleNTT
localparam ST_C = 5'd3; // generate s[i],e[i] via CBD
localparam ST_N = 5'd4; // forward NTT of s[i],e[i] in place
localparam ST_M = 5'd5; // matrix accumulate t_hat = e_hat + sum A o s_hat
localparam ST_E = 5'd6; // byteEncode12 -> ek/dk BRAM
localparam ST_H = 5'd7; // H(ek) via multi-block SHA3-256
// ---- Encaps states ----
localparam ST_ENC_LOAD = 5'd8; // ek already preloaded; placeholder/settle
localparam ST_ENC_H = 5'd9; // H(ek) via multi-block SHA3-256 (reuses ST_H logic)
localparam ST_ENC_G = 5'd10; // (K,r) = G(m || H(ek)), 64-byte single block
localparam ST_ENC_A = 5'd11; // regenerate A_hat (transpose used in U)
localparam ST_ENC_TDEC = 5'd12; // byteDecode12: ek -> t_hat (bank_t)
localparam ST_ENC_C = 5'd13; // sample y,e1,e2 via CBD (eta1/eta2)
localparam ST_ENC_N = 5'd14; // forward NTT of y in place
localparam ST_ENC_U = 5'd15; // u = INTT(sum A^T o y_hat) + e1
localparam ST_ENC_C1 = 5'd16; // Compress_du + byteEncode_du -> ct c1
localparam ST_ENC_V = 5'd17; // v = INTT(sum t_hat o y_hat) + e2 + mu
localparam ST_ENC_C2 = 5'd18; // Compress_dv + byteEncode_dv -> ct c2
localparam ST_DONE = 5'd31;
reg [3:0] st, st_next;
reg [4:0] st, st_next;
reg [255:0] rho_r, sigma_r;
// A-generation bookkeeping: explicit i/j counters (avoid runtime divide)
@@ -270,6 +326,8 @@ module mlkem_top #(
assign done_o = (st == ST_DONE);
assign dbg_rho_o = rho_r;
assign dbg_sigma_o = sigma_r;
assign dbg_r_o = r_r;
assign dbg_hek_o = hek_r;
// ---- sha3_top in G mode: data_i = {K_byte, d} (d byte0 in [7:0]) ----
reg sha3_valid;
@@ -277,7 +335,9 @@ module mlkem_top #(
wire [511:0] sha3_hash;
wire sha3_vo;
reg sha3_ack; // consumer ready for hash
wire [511:0] g_data = {248'b0, 5'b0, k_r, d_i}; // data_i[263:256]=k, [255:0]=d
wire [511:0] kg_g_data = {248'b0, 5'b0, k_r, d_i}; // KeyGen G: [263:256]=k, [255:0]=d
wire [511:0] enc_g_data = {hek_r, m_r}; // Encaps G: m || H(ek), 64 bytes
wire [511:0] g_data = (st == ST_ENC_G) ? enc_g_data : kg_g_data;
// ================================================================
// Shared keccak_core + phase mux (3 consumers -> 1 core)
@@ -301,10 +361,12 @@ module mlkem_top #(
wire sha3_kc_ready_i, snt_kc_ready_i, cbd_kc_ready_i; // all 1'b1
/* verilator lint_on UNUSEDSIGNAL */
// phase selects (mutually exclusive)
wire sel_sha3 = (st == ST_G) || (st == ST_H);
wire sel_snt = (st == ST_A);
wire sel_cbd = (st == ST_C);
// phase selects (mutually exclusive). Encaps adds ST_ENC_H/G (sha3),
// ST_ENC_A (snt), ST_ENC_C (cbd).
wire sel_sha3 = (st == ST_G) || (st == ST_H) ||
(st == ST_ENC_H) || (st == ST_ENC_G);
wire sel_snt = (st == ST_A) || (st == ST_ENC_A);
wire sel_cbd = (st == ST_C) || (st == ST_ENC_C);
// gated output-valid: only the active consumer sees kc_valid_o
wire kc_valid_o_sha3 = kc_valid_o & sel_sha3;
@@ -334,10 +396,13 @@ module mlkem_top #(
// multi-block absorb path (mb_en=1). These phases are disjoint in the
// top FSM, so one sha3_top (one keccak_core) is sufficient. mb_en and
// ready_i are muxed by phase; data_i/mode only matter while mb_en=0.
wire sha3_mb_en = (st == ST_H);
wire sha3_mb_en = (st == ST_H) || (st == ST_ENC_H);
// sha3 single-block mode: 2'b00 = G(33B) for KeyGen, 2'b11 = G(64B) for
// Encaps (m||H(ek)). mode only matters when mb_en=0.
wire [1:0] sha3_mode = (st == ST_ENC_G) ? 2'b11 : 2'b00;
sha3_top_shared u_sha3 (
.clk(clk), .rst_n(rst_n),
.mode(2'b00), // G = SHA3-512 (only used when mb_en=0)
.mode(sha3_mode), // G = SHA3-512 (only used when mb_en=0)
.data_i(g_data),
.valid_i(sha3_valid),
.ready_o(sha3_ready),
@@ -408,7 +473,7 @@ module mlkem_top #(
// dbg_dk ek-region takes priority (its loop runs after the ek loop), so the
// single read port serves dbg_byte_o (ek loop) then dbg_dk_o (dk loop).
wire [11:0] dbgdk_ek_off = dbg_dk_idx_i - dk_bytes_rt; // offset within ek region
assign ek_rd_addr = (st == ST_H) ? h_g_addr[10:0] :
assign ek_rd_addr = (st == ST_H || st == ST_ENC_H) ? h_g_addr[10:0] :
dbgdk_in_ek ? dbgdk_ek_off[10:0]
: dbg_byte_idx_i;
// dkp BRAM read-address mux: dbg_byte (sel=1) or dbg_dk (dkp region).
@@ -600,7 +665,7 @@ module mlkem_top #(
always @(*) begin
st_next = st;
case (st)
ST_IDLE: if (start_i) st_next = ST_G;
ST_IDLE: if (start_i) st_next = op_i ? ST_ENC_H : ST_G;
ST_G: if (sha3_vo) st_next = ST_A;
ST_A: if (a_pair >= kk_rt) st_next = ST_C;
ST_C: if (c_poly >= {1'b0, k_r, 1'b0}) st_next = ST_N;
@@ -608,6 +673,9 @@ module mlkem_top #(
ST_M: if (m_i >= k_r) st_next = ST_E;
ST_E: if (e_done) st_next = ST_H;
ST_H: if (h_phase == 2'd3) st_next = ST_DONE;
// ---- Encaps ----
ST_ENC_H: if (h_phase == 2'd3) st_next = ST_ENC_G;
ST_ENC_G: if (sha3_vo) st_next = ST_DONE; // E0: stop here (K,r captured)
ST_DONE: st_next = ST_IDLE;
default: st_next = ST_IDLE;
endcase
@@ -617,6 +685,10 @@ module mlkem_top #(
if (!rst_n) begin
st <= ST_IDLE;
k_r <= 3'd0;
op_r <= 1'b0;
m_r <= 256'd0;
ss_r <= 256'd0;
r_r <= 256'd0;
rho_r <= 256'd0;
sigma_r <= 256'd0;
sha3_valid <= 1'b0;
@@ -680,11 +752,25 @@ module mlkem_top #(
ek_we <= 1'b0;
dkp_we <= 1'b0;
// Kick off G when entering ST_G
// Kick off when entering from IDLE: KeyGen starts G; Encaps captures
// op/m and arms the H(ek) machinery (ST_ENC_H reuses the ST_H FSM).
if (st == ST_IDLE && start_i) begin
k_r <= k_i; // capture runtime ML-KEM param
sha3_valid <= 1'b1;
sha3_ack <= 1'b1;
k_r <= k_i; // capture runtime ML-KEM param
op_r <= op_i;
if (op_i) begin
m_r <= msg_i; // capture Encaps message
// arm H(ek) (same fields the ST_E->ST_H arming sets)
h_blk <= 3'd0;
h_byte <= 8'd0;
h_phase <= 2'd0; // assemble
h_mbvalid<= 1'b0;
h_mblast <= 1'b0;
h_ack <= 1'b1; // ready to consume final digest
h_wb_vld <= 1'b0;
end else begin
sha3_valid <= 1'b1;
sha3_ack <= 1'b1;
end
end
// Drop valid once accepted
if (sha3_valid && sha3_ready) sha3_valid <= 1'b0;
@@ -954,8 +1040,10 @@ module mlkem_top #(
h_wb_vld <= 1'b0; // no pending writeback yet
end
// ---- ST_H: H(ek) via multi-block SHA3-256 (6/9/12 pre-padded blocks) ----
if (st == ST_H) begin
// ---- ST_H / ST_ENC_H: H(ek) via multi-block SHA3-256 ----
// Same logic for KeyGen (ek just written by ST_E) and Encaps (ek
// preloaded into ek_bram by the TB). hek_r captures H(ek) either way.
if (st == ST_H || st == ST_ENC_H) begin
case (h_phase)
// assemble 136 bytes of block h_blk into h_block_r.
// ek BRAM read is registered: present addr for h_byte this
@@ -1005,6 +1093,19 @@ module mlkem_top #(
default: ; // 2'd3 done: hold
endcase
end
// Arm Encaps G when H(ek) finishes: fire the 64-byte single-block
// G(m||H(ek)). hek_r now holds H(ek); enc_g_data = {hek_r, m_r}.
if (st == ST_ENC_H && st_next == ST_ENC_G) begin
sha3_valid <= 1'b1;
sha3_ack <= 1'b1;
end
// Capture (K, r) when Encaps G completes. ss = K = low half, r = high.
if (st == ST_ENC_G && sha3_vo) begin
ss_r <= sha3_hash[255:0]; // K (shared secret) = G bytes 0..31
r_r <= sha3_hash[511:256]; // r = G bytes 32..63
sha3_ack <= 1'b0;
end
end
end