ML-KEM Sender / Receiver — Primitives

The ML-KEM chiplet proves FIPS 203 key generation, encapsulation and decapsulation at ML-KEM-512, ML-KEM-768 and ML-KEM-1024, with seeds, keys and the shared secret kept in the witness. One pipeline proves up to 256 calls of any of the three kinds in a single proof. It shares hekate-pqc with the ML-DSA chiplet.

§ What the proof states

Each call runs one of the FIPS 203 internal algorithms:

Call	FIPS 203	In	Out
KeyGen	Algorithm 16, `ML-KEM.KeyGen_internal`	seed d	ek, H(ek), dk_PKE
Encaps	Algorithm 17, `ML-KEM.Encaps_internal`	ek, randomness m	shared key K, ciphertext c
Decaps	Algorithm 18, `ML-KEM.Decaps_internal`	dk, c	shared key K, valid

KeyGen leaves z to the caller, who appends it to form dk = dk_PKE ‖ ek ‖ H(ek) ‖ z. Decaps re-encrypts and compares. When c fails the comparison, it returns the implicit-rejection key J(z ‖ c) as FIPS 203 specifies and reports valid = false to the host program.

The FIPS 203 input checks are part of the statement. The Codec table rejects an encapsulation-key coefficient at or above q, the modulus check, and a decapsulation key whose stored H(ek) differs from the hash of its ek breaks the word bus, the hash check. Trace generation refuses both inputs before any proving starts.

The host program decides what becomes public. Both examples publish H(ek) and c and keep the seeds, the keys and K in the witness.

§ How the tables split the work

Encapsulation runs through the tables like this, per the hekate-pqc README:

text

FIPS 203 Encaps_internal(ek, m)               table
────────────────────────────────────────────  ───────────────
h ← H(ek); (K, r) ← G(m ‖ h)                  ctrl, Keccak
t̂ ← ByteDecode12(ek); μ ← Decompress1(m)      Codec
Â ← SampleNTT(ρ); y, e1, e2 ← CBD(r)          Sampler, Keccak
ŷ ← NTT(y); s ← e2 + μ                        NTT
û ← Âᵀ∘ŷ; v̂ ← t̂ᵀ∘ŷ                            PolyArith
u ← NTT⁻¹(û) + e1; v ← NTT⁻¹(v̂) + s           NTT
c1 ← Compress_du(u); c2 ← Compress_dv(v)      Codec

Decaps decrypts c on the same tables and re-encrypts, and a KemSelect table compares c with c′ and returns K′ or the implicit-rejection key K̄. KeyGen samples from G(d ‖ k) and encodes t̂ and ŝ in the Codec. The host program trades 32-bit words with the ctrl table over one service bus, and the tables pass coefficients, words and Keccak lanes over internal buses, each value emitted once by its producer and once by its consumer.

§ Outside the proof

Randomness. The seeds d and z and the encapsulation randomness m enter as witness, and the proof does not show they were drawn uniformly. Drawing them is the job of the external algorithms, ML-KEM.KeyGen (Algorithm 19) and ML-KEM.Encaps (Algorithm 20), which stay with the caller.

§ How it is tested

The tests run KeyGen, Encaps and Decaps at all three parameter sets, match their keys and ciphertexts against RustCrypto's ml-kem crate, and prove the calls. Forgery tests swap one input at a time, such as a key derived for another message, a foreign dk_PKE or a different ρ, and check that each swap breaks the word or lane bus that guards it. A non-canonical encapsulation key breaks the Codec's modulus check, and a decapsulation key with a foreign hash breaks the word bus. NIST's KAT files are not part of the suite.

§ Examples

Two programs model the two sides of a key exchange at the level HEKATE_LEVEL selects:

mlkem_sender.rs proves Encaps of K to the recipient's ek and an AES-256-CTR encryption of a message under K, joined by the key_link bus that carries K into the AES key. It publishes H(ek), c, the 12-byte nonce of each counter block, the message and its ciphertext; ek, m and K stay private. The program also checks that the recipient recovers K from c and decrypts the message.
mlkem_receiver.rs proves KeyGen from the seed d chained into Decaps of a sender's c, through the dk_link bus. It publishes H(ek), c and valid; d, z and dk stay private. The program also checks that the K it decapsulates matches the sender's.

From a hekate checkout:

bash

HEKATE_LEVEL=768 just example mlkem_sender public     # HEKATE_MESSAGE=<text> sets the payload
HEKATE_LEVEL=768 just example mlkem_receiver public   # 512 | 768 | 1024

public builds the variable-time table-math prover that the README's measurements use. Without it the recipe builds the constant-time prover, the one to use when the witness is private.

The programs below are the two examples at ML-KEM-768, as two crates, mlkem-sender and mlkem-receiver. Each takes three files: program.rs declares the host table, the tables it attaches and their buses, and publishes the statement's values; trace.rs runs the host side and fills the tables; main.rs proves and verifies.

