hekate_aes/
lib.rs

1// SPDX-FileCopyrightText: 2026 Andrei Kochergin <andrei@oumuamua.dev>
2// SPDX-FileCopyrightText: 2026 Oumuamua Labs <info@oumuamua.dev>
3// SPDX-License-Identifier: AGPL-3.0-only
4
5//! AES-128 and AES-256 chiplets for Hekate, the Rust zero-knowledge proof engine.
6//!
7//! [`Aes128Chiplet`] and [`Aes256Chiplet`] pair a round table,
8//! [`AesRound128Air`] or [`AesRound256Air`], that proves the FIPS 197 round
9//! function with an S-box ROM for the GF(2^8) inversion. [`trace`] expands
10//! keys and builds the chiplet traces in constant time.
11//!
12//! - [AES-128 and AES-256 Encryption][aes]: what the
13//!   proof states and what stays outside it
14//!
15//! [aes]: https://oumuamua.dev/primitives/encryption/aes
16
17#![cfg_attr(not(feature = "std"), no_std)]
18
19extern crate alloc;
20
21use alloc::vec::Vec;
22use hekate_math::TowerField;
23use hekate_program::FixedShape;
24use hekate_program::constraint::builder::ConstraintSystem;
25
26pub(crate) mod sbox_rom;
27
28pub mod aes128;
29pub mod aes256;
30pub mod trace;
31
32pub use aes128::{
33    Aes128Chiplet, Aes128Columns, AesRound128Air, CpuAes128Columns, PhysAes128Columns,
34};
35pub use aes256::{
36    Aes256Chiplet, Aes256Columns, AesRound256Air, CpuAes256Columns, PhysAes256Columns,
37};
38pub use sbox_rom::{PhysSboxRomColumns, SboxRomColumns};
39
40/// Host emit schedule:
41/// `SELECTOR` fires on each block's input row
42/// at offset 0 and output row at offset 1.
43pub fn host_selector_shape<F: TowerField>(stride: usize, count: usize) -> FixedShape<F> {
44    assert!(stride >= 2);
45
46    FixedShape::Cadence {
47        stride,
48        count,
49        origin: 0,
50        values: (0..stride)
51            .map(|off| if off <= 1 { F::ONE } else { F::ZERO })
52            .collect(),
53    }
54}
55
56/// Host key emit schedule:
57/// `KEY_SELECTOR` fires only on each
58/// block's input row at offset 0.
59pub fn host_key_selector_shape<F: TowerField>(stride: usize, count: usize) -> FixedShape<F> {
60    assert!(stride >= 2);
61
62    FixedShape::Cadence {
63        stride,
64        count,
65        origin: 0,
66        values: (0..stride)
67            .map(|off| if off == 0 { F::ONE } else { F::ZERO })
68            .collect(),
69    }
70}
71
72/// FIPS 197 §5.1.2:
73/// ShiftRows byte permutation. `SHIFT_MAP[j]` = source byte
74/// index for output position j. AES state is column-major 4×4:
75/// byte[i] = state[i%4][i/4].
76#[rustfmt::skip]
77const SHIFT_MAP: [usize; 16] = [
78     0,  5, 10, 15,
79     4,  9, 14,  3,
80     8, 13,  2,  7,
81    12,  1,  6, 11,
82];
83
84/// FIPS 197 §5.1.3:
85/// MixColumns coefficient matrix.
86/// MC[row][col] in GF(2^8).
87#[rustfmt::skip]
88const MC: [[u8; 4]; 4] = [
89    [2, 3, 1, 1],
90    [1, 2, 3, 1],
91    [1, 1, 2, 3],
92    [3, 1, 1, 2],
93];
94
95/// FIPS 197 §5.2:
96/// RotWord byte permutation. Maps S-box
97/// index j (0..4) to the source byte
98/// offset in the key's last word.
99const ROT_MAP: [usize; 4] = [13, 14, 15, 12];
100
101pub const AES_BYTE_LABELS: [&[u8]; 16] = [
102    b"aes_byte_0",
103    b"aes_byte_1",
104    b"aes_byte_2",
105    b"aes_byte_3",
106    b"aes_byte_4",
107    b"aes_byte_5",
108    b"aes_byte_6",
109    b"aes_byte_7",
110    b"aes_byte_8",
111    b"aes_byte_9",
112    b"aes_byte_10",
113    b"aes_byte_11",
114    b"aes_byte_12",
115    b"aes_byte_13",
116    b"aes_byte_14",
117    b"aes_byte_15",
118];
119
120#[rustfmt::skip]
121const SBOX_IN_LABELS: [&[u8]; 16] = [
122    b"aes_sbox_in_0",  b"aes_sbox_in_1",
123    b"aes_sbox_in_2",  b"aes_sbox_in_3",
124    b"aes_sbox_in_4",  b"aes_sbox_in_5",
125    b"aes_sbox_in_6",  b"aes_sbox_in_7",
126    b"aes_sbox_in_8",  b"aes_sbox_in_9",
127    b"aes_sbox_in_10", b"aes_sbox_in_11",
128    b"aes_sbox_in_12", b"aes_sbox_in_13",
129    b"aes_sbox_in_14", b"aes_sbox_in_15",
130];
131
132#[rustfmt::skip]
133const SBOX_OUT_LABELS: [&[u8]; 16] = [
134    b"aes_sbox_out_0",  b"aes_sbox_out_1",
135    b"aes_sbox_out_2",  b"aes_sbox_out_3",
136    b"aes_sbox_out_4",  b"aes_sbox_out_5",
137    b"aes_sbox_out_6",  b"aes_sbox_out_7",
138    b"aes_sbox_out_8",  b"aes_sbox_out_9",
139    b"aes_sbox_out_10", b"aes_sbox_out_11",
140    b"aes_sbox_out_12", b"aes_sbox_out_13",
141    b"aes_sbox_out_14", b"aes_sbox_out_15",
142];
143
144/// FIPS 197 §5.1.2–5.1.4:
145/// SubBytes + ShiftRows + MixColumns + AddRoundKey (full rounds),
146/// SubBytes + ShiftRows + AddRoundKey (final round).
147/// Shared across AES-128 and AES-256, the round
148/// function is identical for all key sizes.
149pub(crate) fn build_round_constraints<F: TowerField>(
150    cs: &ConstraintSystem<F>,
151    state_in: usize,
152    sbox_out: usize,
153    round_key: usize,
154    s_round_col: usize,
155    s_final_col: usize,
156) {
157    let s_round = cs.col(s_round_col);
158    let s_final = cs.col(s_final_col);
159    let two = cs.constant(F::from(2u8));
160    let three = cs.constant(F::from(3u8));
161
162    // Full rounds:
163    // next.state = MixCol(ShiftRows(sbox_out)) + round_key
164    for j in 0..16usize {
165        let aes_col = j / 4;
166        let aes_row = j % 4;
167
168        let mut mc_terms = Vec::with_capacity(4);
169        for k in 0..4 {
170            let src = cs.col(sbox_out + SHIFT_MAP[aes_col * 4 + k]);
171            mc_terms.push(match MC[aes_row][k] {
172                1 => src,
173                2 => two * src,
174                3 => three * src,
175                _ => unreachable!(),
176            });
177        }
178
179        let body = cs.next(state_in + j) + cs.col(round_key + j) + cs.sum(&mc_terms);
180        cs.assert_zero_when(s_round, body);
181    }
182
183    // Final round:
184    // next.state = ShiftRows(sbox_out) + round_key (no MixColumns)
185    for (j, &src_byte) in SHIFT_MAP.iter().enumerate() {
186        let shifted = cs.col(sbox_out + src_byte);
187        let body = cs.next(state_in + j) + cs.col(round_key + j) + shifted;
188
189        cs.assert_zero_when(s_final, body);
190    }
191}
192
193/// FIPS 197 S-box inversion witness:
194/// SubWord(input) = sub via explicit
195/// inverse bit decomposition.
196/// 52 constraints per call (13 per byte).
197pub(crate) fn build_sbox_inversion_constraints<F: TowerField>(
198    cs: &ConstraintSystem<F>,
199    input_cols: [usize; 4],
200    sub_col: usize,
201    inv_bits_col: usize,
202    z_col: usize,
203    gate_col: usize,
204) {
205    let gate = cs.col(gate_col);
206    let one = cs.one();
207    let affine_const = cs.constant(F::from(0x63u8));
208
209    for (j, &in_col) in input_cols.iter().enumerate() {
210        let input = cs.col(in_col);
211        let sub = cs.col(sub_col + j);
212        let z = cs.col(z_col + j);
213
214        cs.assert_boolean(z);
215
216        let bits: [_; 8] = core::array::from_fn(|k| {
217            let b = cs.col(inv_bits_col + j * 8 + k);
218            cs.assert_boolean(b);
219
220            b
221        });
222
223        let inv_terms: Vec<_> = (0..8)
224            .map(|k| cs.scale(F::from(1u8 << k), bits[k]))
225            .collect();
226        let inv_sum = cs.sum(&inv_terms);
227
228        cs.assert_zero_when(gate, input * inv_sum + z + one);
229
230        cs.constrain(z * input);
231        cs.constrain(z * inv_sum);
232
233        let affine_terms: Vec<_> = (0..8)
234            .map(|k| cs.scale(F::from(sbox_rom::AFFINE_COLS[k]), bits[k]))
235            .collect();
236        let affine_sum = cs.sum(&affine_terms);
237
238        cs.assert_zero_when(gate, sub + affine_const + affine_sum);
239    }
240}
241
242#[cfg(test)]
243mod tests {
244    use super::*;
245
246    #[test]
247    fn shift_map_is_permutation() {
248        let mut seen = [false; 16];
249        for &s in &SHIFT_MAP {
250            assert!(!seen[s]);
251            seen[s] = true;
252        }
253    }
254
255    #[test]
256    fn shift_map_row0_identity() {
257        // FIPS 197:
258        // row 0 is not shifted.
259        assert_eq!(SHIFT_MAP[0], 0);
260        assert_eq!(SHIFT_MAP[4], 4);
261        assert_eq!(SHIFT_MAP[8], 8);
262        assert_eq!(SHIFT_MAP[12], 12);
263    }
264}
hekate_aes/lib.rs

hekate_aes/
lib.rs
