AES-128 and AES-256 Encryption — Primitives

The AES chiplet proves FIPS 197 block encryption with AES-128 and AES-256: for every block, the ciphertext is the encryption of the block's input under the key, key expansion included.

§ What the proof states

The round table implements the FIPS 197 round function, SubBytes, ShiftRows, MixColumns and AddRoundKey, as a binary-field AIR. Its SubBytes values come from an S-box ROM table that proves the GF(2^8) inversion and the affine map, and the key schedule, SubWord included, is constrained in the round AIR itself. The host program sends each block's input and its key to the round table over LogUp buses and receives the ciphertext. A block's input is its plaintext XORed with round key 0, the first AddRoundKey, which the chiplet does not check: the host program must.

The host program decides what becomes public. The README's Usage section publishes nothing, and the ML-KEM sender publishes a message and its AES-256-CTR ciphertext while the key stays private.

§ Outside the proof

AES-192.
Decryption.
Modes of operation. The ML-KEM sender builds AES-256-CTR in its host program on top of the block chiplet.

§ Constant-time trace generation

Without table-math, the default, trace generation never indexes a table with a secret. The S-box is field arithmetic, the GF(2^8) inverse x²⁵⁴ followed by the FIPS 197 affine map, with no key-dependent index, branch or memory access, which keeps the key out of reach of cache-timing observation.

§ Examples

hekate/examples/aes.rs proves and verifies AES-128 or AES-256 over a run of blocks, at the key size HEKATE_LEVEL selects. From a hekate checkout:

bash

HEKATE_LEVEL=256 just example aes public   # 128 | 256

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 program below turns the example into a statement a verifier can use: one AES-128 key, kept private, encrypts each of four published plaintext blocks to its published ciphertext block. Block 0 is the FIPS 197 Appendix B example. The program takes three files:

src/program.rs attaches the round table and the S-box ROM, keeps the key unchanged on every host row, constrains each block's input to its plaintext XOR the key (the first AddRoundKey) and publishes every plaintext and ciphertext block.
src/trace.rs expands the key, reads the ciphertexts from the round table and fills the host table.
src/main.rs proves and verifies.

Start from the crates in Installation and add hekate-aes and rand.

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

rust

1use hekate_aes::{Aes128Chiplet, AesRound128Air, host_key_selector_shape, host_selector_shape};
2use hekate_core::config::Config;
3use hekate_core::errors;
4use hekate_math::Block128;
5use hekate_program::circuit::{Circuit, CircuitProgram, Col, ColRange};
6use hekate_program::constraint::builder::ConstraintSystem;
7use hekate_program::define_columns;
8
9pub const NUM_BLOCKS: usize = 4;
10pub const CPU_IO_PER_BLOCK: usize = 2;
11
12pub type F = Block128;
13
14define_columns! {
15    pub CpuColumns {
16        KEY: [B8; 16],
17        KEY_SELECTOR: Bit,
18        DATA: [B8; 16],
19        SELECTOR: Bit,
20        PLAINTEXT: [B8; 16],
21    }
22}
23
24pub const fn in_row(block: usize) -> usize {
25    block * CPU_IO_PER_BLOCK
26}
27
28pub const fn out_row(block: usize) -> usize {
29    in_row(block) + 1
30}
31
32pub fn table_rows(config: &Config, rows_per_block: usize) -> usize {
33    (NUM_BLOCKS * rows_per_block)
34        .next_power_of_two()
35        .max(config.min_table_rows())
36}
37
38pub fn aes128_chiplet(config: &Config) -> errors::Result<Aes128Chiplet> {
39    Aes128Chiplet::new(
40        table_rows(config, AesRound128Air::BLOCK_ROWS),
41        table_rows(config, AesRound128Air::ACTIVE_ROWS),
42        NUM_BLOCKS,
43    )
44}
45
46pub fn build_aes128_program(
47    cpu_rows: usize,
48    aes: &Aes128Chiplet,
49) -> errors::Result<CircuitProgram<F>> {
50    let mut cx = Circuit::<F>::new("Aes128", cpu_rows)?;
51    let cpu = cx.schema(&CpuColumns::build_layout());
52
53    let selector = cpu.at(CpuColumns::SELECTOR);
54    let key_selector = cpu.at(CpuColumns::KEY_SELECTOR);
55
56    let link_values: Vec<Col> = (0..16).map(|j| cpu.at(CpuColumns::DATA + j)).collect();
57
58    cx.call(&AesRound128Air::link_service(), &link_values, selector)?;
59
60    let key_values: Vec<Col> = (0..16).map(|j| cpu.at(CpuColumns::KEY + j)).collect();
61
62    cx.call(&AesRound128Air::key_service(), &key_values, key_selector)?;
63
64    cx.fix(selector, host_selector_shape(CPU_IO_PER_BLOCK, NUM_BLOCKS));
65    cx.fix(
66        key_selector,
67        host_key_selector_shape(CPU_IO_PER_BLOCK, NUM_BLOCKS),
68    );
69
70    cx.attach_namespaced("aes128", aes.defs()?, &Aes128Chiplet::EXTERNAL_BUS_IDS)?;
71
72    for block in 0..NUM_BLOCKS {
73        for j in 0..16 {
74            cx.publish(cpu.at(CpuColumns::PLAINTEXT + j), in_row(block));
75        }
76
77        for j in 0..16 {
78            cx.publish(cpu.at(CpuColumns::DATA + j), out_row(block));
79        }
80    }
81
82    add_constraints(cx.cs(), cpu);
83
84    cx.compile()
85}
86
87fn add_constraints(cs: &ConstraintSystem<F>, cpu: ColRange) {
88    let col = |c: usize| cs.col(cpu.at(c).index());
89    let next = |c: usize| cs.next(cpu.at(c).index());
90    let key_selector = col(CpuColumns::KEY_SELECTOR);
91
92    for j in 0..16 {
93        let key = col(CpuColumns::KEY + j);
94        let plaintext = col(CpuColumns::PLAINTEXT + j);
95        let data = col(CpuColumns::DATA + j);
96
97        cs.constrain(next(CpuColumns::KEY + j) + key);
98        cs.assert_zero_when(key_selector, data + plaintext + key);
99    }
100}