Both links are permutation buses without clocks, which Hekate accepts only with a clock waiver: the written reason their keys are unique, here fixed indices that each side emits once. The engine checks that the reason is present; review checks the argument.

Start from the crates in Installation and add hekate-pqc, ml-kem and rand, plus hekate-aes for the sender.

§ Sender

src/program.rssrc/trace.rssrc/main.rs

rust

1use hekate_aes::{Aes256Chiplet, AesRound256Air, host_key_selector_shape, host_selector_shape};
2use hekate_core::config::Config;
3use hekate_core::errors;
4use hekate_math::{Block128, TowerField};
5use hekate_pqc::mlkem::{self, MLKEM_DATA_BUS_ID, MlKemChiplet, MlKemParams};
6use hekate_program::circuit::{Circuit, CircuitProgram, Col, ColRange};
7use hekate_program::constraint::builder::ConstraintSystem;
8use hekate_program::permutation::{PermutationCheckSpec, Source};
9use hekate_program::{FixedShape, define_columns};
10
11pub const PARAMS: MlKemParams = MlKemParams::ML_KEM_768;
12pub const SEED_WORDS: usize = 8;
13pub const BLOCK: usize = 16;
14pub const NONCE_BYTES: usize = 12;
15
16const KEY_LINK_BUS_ID: &str = "key_link";
17const KEY_LINK_WAIVER: &str = "see mlkem-sender/src/program.rs: \
18     link_idx is a fixed column over K's host rows and a constant \
19     on the first AES row; each index is emitted once on each side";
20
21pub type F = Block128;
22pub type Published = (usize, usize, F);
23
24define_columns! {
25    pub SenderColumns {
26        WORD: B32,
27        SEL: Bit,
28
29        LINK_IDX: B16,
30        LINK_SEL: Bit,
31
32        KEY: [B8; 32],
33        KEY_SELECTOR: Bit,
34        DATA: [B8; 16],
35        SELECTOR: Bit,
36
37        PACKED: [B32; 8],
38        FIRST: Bit,
39
40        CARRY: Bit,
41
42        COUNTER: [B8; 16],
43        PLAINTEXT: [B8; 16],
44        CIPHERTEXT: [B8; 16],
45    }
46}
47
48pub struct Layout {
49    words: usize,
50    pub h: usize,
51    pub key: usize,
52    pub c: usize,
53    pub c_words: usize,
54}
55
56impl Layout {
57    pub fn new(params: MlKemParams) -> Self {
58        let m = 96 * params.k() + SEED_WORDS;
59        let h = m + SEED_WORDS;
60        let key = h + SEED_WORDS;
61        let c = key + SEED_WORDS;
62        let c_words = 8 * (params.du() as usize * params.k() + params.dv() as usize);
63
64        Self {
65            words: c + c_words,
66            h,
67            key,
68            c,
69            c_words,
70        }
71    }
72}
73
74pub fn aes_chiplet(blocks: usize, config: &Config) -> errors::Result<Aes256Chiplet> {
75    let floor = config.min_table_rows();
76
77    Aes256Chiplet::new(
78        (blocks * AesRound256Air::BLOCK_ROWS)
79            .next_power_of_two()
80            .max(floor),
81        (blocks * AesRound256Air::ACTIVE_ROWS)
82            .next_power_of_two()
83            .max(floor),
84        blocks,
85    )
86}
87
88pub fn build(
89    pipeline: &MlKemChiplet<F>,
90    aes: &Aes256Chiplet,
91    layout: &Layout,
92    message_len: usize,
93    rows: usize,
94    statement: &[Published],
95) -> errors::Result<CircuitProgram<F>> {
96    let blocks = message_len.div_ceil(BLOCK);
97
98    let mut cx = Circuit::<F>::new("MlKemSender", rows)?;
99
100    let cols = cx.schema(&SenderColumns::build_layout());
101
102    let at = |col: usize| cols.at(col);
103
104    let word = at(SenderColumns::WORD);
105    let key: [Col; 32] = core::array::from_fn(|j| at(SenderColumns::KEY + j));
106    let data: [Col; BLOCK] = core::array::from_fn(|j| at(SenderColumns::DATA + j));
107
108    fix_shapes(&mut cx, cols, layout, blocks);
109
110    cx.call(&mlkem::service(), &[word], at(SenderColumns::SEL))?;
111    cx.call(
112        &AesRound256Air::link_service(),
113        &data,
114        at(SenderColumns::SELECTOR),
115    )?;
116    cx.call(
117        &AesRound256Air::key_service(),
118        &key,
119        at(SenderColumns::KEY_SELECTOR),
120    )?;
121
122    let link_idx = Source::Column(at(SenderColumns::LINK_IDX).index());
123
124    cx.bus(
125        KEY_LINK_BUS_ID,
126        key_link(link_idx, word, at(SenderColumns::LINK_SEL)),
127    );
128
129    for j in 0..SEED_WORDS {
130        let packed = at(SenderColumns::PACKED + j);
131
132        cx.bus(
133            KEY_LINK_BUS_ID,
134            key_link(Source::Const(j as u128), packed, at(SenderColumns::FIRST)),
135        );
136    }
137
138    add_constraints(cx.cs(), cols);
139
140    for &(c, row, _) in statement {
141        cx.publish(at(c), row);
142    }
143
144    let pad = (BLOCK - message_len % BLOCK) % BLOCK;
145
146    for j in BLOCK - pad..BLOCK {
147        cx.boundary(at(SenderColumns::PLAINTEXT + j), 2 * blocks - 1, F::ZERO);
148    }
149
150    cx.attach_namespaced("mlkem", pipeline.defs()?, &[MLKEM_DATA_BUS_ID])?;
151    cx.attach_namespaced("aes256", aes.defs()?, &Aes256Chiplet::EXTERNAL_BUS_IDS)?;
152
153    cx.compile()
154}
155
156fn fix_shapes(cx: &mut Circuit<F>, cols: ColRange, layout: &Layout, blocks: usize) {
157    let sel = FixedShape::Cadence {
158        stride: 1,
159        count: layout.words,
160        origin: 0,
161        values: vec![F::ONE],
162    };
163    let link_idx = (1..SEED_WORDS).map(|j| (layout.key + j, F::from(j as u32)));
164    let link_sel = (0..SEED_WORDS).map(|j| (layout.key + j, F::ONE));
165    let carry = FixedShape::Cadence {
166        stride: 1,
167        count: 2 * blocks - 1,
168        origin: 0,
169        values: vec![F::ONE],
170    };
171
172    cx.fix(cols.at(SenderColumns::SEL), sel);
173    cx.fix(
174        cols.at(SenderColumns::LINK_IDX),
175        FixedShape::Sparse(link_idx.collect()),
176    );
177    cx.fix(
178        cols.at(SenderColumns::LINK_SEL),
179        FixedShape::Sparse(link_sel.collect()),
180    );
181    cx.fix(
182        cols.at(SenderColumns::SELECTOR),
183        host_selector_shape(2, blocks),
184    );
185    cx.fix(
186        cols.at(SenderColumns::KEY_SELECTOR),
187        host_key_selector_shape(2, blocks),
188    );
189    cx.fix(
190        cols.at(SenderColumns::FIRST),
191        FixedShape::Sparse(vec![(0, F::ONE)]),
192    );
193    cx.fix(cols.at(SenderColumns::CARRY), carry);
194
195    for i in 0..BLOCK - NONCE_BYTES {
196        let index = (0..blocks).map(|b| (2 * b, F::from((b as u32).to_be_bytes()[i])));
197
198        cx.fix(
199            cols.at(SenderColumns::COUNTER + NONCE_BYTES + i),
200            FixedShape::Sparse(index.collect()),
201        );
202    }
203}
204
205fn add_constraints(cs: &ConstraintSystem<F>, cols: ColRange) {
206    let col = |c: usize| cs.col(cols.at(c).index());
207    let next = |c: usize| cs.next(cols.at(c).index());
208
209    let one = cs.one();
210    let carry = col(SenderColumns::CARRY);
211    let selector = col(SenderColumns::SELECTOR);
212    let key_selector = col(SenderColumns::KEY_SELECTOR);
213    let output = selector + key_selector;
214
215    cs.constrain((one + col(SenderColumns::SEL)) * col(SenderColumns::WORD));
216
217    for j in 0..SEED_WORDS {
218        let packed = (0..4).fold(col(SenderColumns::PACKED + j), |acc, i| {
219            let shift = cs.constant(F::from(1u128 << (8 * i)));
220
221            acc + shift * col(SenderColumns::KEY + 4 * j + i)
222        });
223
224        cs.constrain(packed);
225    }
226
227    for j in 0..32 {
228        let key = SenderColumns::KEY + j;
229
230        cs.constrain(carry * (next(key) + col(key)));
231        cs.constrain((one + selector) * col(key));
232    }
233
234    for j in 0..BLOCK {
235        let data = col(SenderColumns::DATA + j);
236        let counter = col(SenderColumns::COUNTER + j);
237        let plaintext = col(SenderColumns::PLAINTEXT + j);
238        let ciphertext = col(SenderColumns::CIPHERTEXT + j);
239
240        let whitened = data + counter + col(SenderColumns::KEY + j);
241        let encrypted = ciphertext + plaintext + data;
242
243        cs.constrain(key_selector * whitened);
244        cs.constrain(output * encrypted);
245
246        cs.constrain((one + selector) * data);
247        cs.constrain((one + key_selector) * counter);
248        cs.constrain((one + output) * plaintext);
249        cs.constrain((one + output) * ciphertext);
250    }
251}
252
253fn key_link(idx: Source, word: Col, selector: Col) -> PermutationCheckSpec {
254    PermutationCheckSpec::new(
255        vec![
256            (idx, b"kappa_key_idx" as &[u8]),
257            (Source::Column(word.index()), b"kappa_key_word" as &[u8]),
258        ],
259        Some(selector.index()),
260    )
261    .with_clock_waiver(KEY_LINK_WAIVER)
262}

