feat(dec): Decaps D7 - implicit-reject compare + end-to-end KAT

Final FO step (FIPS 203 Alg 17 steps 9-11): compare c' to c and select the
shared secret with implicit rejection. Completes full ML-KEM Decaps.

- ST_DEC_CMP walks all ct bytes (no early-out: constant work, matching the
  constant-time spec intent), reading c' (ct_bram) and c (c_in_bram) in
  parallel, OR-ing every byte difference into cmp_neq; the final byte latches
  dec_reject.
- ss_o = dec_reject ? kbar_r (K-bar) : ss_r (K'). KeyGen/Encaps leave
  dec_reject=0 so ss_r passes through unchanged.
- ST_ENC_C2 terminal branches on op_r: Decaps -> ST_DEC_CMP, Encaps -> DONE.

The dec TB now runs end-to-end twice per case (dk loaded once):
  - valid ct  : ss_o == KAT ss   (c'==c -> K'),            dec_reject==0
  - corrupt ct: ss_o == KAT ss_n (c'!=c -> K-bar=J(z||ct_n)), dec_reject==1
exercising both the accept and implicit-reject paths against the KAT's own
ct_n/ss_n reject vectors.

Verified: dec D7 K=2/3/4 all cases PASS (accept + reject); KeyGen + Encaps
unregressed. Full ML-KEM (KeyGen + Encaps + Decaps w/ implicit rejection) now
works in hardware across all three parameter sets.
This commit is contained in:
2026-06-29 21:56:07 +08:00
parent a734eb2cad
commit 2b70431923
3 changed files with 144 additions and 38 deletions

View File

@@ -51,9 +51,11 @@ module tb_mlkem_dec_katK_xsim;
reg [7:0] dk_b [0:DKB-1];
reg [7:0] ct_b [0:CTB-1];
reg [7:0] ctn_b [0:CTB-1];
reg [7:0] ss_b [0:31];
reg [7:0] ssn_b [0:31];
integer c, i, j, errors, casenum;
reg [8*80-1:0] tag, dkfile, ctfile, ssfile;
reg [8*80-1:0] tag, dkfile, ctfile, ssfile, ctnfile, ssnfile;
initial begin
if (!$value$plusargs("CASE=%d", casenum)) casenum = 0;
@@ -61,46 +63,44 @@ module tb_mlkem_dec_katK_xsim;
$sformat(dkfile, "sync_rtl/top/TB/vectors/dec_%0s_c%0d_dk.hex", tag, casenum);
$sformat(ctfile, "sync_rtl/top/TB/vectors/dec_%0s_c%0d_ct.hex", tag, casenum);
$sformat(ssfile, "sync_rtl/top/TB/vectors/dec_%0s_c%0d_ss.hex", tag, casenum);
$sformat(ctnfile, "sync_rtl/top/TB/vectors/dec_%0s_c%0d_ctn.hex", tag, casenum);
$sformat(ssnfile, "sync_rtl/top/TB/vectors/dec_%0s_c%0d_ssn.hex", tag, casenum);
$readmemh(dkfile, dk_b);
$readmemh(ctfile, ct_b);
$readmemh(ssfile, ss_b);
$readmemh(ctnfile, ctn_b);
$readmemh(ssnfile, ssn_b);
k_i = KP[2:0];
$display("=== ML-KEM K=%0d Decaps KAT case %0d (D0: load+parse) ===", KP, casenum);
rst_n=0; repeat(4) @(posedge clk); rst_n=1; @(posedge clk);
// ---- stream dk into the design (1 byte/cycle) ----
// ---- stream dk into the design (1 byte/cycle); persists for both runs ----
for (i = 0; i < DKB; i = i + 1) begin
dk_in_we = 1'b1; dk_in_addr = i[11:0]; dk_in_byte = dk_b[i];
@(posedge clk);
end
dk_in_we = 1'b0;
// ---- stream ct into c_in_bram (1 byte/cycle) ----
for (i = 0; i < CTB; i = i + 1) begin
c_in_we = 1'b1; c_in_addr = i[10:0]; c_in_byte = ct_b[i];
@(posedge clk);
end
c_in_we = 1'b0; @(posedge clk);
errors = 0;
// ---- run Decaps ----
start_i=1; @(posedge clk); start_i=0;
c=0; while(!done_o && c<2000000) begin @(posedge clk); c=c+1; end
if(!done_o) begin $display("FAIL K=%0d case %0d: timeout", KP, casenum); $finish; end
$display("=== Decaps D6 done in %0d cyc ===", c);
// ================= Run 1: valid ciphertext (accept path) =================
// ct == c, so c' == c, no implicit reject: ss == K' == KAT ss.
run_decaps(0); // 0 = use ct_b
verify_d0; // dk parse (errors NOT reset here; see verify_d0 note)
verify_d4; // m' (in m_r, survives the re-encrypt)
verify_d5; // K'/r'/K-bar
verify_d6; // c' == golden
verify_d7(0); // ss_o == KAT ss, dec_reject == 0
// D6 re-encrypt clobbers bank_a/bank_se/bank_t, so the bank-based stage
// checks (D1 v', D2 s_hat/u_hat, D3 w) are no longer valid at end-of-run;
// their correctness was proven on earlier per-stage builds and transitively
// by c'. Here we check the surviving register/BRAM artifacts: dk parse (D0),
// m'/K'/r'/K-bar (D5), and the re-encrypted ciphertext c' (D6).
verify_d0; // also initializes errors = 0
verify_d4; // m' (now in m_r, survives the re-encrypt)
verify_d5;
verify_d6;
if (errors == 0) $display("K=%0d CASE %0d PASS (D6): c' = Encrypt(ek,m',r') OK", KP, casenum);
else $display("K=%0d CASE %0d FAIL (D6): %0d total errors", KP, casenum, errors);
// ================= Run 2: corrupted ciphertext (reject path) =============
// ct_n != c, so c' != c -> implicit reject: ss == K-bar == J(z||ct_n) == ss_n.
run_decaps(1); // 1 = use ctn_b
verify_d7(1); // ss_o == KAT ss_n, dec_reject == 1
if (errors == 0) $display("K=%0d CASE %0d PASS (D7): end-to-end Decaps (accept + reject) OK", KP, casenum);
else $display("K=%0d CASE %0d FAIL (D7): %0d total errors", KP, casenum, errors);
$finish;
end
@@ -117,6 +117,56 @@ module tb_mlkem_dec_katK_xsim;
end
endtask
// Stream a ciphertext into c_in_bram (sel 0 = valid ct_b, 1 = corrupted ctn_b),
// pulse start_i, and wait for done_o. dk is already loaded and persists.
task run_decaps;
input sel;
integer ii;
begin
for (ii = 0; ii < CTB; ii = ii + 1) begin
c_in_we = 1'b1; c_in_addr = ii[10:0];
c_in_byte = sel ? ctn_b[ii] : ct_b[ii];
@(posedge clk);
end
c_in_we = 1'b0; @(posedge clk);
start_i=1; @(posedge clk); start_i=0;
c=0; while(!done_o && c<2000000) begin @(posedge clk); c=c+1; end
if(!done_o) begin $display("FAIL K=%0d case %0d: timeout (sel=%0d)", KP, casenum, sel); $finish; end
$display("=== Decaps run (sel=%0d) done in %0d cyc ===", sel, c);
// settle so dbg taps reflect the finished run
repeat(2) @(posedge clk);
end
endtask
// D7: end-to-end shared secret. sel=0 (accept): ss_o == KAT ss, dec_reject=0.
// sel=1 (reject): ss_o == KAT ss_n (= K-bar = J(z||ct_n)), dec_reject=1.
task verify_d7;
input sel;
integer be;
reg [255:0] exp;
begin
be = 0;
for (j = 0; j < 32; j = j + 1)
exp[8*j +: 8] = sel ? ssn_b[j] : ss_b[j];
for (j = 0; j < 32; j = j + 1)
if (ss_o[8*j +: 8] !== exp[8*j +: 8]) begin
if (be < 4) $display(" ss[%0d] got=%02x exp=%02x", j, ss_o[8*j +: 8], exp[8*j +: 8]);
be = be + 1;
end
if (be == 0)
$display(" PASS: ss == KAT %0s (32 bytes), reject=%b",
sel ? "ss_n (reject)" : "ss (accept)", dut.dec_reject);
else
$display(" FAIL: ss %0d byte mismatches (sel=%0d, reject=%b)", be, sel, dut.dec_reject);
errors = errors + be;
// also check the reject flag matches the expectation
if (dut.dec_reject !== sel) begin
$display(" FAIL: dec_reject=%b expected %b", dut.dec_reject, sel);
errors = errors + 1;
end
end
endtask
// D0: verify dk parse. H(ek)=dk[768K+32:+32], z=dk[768K+64:+32] captured into
// hek_r/z_r (dbg_dech_o/dbg_decz_o). ek_pke=dk[384K:768K+32] in ek_bram
// (dbg_byte sel=0), dk_pke=dk[0:384K] in dkp_bram (sel=1).
@@ -124,12 +174,14 @@ module tb_mlkem_dec_katK_xsim;
integer be;
reg [7:0] got;
begin
errors = 0;
// errors is initialized by the caller (main flow), not here.
// H(ek)
be = 0;
for (j = 0; j < 32; j = j + 1)
if (dbg_dech_o[8*j +: 8] !== dk_b[DKPB + EKB + j]) errors = errors + 1;
if (errors == 0) $display(" PASS: H(ek) parsed == dk[768K+32 ..]");
else $display(" FAIL: H(ek) %0d byte mismatches", errors);
if (dbg_dech_o[8*j +: 8] !== dk_b[DKPB + EKB + j]) be = be + 1;
if (be == 0) $display(" PASS: H(ek) parsed == dk[768K+32 ..]");
else $display(" FAIL: H(ek) %0d byte mismatches", be);
errors = errors + be;
// z
be = 0;

View File

@@ -134,7 +134,10 @@ module mlkem_top #(
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
wire [11:0] ct_bytes_rt = c1_bytes_rt + c2_bytes_rt; // 768/1088/1568
assign ss_o = ss_r;
// Shared secret out: Encaps drives ss_r = K directly. Decaps implicit reject:
// ss = (c' == c) ? K' : K-bar. K' is in ss_r, K-bar in kbar_r; dec_reject is
// the latched c' != c flag from D7 (0 for KeyGen/Encaps, so ss_r passes through).
assign ss_o = dec_reject ? kbar_r : ss_r;
// ---- E1: rho-load + byteDecode12 bookkeeping ----
// rho load: read 32 bytes from ek_bram at offset 384*k into rho_r.
@@ -180,7 +183,8 @@ module mlkem_top #(
.rd_addr(ct_rd_addr), .rd_data(ct_rd_data),
.wr_en(ct_we), .wr_addr(ct_wa), .wr_data(ct_wd)
);
assign ct_rd_addr = dbg_ct_idx_i;
// ct_bram read addr: D7 compare walks c' via cmp_idx; else debug readback.
assign ct_rd_addr = (st == ST_DEC_CMP) ? cmp_idx : dbg_ct_idx_i;
assign dbg_ct_o = ct_rd_data;
// ---- c_in_bram: Decaps input ciphertext c (<=1568 B). Preloaded via
@@ -199,7 +203,8 @@ module mlkem_top #(
// stage) so the registered BRAM read yields the byte 1 cycle later, matching
// the assemble/writeback pipeline. Otherwise the registered cin_rd_addr_r
// (D1 walker / idle) drives it.
assign cin_rd_addr = (st == ST_DEC_J) ? dj_c_idx[10:0] : cin_rd_addr_r;
assign cin_rd_addr = (st == ST_DEC_J) ? dj_c_idx[10:0] :
(st == ST_DEC_CMP) ? cmp_idx : cin_rd_addr_r;
// ================================================================
// Polynomial storage, sized for KMAX (worst case). Runtime k uses a
@@ -488,6 +493,7 @@ module mlkem_top #(
localparam ST_DEC_MENC = 5'd25; // D4: m' = byteEncode_1(Compress_1(w))
localparam ST_DEC_G = 5'd26; // D5: (K',r') = G(m' || h), SHA3-512 single block
localparam ST_DEC_J = 5'd27; // D5: K-bar = J(z || c), multi-block SHAKE256
localparam ST_DEC_CMP = 5'd28; // D7: c' == c compare + implicit-reject ss select
localparam ST_DONE = 5'd31;
reg [4:0] st, st_next;
@@ -970,6 +976,20 @@ module mlkem_top #(
reg men_done;
// D4 packs the recovered message bits directly into m_r (reused by D6's
// V/mu re-encrypt and exposed on dbg_mprime_o).
// ================================================================
// D7: implicit-reject compare (Decaps ST_DEC_CMP).
// Walk all ct bytes, reading c' (ct_bram) and c (c_in_bram) in parallel;
// OR every byte difference into cmp_neq. After the last byte, ss = cmp_neq
// ? K-bar : K'. The walk visits every byte regardless of early differences
// (constant work / no early-out), matching the constant-time spec intent.
// ph0: present both addrs; ph1: compare registered reads, accumulate.
reg [10:0] cmp_idx; // byte index 0..ct_bytes
reg cmp_ph; // 0=present addr, 1=compare
reg cmp_neq; // sticky: any byte differed
reg cmp_done;
reg dec_reject; // latched cmp_neq at end (drives ss select)
wire cmp_byte_neq = (ct_rd_data != cin_rd_data);
wire [13:0] men_rd = UPSUM*256 + men_idx; // bank_t w read addr
wire men_w_bit = (bt_rd_data > 12'd832) && (bt_rd_data <= 12'd2496); // Compress_1
@@ -1151,7 +1171,10 @@ module mlkem_top #(
ST_ENC_TDEC: if (td_done) st_next = ST_ENC_E2MV; // t_hat decoded -> relocate e2
ST_ENC_E2MV: if (em_done) st_next = ST_ENC_V; // e2 relocated -> compute v
ST_ENC_V: if (u_row >= 3'd1) st_next = ST_ENC_C2; // E6: v done -> E7
ST_ENC_C2: if (cp_done) st_next = ST_DONE; // E7: c2 packed -> done
ST_ENC_C2: if (cp_done) st_next = (op_r == 2'd2) ? ST_DEC_CMP : ST_DONE;
// D7: walk ct bytes comparing c' (ct_bram) vs c (c_in_bram); when the
// last byte is compared, select ss = (c'==c) ? K' : K-bar and finish.
ST_DEC_CMP: if (cmp_done) st_next = ST_DONE;
ST_DONE: st_next = ST_IDLE;
default: st_next = ST_IDLE;
endcase
@@ -1275,6 +1298,11 @@ module mlkem_top #(
men_idx <= 8'd0;
men_ph <= 2'd0;
men_done <= 1'b0;
cmp_idx <= 11'd0;
cmp_ph <= 1'b0;
cmp_neq <= 1'b0;
cmp_done <= 1'b0;
dec_reject <= 1'b0;
dj_blk <= 4'd0;
dj_byte <= 8'd0;
dj_phase <= 2'd0;
@@ -2036,8 +2064,34 @@ module mlkem_top #(
endcase
end
// Arm rho-load when entering ST_ENC_LOAD. rho = ek[384k .. 384k+31].
// Entered from Encaps G (ST_ENC_G) or Decaps D6 re-encrypt (ST_DEC_J).
// Arm D7 compare when re-encrypt finishes (ST_ENC_C2 -> ST_DEC_CMP).
if (st == ST_ENC_C2 && st_next == ST_DEC_CMP) begin
cmp_idx <= 11'd0;
cmp_ph <= 1'b0;
cmp_neq <= 1'b0;
cmp_done <= 1'b0;
end
// ---- ST_DEC_CMP (D7): constant-time c' == c compare + ss select ----
// ph0: present cmp_idx to both ct_bram (c') and c_in_bram (c).
// ph1: registered reads valid -> OR any difference into cmp_neq,
// advance. After the last byte, latch dec_reject and ss stays
// K' (ss_r) or switches to K-bar (kbar_r) via the ss_o mux.
if (st == ST_DEC_CMP && !cmp_done) begin
if (!cmp_ph) begin
cmp_ph <= 1'b1; // addr presented; wait for read
end else begin
cmp_neq <= cmp_neq | cmp_byte_neq;
if (cmp_idx == ct_bytes_rt - 12'd1) begin
dec_reject <= cmp_neq | cmp_byte_neq; // include final byte
cmp_done <= 1'b1;
end else begin
cmp_idx <= cmp_idx + 11'd1;
end
cmp_ph <= 1'b0;
end
end
if (st_next == ST_ENC_LOAD &&
(st == ST_ENC_G || st == ST_DEC_J)) begin
rl_idx <= 6'd0;