hekate_aes/
lib.rs

1// SPDX-License-Identifier: Apache-2.0
2// This file is part of the hekate-aes project.
3// Copyright (C) 2026 Andrei Kochergin <andrei@oumuamua.dev>
4// Copyright (C) 2026 Oumuamua Labs <info@oumuamua.dev>. All rights reserved.
5//
6// Licensed under the Apache License, Version 2.0 (the "License");
7// you may not use this file except in compliance with the License.
8// You may obtain a copy of the License at
9//
10//     http://www.apache.org/licenses/LICENSE-2.0
11//
12// Unless required by applicable law or agreed to in writing, software
13// distributed under the License is distributed on an "AS IS" BASIS,
14// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15// See the License for the specific language governing permissions and
16// limitations under the License.
17
18//! AES chiplets: shared constants and
19//! level-specific Air implementations.
20//!
21//! Shared:
22//! SBOX, ShiftRows, MixColumns, RotWord.
23//!
24//! Level-specific:
25//! - aes128 (AES-128)
26//! - aes256 (AES-256).
27
28#![cfg_attr(not(feature = "std"), no_std)]
29
30extern crate alloc;
31
32use alloc::vec::Vec;
33use hekate_math::TowerField;
34use hekate_program::constraint::builder::ConstraintSystem;
35
36pub(crate) mod sbox_rom;
37
38pub mod aes128;
39pub mod aes256;
40pub mod trace;
41
42pub use aes128::{
43    Aes128Chiplet, Aes128Columns, AesRound128Air, CpuAes128Columns, CpuAes128Unit,
44    PhysAes128Columns,
45};
46pub use aes256::{
47    Aes256Chiplet, Aes256Columns, AesRound256Air, CpuAes256Columns, CpuAes256Unit,
48    PhysAes256Columns,
49};
50
51/// FIPS 197 Table 4:
52/// AES S-box.
53/// Maps each input byte
54/// to its SubBytes output.
55///
56/// S(x) = AffineTransform(x^{-1}) in GF(2^8)
57/// with irreducible polynomial x^8+x^4+x^3+x+1.
58#[rustfmt::skip]
59pub const SBOX: [u8; 256] = [
60    0x63, 0x7c, 0x77, 0x7b, 0xf2, 0x6b, 0x6f, 0xc5,
61    0x30, 0x01, 0x67, 0x2b, 0xfe, 0xd7, 0xab, 0x76,
62    0xca, 0x82, 0xc9, 0x7d, 0xfa, 0x59, 0x47, 0xf0,
63    0xad, 0xd4, 0xa2, 0xaf, 0x9c, 0xa4, 0x72, 0xc0,
64    0xb7, 0xfd, 0x93, 0x26, 0x36, 0x3f, 0xf7, 0xcc,
65    0x34, 0xa5, 0xe5, 0xf1, 0x71, 0xd8, 0x31, 0x15,
66    0x04, 0xc7, 0x23, 0xc3, 0x18, 0x96, 0x05, 0x9a,
67    0x07, 0x12, 0x80, 0xe2, 0xeb, 0x27, 0xb2, 0x75,
68    0x09, 0x83, 0x2c, 0x1a, 0x1b, 0x6e, 0x5a, 0xa0,
69    0x52, 0x3b, 0xd6, 0xb3, 0x29, 0xe3, 0x2f, 0x84,
70    0x53, 0xd1, 0x00, 0xed, 0x20, 0xfc, 0xb1, 0x5b,
71    0x6a, 0xcb, 0xbe, 0x39, 0x4a, 0x4c, 0x58, 0xcf,
72    0xd0, 0xef, 0xaa, 0xfb, 0x43, 0x4d, 0x33, 0x85,
73    0x45, 0xf9, 0x02, 0x7f, 0x50, 0x3c, 0x9f, 0xa8,
74    0x51, 0xa3, 0x40, 0x8f, 0x92, 0x9d, 0x38, 0xf5,
75    0xbc, 0xb6, 0xda, 0x21, 0x10, 0xff, 0xf3, 0xd2,
76    0xcd, 0x0c, 0x13, 0xec, 0x5f, 0x97, 0x44, 0x17,
77    0xc4, 0xa7, 0x7e, 0x3d, 0x64, 0x5d, 0x19, 0x73,
78    0x60, 0x81, 0x4f, 0xdc, 0x22, 0x2a, 0x90, 0x88,
79    0x46, 0xee, 0xb8, 0x14, 0xde, 0x5e, 0x0b, 0xdb,
80    0xe0, 0x32, 0x3a, 0x0a, 0x49, 0x06, 0x24, 0x5c,
81    0xc2, 0xd3, 0xac, 0x62, 0x91, 0x95, 0xe4, 0x79,
82    0xe7, 0xc8, 0x37, 0x6d, 0x8d, 0xd5, 0x4e, 0xa9,
83    0x6c, 0x56, 0xf4, 0xea, 0x65, 0x7a, 0xae, 0x08,
84    0xba, 0x78, 0x25, 0x2e, 0x1c, 0xa6, 0xb4, 0xc6,
85    0xe8, 0xdd, 0x74, 0x1f, 0x4b, 0xbd, 0x8b, 0x8a,
86    0x70, 0x3e, 0xb5, 0x66, 0x48, 0x03, 0xf6, 0x0e,
87    0x61, 0x35, 0x57, 0xb9, 0x86, 0xc1, 0x1d, 0x9e,
88    0xe1, 0xf8, 0x98, 0x11, 0x69, 0xd9, 0x8e, 0x94,
89    0x9b, 0x1e, 0x87, 0xe9, 0xce, 0x55, 0x28, 0xdf,
90    0x8c, 0xa1, 0x89, 0x0d, 0xbf, 0xe6, 0x42, 0x68,
91    0x41, 0x99, 0x2d, 0x0f, 0xb0, 0x54, 0xbb, 0x16,
92];
93
94/// FIPS 197 §5.1.2:
95/// ShiftRows byte permutation.
96/// `SHIFT_MAP[j]` = source byte
97/// index for output position j.
98/// AES state is column-major 4×4:
99/// byte[i] = state[i%4][i/4].
100#[rustfmt::skip]
101const SHIFT_MAP: [usize; 16] = [
102     0,  5, 10, 15,
103     4,  9, 14,  3,
104     8, 13,  2,  7,
105    12,  1,  6, 11,
106];
107
108/// FIPS 197 §5.1.3:
109/// MixColumns coefficient matrix.
110/// MC[row][col] in GF(2^8).
111#[rustfmt::skip]
112const MC: [[u8; 4]; 4] = [
113    [2, 3, 1, 1],
114    [1, 2, 3, 1],
115    [1, 1, 2, 3],
116    [3, 1, 1, 2],
117];
118
119/// FIPS 197 §5.2:
120/// RotWord byte permutation. Maps S-box
121/// index j (0..4) to the source byte
122/// offset in the key's last word.
123const ROT_MAP: [usize; 4] = [13, 14, 15, 12];
124
125const AES_BYTE_LABELS: [&[u8]; 16] = [
126    b"aes_byte_0",
127    b"aes_byte_1",
128    b"aes_byte_2",
129    b"aes_byte_3",
130    b"aes_byte_4",
131    b"aes_byte_5",
132    b"aes_byte_6",
133    b"aes_byte_7",
134    b"aes_byte_8",
135    b"aes_byte_9",
136    b"aes_byte_10",
137    b"aes_byte_11",
138    b"aes_byte_12",
139    b"aes_byte_13",
140    b"aes_byte_14",
141    b"aes_byte_15",
142];
143
144#[rustfmt::skip]
145const SBOX_IN_LABELS: [&[u8]; 16] = [
146    b"aes_sbox_in_0",  b"aes_sbox_in_1",
147    b"aes_sbox_in_2",  b"aes_sbox_in_3",
148    b"aes_sbox_in_4",  b"aes_sbox_in_5",
149    b"aes_sbox_in_6",  b"aes_sbox_in_7",
150    b"aes_sbox_in_8",  b"aes_sbox_in_9",
151    b"aes_sbox_in_10", b"aes_sbox_in_11",
152    b"aes_sbox_in_12", b"aes_sbox_in_13",
153    b"aes_sbox_in_14", b"aes_sbox_in_15",
154];
155
156#[rustfmt::skip]
157const SBOX_OUT_LABELS: [&[u8]; 16] = [
158    b"aes_sbox_out_0",  b"aes_sbox_out_1",
159    b"aes_sbox_out_2",  b"aes_sbox_out_3",
160    b"aes_sbox_out_4",  b"aes_sbox_out_5",
161    b"aes_sbox_out_6",  b"aes_sbox_out_7",
162    b"aes_sbox_out_8",  b"aes_sbox_out_9",
163    b"aes_sbox_out_10", b"aes_sbox_out_11",
164    b"aes_sbox_out_12", b"aes_sbox_out_13",
165    b"aes_sbox_out_14", b"aes_sbox_out_15",
166];
167
168/// FIPS 197 §5.1.2–5.1.4:
169/// SubBytes + ShiftRows + MixColumns + AddRoundKey (full rounds),
170/// SubBytes + ShiftRows + AddRoundKey (final round).
171/// Shared across AES-128 and AES-256, the round
172/// function is identical for all key sizes.
173pub(crate) fn build_round_constraints<F: TowerField>(
174    cs: &ConstraintSystem<F>,
175    state_in: usize,
176    sbox_out: usize,
177    round_key: usize,
178    s_round_col: usize,
179    s_final_col: usize,
180) {
181    let s_round = cs.col(s_round_col);
182    let s_final = cs.col(s_final_col);
183    let two = cs.constant(F::from(2u8));
184    let three = cs.constant(F::from(3u8));
185
186    // Full rounds:
187    // next.state = MixCol(ShiftRows(sbox_out)) + round_key
188    for j in 0..16usize {
189        let aes_col = j / 4;
190        let aes_row = j % 4;
191
192        let mut mc_terms = Vec::with_capacity(4);
193        for k in 0..4 {
194            let src = cs.col(sbox_out + SHIFT_MAP[aes_col * 4 + k]);
195            mc_terms.push(match MC[aes_row][k] {
196                1 => src,
197                2 => two * src,
198                3 => three * src,
199                _ => unreachable!(),
200            });
201        }
202
203        let body = cs.next(state_in + j) + cs.col(round_key + j) + cs.sum(&mc_terms);
204        cs.assert_zero_when(s_round, body);
205    }
206
207    // Final round:
208    // next.state = ShiftRows(sbox_out) + round_key (no MixColumns)
209    for (j, &src_byte) in SHIFT_MAP.iter().enumerate() {
210        let shifted = cs.col(sbox_out + src_byte);
211        let body = cs.next(state_in + j) + cs.col(round_key + j) + shifted;
212
213        cs.assert_zero_when(s_final, body);
214    }
215}
216
217/// FIPS 197 S-box inversion witness:
218/// SubWord(input) = sub via explicit
219/// inverse bit decomposition.
220/// 52 constraints per call (13 per byte).
221pub(crate) fn build_sbox_inversion_constraints<F: TowerField>(
222    cs: &ConstraintSystem<F>,
223    input_cols: [usize; 4],
224    sub_col: usize,
225    inv_bits_col: usize,
226    z_col: usize,
227    gate_col: usize,
228) {
229    let gate = cs.col(gate_col);
230    let one = cs.one();
231    let affine_const = cs.constant(F::from(0x63u8));
232
233    for (j, &in_col) in input_cols.iter().enumerate() {
234        let input = cs.col(in_col);
235        let sub = cs.col(sub_col + j);
236        let z = cs.col(z_col + j);
237
238        cs.assert_boolean(z);
239
240        let bits: [_; 8] = core::array::from_fn(|k| {
241            let b = cs.col(inv_bits_col + j * 8 + k);
242            cs.assert_boolean(b);
243
244            b
245        });
246
247        let inv_terms: Vec<_> = (0..8)
248            .map(|k| cs.scale(F::from(1u8 << k), bits[k]))
249            .collect();
250        let inv_sum = cs.sum(&inv_terms);
251
252        cs.assert_zero_when(gate, input * inv_sum + z + one);
253
254        cs.constrain(z * input);
255        cs.constrain(z * inv_sum);
256
257        let affine_terms: Vec<_> = (0..8)
258            .map(|k| cs.scale(F::from(sbox_rom::AFFINE_COLS[k]), bits[k]))
259            .collect();
260        let affine_sum = cs.sum(&affine_terms);
261
262        cs.assert_zero_when(gate, sub + affine_const + affine_sum);
263    }
264}
265
266#[cfg(test)]
267mod tests {
268    use super::*;
269
270    #[test]
271    fn shift_map_is_permutation() {
272        let mut seen = [false; 16];
273        for &s in &SHIFT_MAP {
274            assert!(!seen[s]);
275            seen[s] = true;
276        }
277    }
278
279    #[test]
280    fn shift_map_row0_identity() {
281        // FIPS 197:
282        // row 0 is not shifted.
283        assert_eq!(SHIFT_MAP[0], 0);
284        assert_eq!(SHIFT_MAP[4], 4);
285        assert_eq!(SHIFT_MAP[8], 8);
286        assert_eq!(SHIFT_MAP[12], 12);
287    }
288}
hekate_aes/lib.rs

hekate_aes/
lib.rs