1use crate::program::{BLOCK, F, Layout, NONCE_BYTES, Published, SEED_WORDS, SenderColumns};
2use hekate_aes::Aes256Chiplet;
3use hekate_aes::trace::{Aes256Call, aes256_encrypt_block, expand_key_256};
4use hekate_core::errors;
5use hekate_core::trace::{ColumnTrace, TraceBuilder};
6use hekate_math::{Bit, Block8, Block16, Block32, TowerField};
7use std::error::Error;
8
9pub struct Ctr {
10    pub counters: Vec<[u8; BLOCK]>,
11    pub keystream: Vec<[u8; BLOCK]>,
12    pub plaintext: Vec<u8>,
13    pub ciphertext: Vec<u8>,
14}
15
16impl Ctr {
17    pub fn encrypt(key: &[u8; 32], nonce: &[u8; NONCE_BYTES], message: &[u8]) -> Self {
18        let round_keys = expand_key_256(key);
19        let blocks = message.len().div_ceil(BLOCK);
20
21        let counters: Vec<[u8; BLOCK]> = (0..blocks as u32)
22            .map(|b| {
23                let mut block = [0u8; BLOCK];
24                block[..NONCE_BYTES].copy_from_slice(nonce);
25                block[NONCE_BYTES..].copy_from_slice(&b.to_be_bytes());
26
27                block
28            })
29            .collect();
30
31        let keystream: Vec<[u8; BLOCK]> = counters
32            .iter()
33            .map(|counter| aes256_encrypt_block(&round_keys, counter))
34            .collect();
35
36        let mut plaintext = message.to_vec();
37        plaintext.resize(blocks * BLOCK, 0);
38
39        let ciphertext = plaintext
40            .iter()
41            .zip(keystream.iter().flatten())
42            .map(|(p, k)| p ^ k)
43            .collect();
44
45        Self {
46            counters,
47            keystream,
48            plaintext,
49            ciphertext,
50        }
51    }
52
53    pub fn blocks(&self) -> usize {
54        self.counters.len()
55    }
56}
57
58pub fn statement(
59    layout: &Layout,
60    words: &[u32],
61    ctr: &Ctr,
62    message_len: usize,
63) -> Result<Vec<Published>, Box<dyn Error>> {
64    let h = layout.h..layout.h + SEED_WORDS;
65    let c = layout.c..layout.c + layout.c_words;
66
67    let mut statement = Vec::new();
68
69    for row in h.chain(c) {
70        let word = words
71            .get(row)
72            .ok_or("the host words end before the published rows")?;
73
74        statement.push((SenderColumns::WORD, row, F::from(u128::from(*word))));
75    }
76
77    for (b, counter) in ctr.counters.iter().enumerate() {
78        for (j, &byte) in counter.iter().take(NONCE_BYTES).enumerate() {
79            statement.push((SenderColumns::COUNTER + j, 2 * b, F::from(byte)));
80        }
81    }
82
83    let bytes = ctr.plaintext.iter().zip(&ctr.ciphertext).take(message_len);
84
85    for (i, (&plain, &cipher)) in bytes.enumerate() {
86        let (row, j) = (2 * (i / BLOCK) + 1, i % BLOCK);
87
88        statement.push((SenderColumns::PLAINTEXT + j, row, F::from(plain)));
89        statement.push((SenderColumns::CIPHERTEXT + j, row, F::from(cipher)));
90    }
91
92    Ok(statement)
93}
94
95pub fn host_trace(
96    words: &[u32],
97    layout: &Layout,
98    key: &[u8; 32],
99    ctr: &Ctr,
100    rows: usize,
101) -> errors::Result<ColumnTrace> {
102    let mut tb = TraceBuilder::new_secret(
103        &SenderColumns::build_layout(),
104        rows.trailing_zeros() as usize,
105    )?;
106
107    for (r, &word) in words.iter().enumerate() {
108        tb.set_b32(SenderColumns::WORD, r, Block32::from(word))?;
109        tb.set_bit(SenderColumns::SEL, r, Bit::ONE)?;
110    }
111
112    for j in 0..SEED_WORDS {
113        tb.set_b16(SenderColumns::LINK_IDX, layout.key + j, Block16(j as u16))?;
114        tb.set_bit(SenderColumns::LINK_SEL, layout.key + j, Bit::ONE)?;
115    }
116
117    let aes_rows = 2 * ctr.blocks();
118    let (packed, _) = key.as_chunks::<4>();
119
120    for r in 0..aes_rows {
121        for (j, &byte) in key.iter().enumerate() {
122            tb.set_b8(SenderColumns::KEY + j, r, Block8(byte))?;
123        }
124
125        for (j, &bytes) in packed.iter().enumerate() {
126            let word = Block32::from(u32::from_le_bytes(bytes));
127
128            tb.set_b32(SenderColumns::PACKED + j, r, word)?;
129        }
130
131        tb.set_bit(SenderColumns::SELECTOR, r, Bit::ONE)?;
132        tb.set_bit(
133            SenderColumns::CARRY,
134            r,
135            Bit::from(u8::from(r + 1 < aes_rows)),
136        )?;
137    }
138
139    tb.set_bit(SenderColumns::FIRST, 0, Bit::ONE)?;
140
141    for (b, counter) in ctr.counters.iter().enumerate() {
142        tb.set_bit(SenderColumns::KEY_SELECTOR, 2 * b, Bit::ONE)?;
143
144        for (j, (&byte, &k)) in counter.iter().zip(key).enumerate() {
145            tb.set_b8(SenderColumns::COUNTER + j, 2 * b, Block8(byte))?;
146            tb.set_b8(SenderColumns::DATA + j, 2 * b, Block8(byte ^ k))?;
147        }
148    }
149
150    let output = ctr
151        .keystream
152        .iter()
153        .flatten()
154        .zip(ctr.plaintext.iter())
155        .zip(&ctr.ciphertext);
156
157    for (i, ((&keystream, &plain), &cipher)) in output.enumerate() {
158        let (row, j) = (2 * (i / BLOCK) + 1, i % BLOCK);
159
160        tb.set_b8(SenderColumns::DATA + j, row, Block8(keystream))?;
161        tb.set_b8(SenderColumns::PLAINTEXT + j, row, Block8(plain))?;
162        tb.set_b8(SenderColumns::CIPHERTEXT + j, row, Block8(cipher))?;
163    }
164
165    Ok(tb.build())
166}
167
168pub fn aes_traces(
169    aes: &Aes256Chiplet,
170    key: &[u8; 32],
171    ctr: &Ctr,
172) -> errors::Result<Vec<ColumnTrace>> {
173    let round_keys = expand_key_256(key);
174
175    let calls: Vec<Aes256Call> = ctr
176        .counters
177        .iter()
178        .map(|counter| Aes256Call {
179            key: *key,
180            plaintext: *counter,
181            round_keys,
182        })
183        .collect();
184
185    aes.generate_traces(&calls)
186}

