1//! For LICENSE check out https://github.com/0xPARC/plonky2-aes/blob/main/LICENSE
  2//!
  3//! Rust native implementation of [AES-GCM (Galois Counter
  4//! Mode)](https://nvlpubs.nist.gov/nistpubs/Legacy/SP/nistspecialpublication800-38d.pdf)
  5//! having in mind the approach that will be done in-circuit; this means that it
  6//! is not written in the most rust idiomatic way, nor memory-efficient, but
  7//! simulating the behavior that we will do later inside the circuit.
  8//!
  9
 10use crate::native_aes::{encrypt_block, flatten_state, key_expansion, State};
 11
 12// supported tag length
 13const TAG_LEN: usize = 128;
 14
 15/// Section 7.1 from https://nvlpubs.nist.gov/nistpubs/Legacy/SP/nistspecialpublication800-38d.pdf
 16pub fn encrypt<const NK: usize, const NB: usize, const NR: usize>(
 17    key: &[u8; NK * NB],
 18    nonce: &[u8; 12],
 19    pt: &[u8], // plaintext
 20) -> (Vec<u8>, Vec<u8>)
 21where
 22    [(); 4 * (NR + 1)]:,
 23{
 24    let a: &[u8] = &[]; // additional authenticated data // TODO maybe as input
 25    let expanded_key: [[u8; 4]; 4 * (NR + 1)] = key_expansion::<NK, NB, NR>(&key);
 26
 27    // 1. CIPH_K
 28    let h = encrypt_block::<NR>(&[0u8; 16], &expanded_key);
 29    let h = flatten_state(h);
 30
 31    // 2. J_0
 32    let j0: [u8; 16] = if nonce.len() * 8 == 96 {
 33        // J_0 = IV || 0^31 || 1
 34        let mut out = [0u8; 16];
 35        out[..12].copy_from_slice(&nonce.as_slice());
 36        out[12..16].copy_from_slice(&[0x00, 0x00, 0x00, 0x01]);
 37        out
 38    } else {
 39        panic!("unsuported at initial version; nonce.len()=12 (96 bits)");
 40    };
 41
 42    // 3. C=GCTR()
 43    let c = gctr(expanded_key, &inc32(j0), &pt);
 44
 45    // 4. u, v
 46    // let u = (16 - ceil(c.len() % 16)) % 16 // TODO WIP
 47    let u: usize = 128 * ((c.len() * 8) as f64 / 128_f64).ceil() as usize - c.len() * 8;
 48    let v: usize = 128 * ((a.len() * 8) as f64 / 128_f64).ceil() as usize - a.len() * 8;
 49
 50    // 5. S = GHASH()
 51    let a_len: [u8; 8] = (a.len() * 8).to_be_bytes();
 52    let c_len: [u8; 8] = (c.len() * 8).to_be_bytes();
 53    let ghash_input = &[
 54        a.to_vec(),
 55        vec![0u8; v],
 56        c.clone(),
 57        vec![0u8; u],
 58        a_len.to_vec(),
 59        c_len.to_vec(),
 60    ]
 61    .concat();
 62    let s = ghash(h, ghash_input);
 63
 64    // 6. T=MSB(GCTR()))
 65    let t = msb_t(TAG_LEN, &gctr(expanded_key, &j0, &s));
 66
 67    (c, t)
 68}
 69
 70/// GCM-Ctr, Section 6.5, Algorithm 3
 71fn gctr<const NR: usize>(key: [[u8; 4]; 4 * (NR + 1)], icb: &[u8; 16], x: &[u8]) -> Vec<u8> {
 72    if x.is_empty() {
 73        return x.to_vec();
 74    }
 75
 76    let n = ((x.len() * 8) as f64 / 128_f64).ceil() as usize;
 77
 78    let mut y = vec![];
 79    let mut cb_i = icb.clone();
 80    for (i, x_i_raw) in x.chunks(16).enumerate() {
 81        // 5.
 82        if i > 1 {
 83            cb_i = inc32(cb_i);
 84        }
 85
 86        let l = x_i_raw.len().min(16);
 87        let mut x_i = [0u8; 16];
 88        x_i[..l].copy_from_slice(x_i_raw);
 89
 90        let ciph_cb_i = flatten_state(encrypt_block::<NR>(&cb_i, &key));
 91
 92        let y_i = if i < n {
 93            xor_blocks(x_i, ciph_cb_i)
 94        } else {
 95            // last chunk, might be smaller than 16 bytes
 96            let msb_res = msb_t(x_i.len(), ciph_cb_i.as_slice());
 97            let l = msb_res.len().min(16);
 98            let mut m = [0u8; 16];
 99            m[..l].copy_from_slice(&msb_res);
100            xor_blocks(x_i, m)
101        };
102        y.push(y_i);
103    }
104    y.concat()
105}
106
107/// Section 6.4, Algorithm 2
108fn ghash(h: [u8; 16], x: &[u8]) -> [u8; 16] {
109    assert!(x.len() % 16 == 0); // multiple of 128 bits
110    let m = x.len() / 16;
111
112    let mut y = [0u8; 16]; // (128 bits)
113    for i in 0..m {
114        let mut xi = [0u8; 16];
115        xi.clone_from_slice(&x[i * 16..i * 16 + 16]);
116        let y_xi = xor_blocks(y, xi);
117        y = gf_2_128_mul(y_xi, h);
118    }
119    y
120}
121fn xor_blocks(b1: [u8; 16], b2: [u8; 16]) -> [u8; 16] {
122    let mut r = [0u8; 16];
123    for i in 0..16 {
124        r[i] = b1[i] ^ b2[i];
125    }
126    r
127}
128
129/// multiplication of blocks (in GF(2^128)), Algorithm 1, Section 6.3
130pub fn gf_2_128_mul(x: [u8; 16], y: [u8; 16]) -> [u8; 16] {
131    // R = 11100001 || 0^120
132    const R: [u8; 16] = [
133        0xE1, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
134        0x00,
135    ];
136
137    let mut z = [0u8; 16];
138    let mut v = y;
139    for i in 0..128 {
140        // xi: i-th bit of x
141        let byte_index = i / 8;
142        let bit_index = 7 - (i % 8);
143        let xi = (x[byte_index] >> bit_index) & 1;
144
145        if xi == 1 {
146            z = xor_blocks(z, v);
147        }
148        let lsb = v[15] & 1;
149        right_shift_one(&mut v);
150        if lsb == 1 {
151            v = xor_blocks(v, R);
152        }
153    }
154    z
155}
156fn right_shift_one(block: &mut [u8; 16]) {
157    let mut carry = 0u8;
158
159    for byte in block.iter_mut() {
160        let new_carry = *byte & 1;
161        *byte = (*byte >> 1) | (carry << 7);
162        carry = new_carry;
163    }
164}
165
166/// increment the right-most 32 bits of the given block (128 bits). Section 2.
167pub fn inc32(b: [u8; 16]) -> [u8; 16] {
168    let mut r = b.clone();
169    // counter = last 32bits
170    let counter = u32::from_be_bytes([r[12], r[13], r[14], r[15]]);
171    let incremented = counter.wrapping_add(1);
172    let bytes = incremented.to_be_bytes();
173    r[12..16].copy_from_slice(&bytes);
174    r
175}
176
177/// returns the t left-most (most-significant in BE) bits of the block
178pub fn msb_t(t: usize, block: &[u8]) -> Vec<u8> {
179    assert!(t <= 128, "MSB_t err: max 128 bits");
180
181    let mut out = Vec::with_capacity((t + 7) / 8); // ceiling bytes
182    let mut bits_left = t;
183    for byte in block.iter() {
184        if bits_left == 0 {
185            break;
186        }
187        if bits_left >= 8 {
188            out.push(*byte);
189            bits_left -= 8;
190        } else {
191            // take the top `bits_left` bits of the byte
192            let mask = 0xFFu8 << (8 - bits_left);
193            out.push(byte & mask);
194            break;
195        }
196    }
197    out
198}
199
200#[cfg(test)]
201mod tests {
202    use super::*;
203
204    use aes_gcm::{
205        aead::{Aead, AeadCore, Key, KeyInit, OsRng},
206        Aes128Gcm, Nonce,
207    };
208
209    #[test]
210    fn test_aes_gcm_encrypt_nist_vector() {
211        // https://csrc.nist.gov/Projects/Cryptographic-Algorithm-Validation-Program/CAVP-TESTING-BLOCK-CIPHER-MODES
212
213        // test vector at line 7
214        let key: [u8; 16] = [
215            0xcf, 0x06, 0x3a, 0x34, 0xd4, 0xa9, 0xa7, 0x6c, 0x2c, 0x86, 0x78, 0x7d, 0x3f, 0x96,
216            0xdb, 0x71,
217        ];
218        let nonce: [u8; 12] = [
219            0x11, 0x3b, 0x97, 0x85, 0x97, 0x18, 0x64, 0xc8, 0x3b, 0x01, 0xc7, 0x87,
220        ];
221        let plaintext = [];
222        let _aad: &[u8] = &[];
223        let expected_ciphertext: &[u8] = &[];
224        let expected_tag: &[u8] = &[
225            0x72, 0xac, 0x84, 0x93, 0xe3, 0xa5, 0x22, 0x8b, 0x5d, 0x13, 0x0a, 0x69, 0xd2, 0x51,
226            0x0e, 0x42,
227        ];
228        let (c, tag) = encrypt::<4, 4, 10>(&key, &nonce, &plaintext);
229        assert_eq!(c, expected_ciphertext);
230        assert_eq!(tag, expected_tag);
231
232        // test vector at line 4417
233        let key: [u8; 16] = [
234            0xe9, 0x8b, 0x72, 0xa9, 0x88, 0x1a, 0x84, 0xca, 0x6b, 0x76, 0xe0, 0xf4, 0x3e, 0x68,
235            0x64, 0x7a,
236        ];
237        let nonce: [u8; 12] = [
238            0x8b, 0x23, 0x29, 0x9f, 0xde, 0x17, 0x40, 0x53, 0xf3, 0xd6, 0x52, 0xba,
239        ];
240        let plaintext = [
241            0x28, 0x28, 0x6a, 0x32, 0x12, 0x93, 0x25, 0x3c, 0x3e, 0x0a, 0xa2, 0x70, 0x4a, 0x27,
242            0x80, 0x32,
243        ];
244        let _aad: &[u8] = &[];
245        let expected_ciphertext: &[u8] = &[
246            0x5a, 0x3c, 0x1c, 0xf1, 0x98, 0x5d, 0xbb, 0x8b, 0xed, 0x81, 0x80, 0x36, 0xfd, 0xd5,
247            0xab, 0x42,
248        ];
249        let expected_tag: &[u8] = &[
250            0x23, 0xc7, 0xab, 0x0f, 0x95, 0x2b, 0x70, 0x91, 0xcd, 0x32, 0x48, 0x35, 0x04, 0x3b,
251            0x5e, 0xb5,
252        ];
253        let (c, tag) = encrypt::<4, 4, 10>(&key, &nonce, &plaintext);
254        assert_eq!(c, expected_ciphertext);
255        assert_eq!(tag, expected_tag);
256    }
257
258    /// test checking against external aes-gcm lib
259    #[test]
260    fn test_with_external_lib() -> anyhow::Result<()> {
261        let key: &[u8; 16] = &[42; 16];
262        let nonce: &[u8; 12] = &[111; 12];
263        let pt: &[u8] = &[42u8; 16];
264
265        let key_ext: &Key<Aes128Gcm> = key.into();
266        let nonce_ext = Nonce::from_slice(nonce);
267        let cipher = Aes128Gcm::new(&key_ext);
268
269        let ciphertext = cipher.encrypt(&nonce_ext, pt.as_ref()).unwrap();
270
271        let plaintext = cipher.decrypt(&nonce_ext, ciphertext.as_ref()).unwrap();
272        assert_eq!(plaintext, pt);
273
274        let (c, t) = encrypt::<4, 4, 10>(key, nonce, &pt);
275        assert_eq!([c, t].concat(), ciphertext);
276        Ok(())
277    }
278}