nx_kyber_sample_ntt_wasm.nx source

1// nx_kyber_sample_ntt_wasm.nx -- FIPS 203 §4.2.1 SampleNTT (rejection sampling). 2// 3// Turns a SHAKE128 byte stream into a uniformly-random NTT-domain 4// polynomial: 256 coefficients in Z_3329. The matrix A used in 5// ML-KEM-768 K-PKE keygen / encrypt is built from 9 such polynomials 6// (k * k = 9, k=3 for ML-KEM-768), one per (i, j) pair seeded by the 7// (rho || j || i) byte string SHAKE128'd into bytes. 8// 9// Algorithm (FIPS 203 §4.2.1 Algorithm 4): 10// Read 3 bytes at a time from the SHAKE stream. Split into two 12-bit 11// values d1 = b0 | ((b1 & 0xf) << 8), d2 = (b1 >> 4) | (b2 << 4). 12// If d < q, accept as next coefficient. Otherwise discard and continue. 13// Stop when 256 coefficients have been accepted. 14// 15// Browser-side glue: 16// 1. SHAKE128.absorb(rho || j_byte || i_byte) 17// 2. SHAKE128.squeeze(buf, buf_len) -- ~600 bytes usually enough 18// 3. nx_kyber_sample_ntt(poly_out, buf, buf_len) -> n_consumed (>=0) 19// OR -1 if buffer too short (caller should squeeze more + retry) 20// 21// API: 22// nx_kyber_sample_ntt(poly_out, buf, buf_len) -> i64 23// returns number of bytes consumed on success, -1 if buffer 24// too short to reach 256 accepted coefficients. 25// 26// Verified: KAT against PQClean reference (a known SHAKE stream 27// produces a known polynomial). 28// 29// license_tier: INDEPENDENT_REDERIVE 30// genealogy_id: international-research-sources/nist/fips_203 31// lineage_id: nishi_kyber_sample_ntt_wasm_q1 32// safe_shift_audit: no 64-bit rotations in this module 33 34const KYBER_Q: i64 = 3329 35const KYBER_N: i64 = 256 36 37func _pstore(p: *u8, i: i64, v: i64) -> i64 { 38 var vv: i64 = v 39 if vv < 0 { vv = vv + 65536 } 40 p[i * 2] = vv & 0xff 41 p[i * 2 + 1] = (vv >> 8) & 0xff 42 return 0 43} 44 45// SampleNTT: parse 3 bytes at a time, accept candidates < q. 46// Returns bytes consumed (>=3*128 = 384 minimum) on success, -1 if 47// buf_len exhausted before 256 acceptances. 48func nx_kyber_sample_ntt(poly: *u8, buf: *u8, buf_len: i64) -> i64 { 49 var i: i64 = 0 // accepted coefficient index (0..256) 50 var pos: i64 = 0 // byte cursor into buf 51 52 while i < KYBER_N { 53 if pos + 3 > buf_len { return -1 } 54 let b0: i64 = buf[pos] 55 let b1: i64 = buf[pos + 1] 56 let b2: i64 = buf[pos + 2] 57 pos = pos + 3 58 59 let d1: i64 = b0 | ((b1 & 0x0f) << 8) 60 let d2: i64 = (b1 >> 4) | (b2 << 4) 61 62 if d1 < KYBER_Q { 63 _pstore(poly, i, d1) 64 i = i + 1 65 } 66 if i < KYBER_N { 67 if d2 < KYBER_Q { 68 _pstore(poly, i, d2) 69 i = i + 1 70 } 71 } 72 } 73 return pos 74}