1mod program;
2mod trace;
3
4use hekate_core::config::Config;
5use hekate_crypto::DefaultHasher;
6use hekate_crypto::transcript::Transcript;
7use hekate_pqc::mlkem::{MlKemCall, MlKemChiplet, MlKemInput, MlKemOutput};
8use hekate_program::digest::program_id;
9use hekate_program::{ProgramInstance, ProgramWitness};
10use hekate_prover_sys::prove;
11use hekate_sdk::preflight;
12use hekate_verifier::HekateVerifier;
13use ml_kem::{Decapsulate, DecapsulationKey768, Generate, KeyExport};
14use program::{F, Layout, NONCE_BYTES, PARAMS, aes_chiplet, build};
15use rand::{TryRng, rngs::SysRng};
16use std::error::Error;
17use trace::{Ctr, aes_traces, host_trace, statement};
18
19const LABEL: &[u8] = b"mlkem-sender";
20const MESSAGE: &str =
21    "Hekate ML-KEM sender: this message travels under AES-256-CTR with the encapsulated key.";
22
23fn main() -> Result<(), Box<dyn Error>> {
24    let recipient = DecapsulationKey768::try_generate_from_rng(&mut SysRng)?;
25    let ek = recipient.encapsulation_key().to_bytes();
26    let message = MESSAGE.as_bytes();
27
28    let mut m = [0; 32];
29    let mut nonce = [0; NONCE_BYTES];
30
31    SysRng.try_fill_bytes(&mut m)?;
32    SysRng.try_fill_bytes(&mut nonce)?;
33
34    let pipeline = MlKemChiplet::<F>::new(PARAMS, &[MlKemCall::Encaps])?;
35    let traced = pipeline.trace(&[MlKemInput::Encaps { ek: &ek, m: &m }])?;
36
37    let [MlKemOutput::Encaps { key, c }] = traced.outputs.as_slice() else {
38        return Err("the pipeline returns one Encaps output".into());
39    };
40
41    let config = Config::default();
42    let ctr = Ctr::encrypt(key, &nonce, message);
43    let aes = aes_chiplet(ctr.blocks(), &config)?;
44
45    let layout = Layout::new(PARAMS);
46    let rows = traced.words.len().max(2 * ctr.blocks()).next_power_of_two();
47    let statement = statement(&layout, &traced.words, &ctr, message.len())?;
48    let program = build(&pipeline, &aes, &layout, message.len(), rows, &statement)?;
49
50    let inputs = statement.iter().map(|&(_, _, value)| value);
51    let instance = ProgramInstance::new(rows, inputs.collect());
52
53    let mut chiplets = traced.traces;
54    chiplets.extend(aes_traces(&aes, key, &ctr)?);
55
56    let host = host_trace(&traced.words, &layout, key, &ctr, rows)?;
57    let witness = ProgramWitness::new(host).with_chiplets(chiplets);
58
59    let report = preflight(&program, &instance, &witness)?;
60
61    if !report.is_clean() {
62        eprintln!("{report}");
63
64        return Err("the trace fails preflight".into());
65    }
66
67    let mut seed = [0; 32];
68    SysRng.try_fill_bytes(&mut seed)?;
69
70    let proof = prove(LABEL, &program, &instance, &witness, &config, seed, None)?;
71
72    let id = program_id::<F, _>(&program)?;
73    let mut transcript = Transcript::<DefaultHasher>::new(LABEL);
74
75    let accepted = HekateVerifier::<F, DefaultHasher>::verify(
76        &id,
77        &program,
78        &instance,
79        &proof,
80        &mut transcript,
81        &config,
82    )?;
83
84    if !accepted {
85        return Err("the verifier rejects the encapsulation and the encryption".into());
86    }
87
88    let received = recipient.decapsulate_slice(c)?;
89    let decrypted = Ctr::encrypt(received.as_ref(), &nonce, &ctr.ciphertext[..message.len()]);
90
91    if *key != received || !decrypted.ciphertext.starts_with(message) {
92        return Err("the recipient misses K or the message".into());
93    }
94
95    println!("c          {} bytes, public", c.len());
96    println!("recipient  recovers K and decrypts the message");
97    println!("the verifier accepts the encapsulation and the encryption");
98
99    Ok(())
100}

