# AES-128 and AES-256 Encryption

> Prove AES-128 and AES-256 encryption in zero knowledge, key expansion included, with the key kept private and constant-time trace generation by default.

- URL: https://oumuamua.dev/primitives/encryption/aes
- Published: 2026-10-03T10:49:16+07:00
- Updated: 2026-10-11T12:51:24+07:00

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][usage] publishes nothing, and the [ML-KEM
sender](https://oumuamua.dev/primitives/encryption/mlkem#examples) 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`][example] 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](https://oumuamua.dev/hekate/docs/getting-started/installation) and add `hekate-aes` and `rand`.

```rust src/program.rs:1
use hekate_aes::{Aes128Chiplet, AesRound128Air, host_key_selector_shape, host_selector_shape};
use hekate_core::config::Config;
use hekate_core::errors;
use hekate_math::Block128;
use hekate_program::circuit::{Circuit, CircuitProgram, Col, ColRange};
use hekate_program::constraint::builder::ConstraintSystem;
use hekate_program::define_columns;

pub const NUM_BLOCKS: usize = 4;
pub const CPU_IO_PER_BLOCK: usize = 2;

pub type F = Block128;

define_columns! {
    pub CpuColumns {
        KEY: [B8; 16],
        KEY_SELECTOR: Bit,
        DATA: [B8; 16],
        SELECTOR: Bit,
        PLAINTEXT: [B8; 16],
    }
}

pub const fn in_row(block: usize) -> usize {
    block * CPU_IO_PER_BLOCK
}

pub const fn out_row(block: usize) -> usize {
    in_row(block) + 1
}

pub fn table_rows(config: &Config, rows_per_block: usize) -> usize {
    (NUM_BLOCKS * rows_per_block)
        .next_power_of_two()
        .max(config.min_table_rows())
}

pub fn aes128_chiplet(config: &Config) -> errors::Result<Aes128Chiplet> {
    Aes128Chiplet::new(
        table_rows(config, AesRound128Air::BLOCK_ROWS),
        table_rows(config, AesRound128Air::ACTIVE_ROWS),
        NUM_BLOCKS,
    )
}

pub fn build_aes128_program(
    cpu_rows: usize,
    aes: &Aes128Chiplet,
) -> errors::Result<CircuitProgram<F>> {
    let mut cx = Circuit::<F>::new("Aes128", cpu_rows)?;
    let cpu = cx.schema(&CpuColumns::build_layout());

    let selector = cpu.at(CpuColumns::SELECTOR);
    let key_selector = cpu.at(CpuColumns::KEY_SELECTOR);

    let link_values: Vec<Col> = (0..16).map(|j| cpu.at(CpuColumns::DATA + j)).collect();

    cx.call(&AesRound128Air::link_service(), &link_values, selector)?;

    let key_values: Vec<Col> = (0..16).map(|j| cpu.at(CpuColumns::KEY + j)).collect();

    cx.call(&AesRound128Air::key_service(), &key_values, key_selector)?;

    cx.fix(selector, host_selector_shape(CPU_IO_PER_BLOCK, NUM_BLOCKS));
    cx.fix(
        key_selector,
        host_key_selector_shape(CPU_IO_PER_BLOCK, NUM_BLOCKS),
    );

    cx.attach_namespaced("aes128", aes.defs()?, &Aes128Chiplet::EXTERNAL_BUS_IDS)?;

    for block in 0..NUM_BLOCKS {
        for j in 0..16 {
            cx.publish(cpu.at(CpuColumns::PLAINTEXT + j), in_row(block));
        }

        for j in 0..16 {
            cx.publish(cpu.at(CpuColumns::DATA + j), out_row(block));
        }
    }

    add_constraints(cx.cs(), cpu);

    cx.compile()
}

fn add_constraints(cs: &ConstraintSystem<F>, cpu: ColRange) {
    let col = |c: usize| cs.col(cpu.at(c).index());
    let next = |c: usize| cs.next(cpu.at(c).index());
    let key_selector = col(CpuColumns::KEY_SELECTOR);

    for j in 0..16 {
        let key = col(CpuColumns::KEY + j);
        let plaintext = col(CpuColumns::PLAINTEXT + j);
        let data = col(CpuColumns::DATA + j);

        cs.constrain(next(CpuColumns::KEY + j) + key);
        cs.assert_zero_when(key_selector, data + plaintext + key);
    }
}
```

```rust src/trace.rs:1
use crate::program::{CpuColumns, NUM_BLOCKS, in_row, out_row};
use hekate_aes::trace::{Aes128Call, expand_key};
use hekate_aes::{AesRound128Air, PhysAes128Columns};
use hekate_core::errors;
use hekate_core::trace::{ColumnTrace, TraceBuilder, TraceColumn};
use hekate_math::{Bit, Block8, TowerField};
use std::error::Error;

pub fn aes128_calls(key: &[u8; 16], plaintexts: &[[u8; 16]]) -> Vec<Aes128Call> {
    let round_keys = expand_key(key);

    plaintexts
        .iter()
        .map(|&plaintext| Aes128Call {
            key: *key,
            plaintext,
            round_keys,
        })
        .collect()
}

pub fn extract_ciphertexts(
    chiplet_traces: &[ColumnTrace],
) -> Result<[[u8; 16]; NUM_BLOCKS], Box<dyn Error>> {
    let round_trace = chiplet_traces
        .first()
        .ok_or("the AES chiplet traced no round table")?;

    let mut ciphertexts = [[0; 16]; NUM_BLOCKS];

    for (block, ciphertext) in ciphertexts.iter_mut().enumerate() {
        let output_row = block * AesRound128Air::BLOCK_ROWS + AesRound128Air::ACTIVE_ROWS;

        for (j, byte) in ciphertext.iter_mut().enumerate() {
            let cell = round_trace
                .columns
                .get(PhysAes128Columns::P_STATE_IN + j)
                .and_then(TraceColumn::as_b8_slice)
                .and_then(|column| column.get(output_row))
                .ok_or("the round table holds no output row for the block")?;

            *byte = cell.to_tower().0;
        }
    }

    Ok(ciphertexts)
}

pub fn build_cpu128_trace(
    key: &[u8; 16],
    calls: &[Aes128Call],
    ciphertexts: &[[u8; 16]],
    num_rows: usize,
) -> errors::Result<ColumnTrace> {
    let num_vars = num_rows.trailing_zeros() as usize;
    let mut tb = TraceBuilder::new_secret(&CpuColumns::build_layout(), num_vars)?;
    let key = key.map(Block8);

    for row in 0..num_rows {
        tb.set_b8_array(CpuColumns::KEY, row, &key)?;
    }

    for (block, (call, ciphertext)) in calls.iter().zip(ciphertexts).enumerate() {
        let input_row = in_row(block);
        let output_row = out_row(block);

        let mut whitened = call.plaintext;

        for (byte, k0) in whitened.iter_mut().zip(&call.round_keys[0]) {
            *byte ^= k0;
        }

        tb.set_b8_array(
            CpuColumns::PLAINTEXT,
            input_row,
            &call.plaintext.map(Block8),
        )?;
        tb.set_b8_array(CpuColumns::DATA, input_row, &whitened.map(Block8))?;
        tb.set_bit(CpuColumns::SELECTOR, input_row, Bit::ONE)?;
        tb.set_bit(CpuColumns::KEY_SELECTOR, input_row, Bit::ONE)?;

        tb.set_b8_array(CpuColumns::DATA, output_row, &ciphertext.map(Block8))?;
        tb.set_bit(CpuColumns::SELECTOR, output_row, Bit::ONE)?;
    }

    Ok(tb.build())
}
```

```rust src/main.rs:1
mod program;
mod trace;

use hekate_core::config::Config;
use hekate_crypto::DefaultHasher;
use hekate_crypto::transcript::Transcript;
use hekate_program::digest::program_id;
use hekate_program::{ProgramInstance, ProgramWitness};
use hekate_prover_sys::prove;
use hekate_sdk::preflight;
use hekate_verifier::HekateVerifier;
use program::{CPU_IO_PER_BLOCK, F, NUM_BLOCKS, aes128_chiplet, build_aes128_program, table_rows};
use rand::{TryRng, rngs::SysRng};
use std::error::Error;
use trace::{aes128_calls, build_cpu128_trace, extract_ciphertexts};

const LABEL: &[u8] = b"aes128";
const FIPS128_KEY: [u8; 16] = 0x2b7e1516_28aed2a6_abf71588_09cf4f3c_u128.to_be_bytes();
const FIPS128_PLAIN: [u8; 16] = 0x3243f6a8_885a308d_313198a2_e0370734_u128.to_be_bytes();

fn main() -> Result<(), Box<dyn Error>> {
    let plaintexts = generate_plaintexts()?;
    let calls = aes128_calls(&FIPS128_KEY, &plaintexts);

    let config = Config::default();
    let cpu_rows = table_rows(&config, CPU_IO_PER_BLOCK);
    let aes = aes128_chiplet(&config)?;
    let program = build_aes128_program(cpu_rows, &aes)?;

    let chiplet_traces = aes.generate_traces(&calls)?;
    let ciphertexts = extract_ciphertexts(&chiplet_traces)?;
    let cpu_trace = build_cpu128_trace(&FIPS128_KEY, &calls, &ciphertexts, cpu_rows)?;

    let inputs = plaintexts
        .iter()
        .zip(&ciphertexts)
        .flat_map(|(plaintext, ciphertext)| plaintext.iter().chain(ciphertext))
        .map(|&byte| F::from(byte));
    let instance = ProgramInstance::new(cpu_rows, inputs.collect());

    let witness = ProgramWitness::new(cpu_trace).with_chiplets(chiplet_traces);

    let report = preflight(&program, &instance, &witness)?;

    if !report.is_clean() {
        eprintln!("{report}");

        return Err("the trace fails preflight".into());
    }

    let mut seed = [0; 32];
    SysRng.try_fill_bytes(&mut seed)?;

    let proof = prove(LABEL, &program, &instance, &witness, &config, seed, None)?;

    let id = program_id::<F, _>(&program)?;
    let mut transcript = Transcript::<DefaultHasher>::new(LABEL);

    let valid = HekateVerifier::<F, DefaultHasher>::verify(
        &id,
        &program,
        &instance,
        &proof,
        &mut transcript,
        &config,
    )?;

    if !valid {
        return Err("the verifier rejects the AES-128 encryptions".into());
    }

    let [first, ..] = ciphertexts;

    println!("block 0 ciphertext {}", hex(&first));
    println!("the verifier accepts {NUM_BLOCKS} AES-128 encryptions");

    Ok(())
}

fn generate_plaintexts() -> Result<[[u8; 16]; NUM_BLOCKS], Box<dyn Error>> {
    let mut plaintexts = [FIPS128_PLAIN; NUM_BLOCKS];

    for plaintext in plaintexts.iter_mut().skip(1) {
        SysRng.try_fill_bytes(plaintext)?;
    }

    Ok(plaintexts)
}

fn hex(block: &[u8; 16]) -> String {
    block.iter().map(|byte| format!("{byte:02x}")).collect()
}
```

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](https://oumuamua.dev/hekate/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](mailto:info@oumuamua.dev). The [README's Usage section][usage] builds an AES-256 host program, and
the API reference is on docs.rs.

[example]: https://github.com/oumuamua-labs/hekate/blob/main/hekate/examples/aes.rs
[usage]: https://github.com/oumuamua-labs/hekate/blob/main/hekate-aes/README.md#usage