1use crate::program::{CpuColumns, NUM_BLOCKS, in_row, out_row};
2use hekate_aes::trace::{Aes128Call, expand_key};
3use hekate_aes::{AesRound128Air, PhysAes128Columns};
4use hekate_core::errors;
5use hekate_core::trace::{ColumnTrace, TraceBuilder, TraceColumn};
6use hekate_math::{Bit, Block8, TowerField};
7use std::error::Error;
8
9pub fn aes128_calls(key: &[u8; 16], plaintexts: &[[u8; 16]]) -> Vec<Aes128Call> {
10    let round_keys = expand_key(key);
11
12    plaintexts
13        .iter()
14        .map(|&plaintext| Aes128Call {
15            key: *key,
16            plaintext,
17            round_keys,
18        })
19        .collect()
20}
21
22pub fn extract_ciphertexts(
23    chiplet_traces: &[ColumnTrace],
24) -> Result<[[u8; 16]; NUM_BLOCKS], Box<dyn Error>> {
25    let round_trace = chiplet_traces
26        .first()
27        .ok_or("the AES chiplet traced no round table")?;
28
29    let mut ciphertexts = [[0; 16]; NUM_BLOCKS];
30
31    for (block, ciphertext) in ciphertexts.iter_mut().enumerate() {
32        let output_row = block * AesRound128Air::BLOCK_ROWS + AesRound128Air::ACTIVE_ROWS;
33
34        for (j, byte) in ciphertext.iter_mut().enumerate() {
35            let cell = round_trace
36                .columns
37                .get(PhysAes128Columns::P_STATE_IN + j)
38                .and_then(TraceColumn::as_b8_slice)
39                .and_then(|column| column.get(output_row))
40                .ok_or("the round table holds no output row for the block")?;
41
42            *byte = cell.to_tower().0;
43        }
44    }
45
46    Ok(ciphertexts)
47}
48
49pub fn build_cpu128_trace(
50    key: &[u8; 16],
51    calls: &[Aes128Call],
52    ciphertexts: &[[u8; 16]],
53    num_rows: usize,
54) -> errors::Result<ColumnTrace> {
55    let num_vars = num_rows.trailing_zeros() as usize;
56    let mut tb = TraceBuilder::new_secret(&CpuColumns::build_layout(), num_vars)?;
57    let key = key.map(Block8);
58
59    for row in 0..num_rows {
60        tb.set_b8_array(CpuColumns::KEY, row, &key)?;
61    }
62
63    for (block, (call, ciphertext)) in calls.iter().zip(ciphertexts).enumerate() {
64        let input_row = in_row(block);
65        let output_row = out_row(block);
66
67        let mut whitened = call.plaintext;
68
69        for (byte, k0) in whitened.iter_mut().zip(&call.round_keys[0]) {
70            *byte ^= k0;
71        }
72
73        tb.set_b8_array(
74            CpuColumns::PLAINTEXT,
75            input_row,
76            &call.plaintext.map(Block8),
77        )?;
78        tb.set_b8_array(CpuColumns::DATA, input_row, &whitened.map(Block8))?;
79        tb.set_bit(CpuColumns::SELECTOR, input_row, Bit::ONE)?;
80        tb.set_bit(CpuColumns::KEY_SELECTOR, input_row, Bit::ONE)?;
81
82        tb.set_b8_array(CpuColumns::DATA, output_row, &ciphertext.map(Block8))?;
83        tb.set_bit(CpuColumns::SELECTOR, output_row, Bit::ONE)?;
84    }
85
86    Ok(tb.build())
87}