§ Receiver

src/program.rssrc/trace.rssrc/main.rs

rust

1use hekate_core::errors;
2use hekate_math::{Block128, TowerField};
3use hekate_pqc::mlkem::{self, MLKEM_DATA_BUS_ID, MlKemChiplet, MlKemParams};
4use hekate_program::circuit::{Circuit, CircuitProgram, Col};
5use hekate_program::permutation::{PermutationCheckSpec, Source};
6use hekate_program::{FixedShape, define_columns};
7
8pub const PARAMS: MlKemParams = MlKemParams::ML_KEM_768;
9
10const SEED_WORDS: usize = 8;
11
12const LINK_BUS_ID: &str = "dk_link";
13const LINK_WAIVER: &str = "see mlkem-receiver/src/program.rs: \
14     link_idx is a fixed column; every index is emitted once \
15     on a KeyGen output row and once on a Decaps input row";
16
17pub type F = Block128;
18
19define_columns! {
20    pub ReceiverColumns {
21        WORD: B32,
22        SEL: Bit,
23
24        LINK_IDX: B16,
25        LINK_SEL: Bit,
26    }
27}
28
29pub struct Layout {
30    words: usize,
31    h: usize,
32    c: usize,
33    c_words: usize,
34    valid: usize,
35    pub links: Vec<(usize, u16)>,
36}
37
38impl Layout {
39    pub fn new(params: MlKemParams) -> Self {
40        let w = 96 * params.k();
41        let c_words = 8 * (params.du() as usize * params.k() + params.dv() as usize);
42
43        let ek_out = SEED_WORDS;
44        let rho_out = ek_out + w;
45        let h_out = rho_out + SEED_WORDS;
46        let dk_out = h_out + SEED_WORDS;
47
48        let dk_in = dk_out + w;
49        let ek_in = dk_in + w;
50        let rho_in = ek_in + w;
51        let h_in = rho_in + SEED_WORDS;
52        let c = h_in + 2 * SEED_WORDS;
53        let valid = c + c_words + SEED_WORDS;
54
55        let mut links = Vec::with_capacity(4 * w + 4 * SEED_WORDS);
56
57        for (keygen, decaps, first, len) in [
58            (dk_out, dk_in, 0, w),
59            (ek_out, ek_in, w, w),
60            (rho_out, rho_in, 2 * w, SEED_WORDS),
61            (h_out, h_in, 2 * w + SEED_WORDS, SEED_WORDS),
62        ] {
63            for j in 0..len {
64                let idx = (first + j) as u16;
65
66                links.push((keygen + j, idx));
67                links.push((decaps + j, idx));
68            }
69        }
70
71        links.sort_unstable();
72
73        Self {
74            words: valid + 1,
75            h: h_out,
76            c,
77            c_words,
78            valid,
79            links,
80        }
81    }
82
83    pub fn published(&self) -> Vec<usize> {
84        (self.h..self.h + SEED_WORDS)
85            .chain(self.c..self.c + self.c_words)
86            .chain([self.valid])
87            .collect()
88    }
89}
90
91pub fn build(
92    pipeline: &MlKemChiplet<F>,
93    layout: &Layout,
94    rows: usize,
95) -> errors::Result<CircuitProgram<F>> {
96    let mut cx = Circuit::<F>::new("MlKemReceiver", rows)?;
97
98    let cols = cx.schema(&ReceiverColumns::build_layout());
99
100    let word = cols.at(ReceiverColumns::WORD);
101    let sel = cols.at(ReceiverColumns::SEL);
102    let link_idx = cols.at(ReceiverColumns::LINK_IDX);
103    let link_sel = cols.at(ReceiverColumns::LINK_SEL);
104
105    let indices = layout
106        .links
107        .iter()
108        .map(|&(row, idx)| (row, F::from(u32::from(idx))))
109        .filter(|&(_, v)| v != F::ZERO);
110
111    cx.fix(
112        sel,
113        FixedShape::Cadence {
114            stride: 1,
115            count: layout.words,
116            origin: 0,
117            values: vec![F::ONE],
118        },
119    );
120    cx.fix(link_idx, FixedShape::Sparse(indices.collect()));
121    cx.fix(
122        link_sel,
123        FixedShape::Sparse(layout.links.iter().map(|&(row, _)| (row, F::ONE)).collect()),
124    );
125
126    let cs = cx.cs();
127    let col = |c: Col| cs.col(c.index());
128
129    cs.constrain((cs.one() + col(sel)) * col(word));
130
131    cx.call(&mlkem::service(), &[word], sel)?;
132
133    cx.bus(
134        LINK_BUS_ID,
135        PermutationCheckSpec::new(
136            vec![
137                (Source::Column(link_idx.index()), b"kappa_link_idx" as &[u8]),
138                (Source::Column(word.index()), b"kappa_link_word" as &[u8]),
139            ],
140            Some(link_sel.index()),
141        )
142        .with_clock_waiver(LINK_WAIVER),
143    );
144
145    for row in layout.published() {
146        cx.publish(word, row);
147    }
148
149    cx.attach_namespaced("mlkem", pipeline.defs()?, &[MLKEM_DATA_BUS_ID])?;
150
151    cx.compile()
152}

