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

fix(sample_ntt,sha3): FIPS-203 SHAKE-128 squeeze + self-checking sha3 TBs

Browse Source 
sample_ntt was non-conformant: both RTL and the test reference re-ran
keccak_p after every 3-byte squeeze instead of consuming the full
1344-bit SHAKE-128 rate. Only coeff[0] matched a standard sampler, so
the generated A matrix would not interoperate with any compliant ML-KEM.

- sample_ntt_sync{,_shared}.v: walk all 56 groups of the rate block via
  grp_ptr_r; re-permute only when the block is exhausted. Verified
  256/256 against ml-kem-r Rust sample_ntt on two seeds, and 1536/1536
  in the Verilator framework (runtime ~128x faster per poly).
- gen_vectors.py: use a self-contained hashlib.shake_128 oracle.

sha3 testbench fixes (all now self-check hash_o against verified vectors,
cross-checked with hashlib and ml-kem-r mlkem_G):
- tb_sha3_xsim_simple.v: test G/H/J modes, not just G.
- tb_keccak_core_xsim.v: correct the wrong EXPECTED_STATE constant
  (RTL was correct; lane0 = 0xf1258f7940e1dde7 per FIPS 202).
- tb_sha3_xsim.v: read expected file and self-check per vector; add
  vectors/g_basic_{input,expected}.hex (3 G / 2 H / 2 J).

Remove stale sha3_chain test (its RTL was deleted in 1cace51) and its
README references. Extend .gitignore for XSIM artifacts and result dumps.
This commit is contained in:
FallenSigh
2026-06-27 17:23:28 +08:00
parent 5d86000231
commit 4d7ce69405
12 changed files with 318 additions and 295 deletions
Download Patch File Download Diff File Expand all files Collapse all files

77
sync_rtl/sample_ntt/sample_ntt_sync.v
View File