1mod program;
2mod trace;
3
4use hekate_core::config::Config;
5use hekate_crypto::DefaultHasher;
6use hekate_crypto::transcript::Transcript;
7use hekate_program::digest::program_id;
8use hekate_program::{ProgramInstance, ProgramWitness};
9use hekate_prover_sys::prove;
10use hekate_sdk::preflight;
11use hekate_verifier::HekateVerifier;
12use program::{CPU_IO_PER_BLOCK, F, NUM_BLOCKS, aes128_chiplet, build_aes128_program, table_rows};
13use rand::{TryRng, rngs::SysRng};
14use std::error::Error;
15use trace::{aes128_calls, build_cpu128_trace, extract_ciphertexts};
16
17const LABEL: &[u8] = b"aes128";
18const FIPS128_KEY: [u8; 16] = 0x2b7e1516_28aed2a6_abf71588_09cf4f3c_u128.to_be_bytes();
19const FIPS128_PLAIN: [u8; 16] = 0x3243f6a8_885a308d_313198a2_e0370734_u128.to_be_bytes();
20
21fn main() -> Result<(), Box<dyn Error>> {
22    let plaintexts = generate_plaintexts()?;
23    let calls = aes128_calls(&FIPS128_KEY, &plaintexts);
24
25    let config = Config::default();
26    let cpu_rows = table_rows(&config, CPU_IO_PER_BLOCK);
27    let aes = aes128_chiplet(&config)?;
28    let program = build_aes128_program(cpu_rows, &aes)?;
29
30    let chiplet_traces = aes.generate_traces(&calls)?;
31    let ciphertexts = extract_ciphertexts(&chiplet_traces)?;
32    let cpu_trace = build_cpu128_trace(&FIPS128_KEY, &calls, &ciphertexts, cpu_rows)?;
33
34    let inputs = plaintexts
35        .iter()
36        .zip(&ciphertexts)
37        .flat_map(|(plaintext, ciphertext)| plaintext.iter().chain(ciphertext))
38        .map(|&byte| F::from(byte));
39    let instance = ProgramInstance::new(cpu_rows, inputs.collect());
40
41    let witness = ProgramWitness::new(cpu_trace).with_chiplets(chiplet_traces);
42
43    let report = preflight(&program, &instance, &witness)?;
44
45    if !report.is_clean() {
46        eprintln!("{report}");
47
48        return Err("the trace fails preflight".into());
49    }
50
51    let mut seed = [0; 32];
52    SysRng.try_fill_bytes(&mut seed)?;
53
54    let proof = prove(LABEL, &program, &instance, &witness, &config, seed, None)?;
55
56    let id = program_id::<F, _>(&program)?;
57    let mut transcript = Transcript::<DefaultHasher>::new(LABEL);
58
59    let valid = HekateVerifier::<F, DefaultHasher>::verify(
60        &id,
61        &program,
62        &instance,
63        &proof,
64        &mut transcript,
65        &config,
66    )?;
67
68    if !valid {
69        return Err("the verifier rejects the AES-128 encryptions".into());
70    }
71
72    let [first, ..] = ciphertexts;
73
74    println!("block 0 ciphertext {}", hex(&first));
75    println!("the verifier accepts {NUM_BLOCKS} AES-128 encryptions");
76
77    Ok(())
78}
79
80fn generate_plaintexts() -> Result<[[u8; 16]; NUM_BLOCKS], Box<dyn Error>> {
81    let mut plaintexts = [FIPS128_PLAIN; NUM_BLOCKS];
82
83    for plaintext in plaintexts.iter_mut().skip(1) {
84        SysRng.try_fill_bytes(plaintext)?;
85    }
86
87    Ok(plaintexts)
88}
89
90fn hex(block: &[u8; 16]) -> String {
91    block.iter().map(|byte| format!("{byte:02x}")).collect()
92}

The printed block 0 ciphertext is the one FIPS 197 Appendix B gives, 3925841d02dc09fbdc118597196a0b32, and the verifier checks it as a public value.

§ 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
AES-128	31,250 blocks	885 ms	6.4 ms	2,967 KiB	777 MiB
AES-256	31,250 blocks	981 ms	6.4 ms	3,198 KiB	963 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, below both AES key sizes, which makes the proof the weaker link for AES-128 and AES-256 alike; the ciphertext is still full AES.

§ License

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