1use crate::program::{F, Layout, ReceiverColumns};
2use hekate_core::errors;
3use hekate_core::trace::{ColumnTrace, TraceBuilder};
4use hekate_math::{Bit, Block16, Block32, TowerField};
5use hekate_pqc::mlkem::{MlKemCall, MlKemChiplet, MlKemInput, MlKemOutput, MlKemParams};
6use std::error::Error;
7
8pub struct KeyPair {
9    pub ek: Vec<u8>,
10    pub dk: Vec<u8>,
11}
12
13impl KeyPair {
14    pub fn generate(
15        params: MlKemParams,
16        d: &[u8; 32],
17        z: &[u8; 32],
18    ) -> Result<Self, Box<dyn Error>> {
19        let keygen = MlKemChiplet::<F>::new(params, &[MlKemCall::KeyGen])?;
20        let generated = keygen.trace(&[MlKemInput::KeyGen { d }])?;
21
22        let [MlKemOutput::KeyGen { ek, dk_pke, h }] = generated.outputs.as_slice() else {
23            return Err("the KeyGen pipeline returns one KeyGen output".into());
24        };
25
26        Ok(Self {
27            ek: ek.clone(),
28            dk: [dk_pke.as_slice(), ek, h, z].concat(),
29        })
30    }
31}
32
33pub fn host_trace(words: &[u32], layout: &Layout, rows: usize) -> errors::Result<ColumnTrace> {
34    let mut tb = TraceBuilder::new_secret(
35        &ReceiverColumns::build_layout(),
36        rows.trailing_zeros() as usize,
37    )?;
38
39    for (r, &word) in words.iter().enumerate() {
40        tb.set_b32(ReceiverColumns::WORD, r, Block32::from(word))?;
41        tb.set_bit(ReceiverColumns::SEL, r, Bit::ONE)?;
42    }
43
44    for &(row, idx) in &layout.links {
45        tb.set_b16(ReceiverColumns::LINK_IDX, row, Block16(idx))?;
46        tb.set_bit(ReceiverColumns::LINK_SEL, row, Bit::ONE)?;
47    }
48
49    Ok(tb.build())
50}