@@ -3,12 +3,14 @@
// Generates one k×k polynomial (256 coefficients) via SHAKE-128 XOF
// rejection sampling from seed rho || j || i.
//
// Matches Python reference (sample.py/SHA_3.py) bit-exactly:
// FIPS 202/203 conformant SHAKE-128 squeeze:
// - Absorb: S = Keccak-p(padded(rho || j || i))
// - For each squeeze: take S[23:0] (3 bytes), extract d1[11:0], d2[23:12]
// - Squeeze the full 1344-bit (168-byte) rate as 56 groups of 3 bytes,
// each group g read from S[24*g +: 24]; extract d1[11:0], d2[23:12]
// - Accept d if d < Q=3329
// - S = Keccak-p(S) (permute between every 3-byte squeeze)
// - Repeat until 256 coefficients collected
// - Only after all 56 groups of the block are consumed: S = Keccak-p(S)
// (re-permute once per rate block, NOT per 3-byte group)
// - Repeat until 256 coefficients collected (~3 blocks)
//
// Parameters:
// K = 4 (ML-KEM parameter)
@@ -52,6 +54,10 @@ module sample_ntt_sync #(parameter K = 4) (
// ================================================================
localparam Q = `Q; // 3329
// SHAKE-128 rate = 1344 bits = 168 bytes = 56 groups of 3 bytes.
// After consuming all 56 groups of a block, re-permute the state.
localparam GRP_MAX = 6'd55;
// ================================================================
// FSM state encoding
// ================================================================
@@ -84,6 +90,13 @@ module sample_ntt_sync #(parameter K = 4) (
// ================================================================
reg [1599:0] squeeze_state_r;
// ================================================================
// Group pointer: which 3-byte group within the 1344-bit rate
// block is currently being consumed (0..GRP_MAX). Re-permute when
// it would exceed GRP_MAX.
// ================================================================
reg [5:0] grp_ptr_r;
// ================================================================
// Registered d1, d2 and acceptance flags
// ================================================================
@@ -121,15 +134,21 @@ module sample_ntt_sync #(parameter K = 4) (
};
// ================================================================
// Comb: extract d1,d2 from squeeze state
// Comb: extract d1,d2 from the current 3-byte group of squeeze state
// ================================================================
// squeeze_state_r[7:0]=c0, [15:8]=c1, [23:16]=c2
// group g occupies bits [24*g +: 24]: c0=byte0, c1=byte1, c2=byte2
// d1 = {c1[3:0], c0}
// d2 = {c2, c1[7:4]}
wire [10:0] grp_bit_off;
assign grp_bit_off = grp_ptr_r * 11'd24; // 0..1320, +24 ≤ 1344 (rate)
wire [23:0] grp_bits;
assign grp_bits = squeeze_state_r[ grp_bit_off +: 24 ];
wire [7:0] c0, c1, c2;
assign c0 = squeeze_state_r[7:0];
assign c1 = squeeze_state_r[15:8];
assign c2 = squeeze_state_r[23:16];
assign c0 = grp_bits[7:0];
assign c1 = grp_bits[15:8];
assign c2 = grp_bits[23:16];
wire [11:0] d1_comb, d2_comb;
assign d1_comb = {c1[3:0], c0};
@@ -161,10 +180,10 @@ module sample_ntt_sync #(parameter K = 4) (
.ready_i (1'b1)
);
// kc_valid_i: asserted during ABSORB and first phase of SQUEEZE.
// Keccak captures it on the transition (when ready_o=1).
// kc_valid_i: asserted on the ABSORB load, and for one cycle when the
// squeeze block is exhausted (SQUEEZE → WAIT) to re-permute the state.
assign kc_valid_i = (state_next == ST_ABSORB) ||
(state_r == ST_SQUEEZE && sq_phase_r == 2'd0);
(state_r == ST_SQUEEZE && state_next == ST_WAIT);
// kc_state_i: absorb_state in ABSORB, squeeze_state_r otherwise
assign kc_state_i = (state_next == ST_ABSORB) ? absorb_state : squeeze_state_r;
@@ -186,6 +205,11 @@ module sample_ntt_sync #(parameter K = 4) (
assign ready_o = (state_r == ST_IDLE);
wire need_more = (coeff_cnt_r < 9'd256);
// grp_done: current 3-byte group fully consumed (phase 2 resolved):
// d2 rejected, no longer need coefficients, or d2 was just accepted out.
wire grp_done = (state_r == ST_SQUEEZE) && (sq_phase_r == 2'd2) &&
(!d2_acc_r || !need_more || (valid_o_r && ready_i));
// ================================================================
// FSM: state_next (combinational)
// ================================================================
@@ -205,11 +229,18 @@ module sample_ntt_sync #(parameter K = 4) (
end
ST_SQUEEZE: begin
// Sub-phase transitions managed in sequential logic.
// Only transitions to ST_WAIT from phase 2 when done.
if (sq_phase_r == 2'd2 &&
(!d2_acc_r || !need_more || (valid_o_r && ready_i)))
state_next = ST_WAIT;
// A group is fully consumed once phase 2 resolves (d2 output,
// rejected, or no longer needed). Then either advance to the
// next group in this block, re-permute (block exhausted), or
// finish.
if (grp_done) begin
if (!need_more)
state_next = ST_DONE;
else if (grp_ptr_r < GRP_MAX)
state_next = ST_SQUEEZE; // next group, no re-permute
else
state_next = ST_WAIT; // block exhausted: re-permute
end
end
ST_WAIT: begin
@@ -239,6 +270,7 @@ module sample_ntt_sync #(parameter K = 4) (
sq_phase_r <= 2'd0;
coeff_cnt_r <= 9'd0;
squeeze_state_r <= 1600'd0;
grp_ptr_r <= 6'd0;
d1_r <= 12'd0;
d2_r <= 12'd0;
d1_acc_r <= 1'b0;
@@ -265,10 +297,12 @@ module sample_ntt_sync #(parameter K = 4) (
end
// ---------------------------------------------------------
// Latch keccak output when valid_o fires
// Latch keccak output when valid_o fires. A fresh block
// starts at group 0 (ABSORB load or WAIT re-permute result).
// ---------------------------------------------------------
if (kc_valid_o) begin
squeeze_state_r <= kc_state_o;
grp_ptr_r <= 6'd0;
end
// ---------------------------------------------------------
@@ -347,6 +381,13 @@ module sample_ntt_sync #(parameter K = 4) (
valid_o_r <= 1'b0;
end
endcase
// Group consumed but block not exhausted: advance to the
// next 3-byte group within the same block (no re-permute).
if (grp_done && need_more && grp_ptr_r < GRP_MAX) begin
grp_ptr_r <= grp_ptr_r + 6'd1;
sq_phase_r <= 2'd0;
end
end else if (state_r != ST_SQUEEZE && state_next == ST_SQUEEZE) begin
// About to enter SQUEEZE: reset phase and output
sq_phase_r <= 2'd0;