1mod program;
2mod trace;
3
4use hekate_core::config::Config;
5use hekate_crypto::DefaultHasher;
6use hekate_crypto::transcript::Transcript;
7use hekate_pqc::mlkem::{MlKemCall, MlKemChiplet, MlKemInput, MlKemOutput};
8use hekate_program::digest::program_id;
9use hekate_program::{ProgramInstance, ProgramWitness};
10use hekate_prover_sys::prove;
11use hekate_sdk::preflight;
12use hekate_verifier::HekateVerifier;
13use ml_kem::{Encapsulate, EncapsulationKey768, TryKeyInit};
14use program::{F, Layout, PARAMS, build};
15use rand::rngs::{StdRng, SysRng};
16use rand::{SeedableRng, TryRng};
17use std::error::Error;
18use trace::{KeyPair, host_trace};
19
20const LABEL: &[u8] = b"mlkem-receiver";
21
22fn main() -> Result<(), Box<dyn Error>> {
23    let mut d = [0; 32];
24    let mut z = [0; 32];
25
26    SysRng.try_fill_bytes(&mut d)?;
27    SysRng.try_fill_bytes(&mut z)?;
28
29    let keys = KeyPair::generate(PARAMS, &d, &z)?;
30
31    let sender = EncapsulationKey768::new_from_slice(&keys.ek)?;
32    let (c, sent) = sender.encapsulate_with_rng(&mut StdRng::try_from_rng(&mut SysRng)?);
33
34    let pipeline = MlKemChiplet::<F>::new(PARAMS, &[MlKemCall::KeyGen, MlKemCall::Decaps])?;
35    let traced = pipeline.trace(&[
36        MlKemInput::KeyGen { d: &d },
37        MlKemInput::Decaps {
38            dk: &keys.dk,
39            c: &c,
40        },
41    ])?;
42
43    let [_, MlKemOutput::Decaps { key, valid }] = traced.outputs.as_slice() else {
44        return Err("the pipeline returns KeyGen then Decaps outputs".into());
45    };
46
47    if *key != sent {
48        return Err("the decapsulated key differs from the sender's".into());
49    }
50
51    let config = Config::default();
52    let layout = Layout::new(PARAMS);
53    let rows = traced.words.len().next_power_of_two();
54    let program = build(&pipeline, &layout, rows)?;
55
56    let public_inputs = layout
57        .published()
58        .iter()
59        .map(|&row| traced.words.get(row).map(|&word| F::from(u128::from(word))))
60        .collect::<Option<Vec<F>>>()
61        .ok_or("the host words end before the published rows")?;
62    let instance = ProgramInstance::new(rows, public_inputs);
63
64    let host = host_trace(&traced.words, &layout, rows)?;
65    let witness = ProgramWitness::new(host).with_chiplets(traced.traces);
66
67    let report = preflight(&program, &instance, &witness)?;
68
69    if !report.is_clean() {
70        eprintln!("{report}");
71
72        return Err("the trace fails preflight".into());
73    }
74
75    let mut seed = [0; 32];
76    SysRng.try_fill_bytes(&mut seed)?;
77
78    let proof = prove(LABEL, &program, &instance, &witness, &config, seed, None)?;
79
80    let id = program_id::<F, _>(&program)?;
81    let mut transcript = Transcript::<DefaultHasher>::new(LABEL);
82
83    let accepted = HekateVerifier::<F, DefaultHasher>::verify(
84        &id,
85        &program,
86        &instance,
87        &proof,
88        &mut transcript,
89        &config,
90    )?;
91
92    if !accepted {
93        return Err("the verifier rejects the key generation and the decapsulation".into());
94    }
95
96    println!("c          {} bytes, public", c.len());
97    println!("valid      {valid}, K matches the sender's");
98    println!("the verifier accepts the key generation and the decapsulation");
99
100    Ok(())
101}

§ Performance

Zero-knowledge proofs, measured on Apple M3 Max, 16 cores. Build: --release, std parallel blake3 table-math public. Best of three on an otherwise idle machine. The Benchmarks page has the base-protocol figures and the command for each run.

Workload	Input	Proving	Verify	Proof size	Peak memory
ML-KEM-768 sender	FIPS 203	413 ms	11.7 ms	1,541 KiB	190 MiB
ML-KEM-768 receiver	FIPS 203	558 ms	13.9 ms	1,884 KiB	269 MiB

§ Security

The crate has not been audited and may contain bugs and security flaws, as its README states. Hekate proves 100 bits of soundness at every parameter set, below the security level of each, which makes the proof the weaker link; the proven operations still run the full FIPS 203 parameter sets.

§ License

hekate-pqc is AGPL-3.0-only. Commercial licenses are available from Oumuamua Labs at info@oumuamua.dev. The README's Usage section builds an ML-KEM-768 encapsulation host program, and the API reference is on docs.rs.