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(&inc32(j0), &pt);
 44
 45    // 4. u, v
 46    // let u = (16 - ceil(c.len() % 16)) % 16 // TODO WIP
 47    let u: usize = 16 * (c.len() as f64 / 16_f64).ceil() as usize - c.len();
 48    let v: usize = 16 * (a.len() as f64 / 16_f64).ceil() as usize - a.len();
 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(&j0, &s));
 66
 67    (c, t)
 68}
 69
 70/// Section 6.5, Algorithm 3
 71fn gctr(icb: &[u8; 16], x: &[u8]) -> Vec<u8> {
 72    todo!()
 73}
 74
 75/// Section 6.4, Algorithm 2
 76fn ghash(h: [u8; 16], x: &[u8]) -> [u8; 16] {
 77    assert!(x.len() % 16 == 0); // multiple of 128 bits
 78    let m = x.len() / 16;
 79
 80    let mut y = [0u8; 16]; // (128 bits)
 81    for i in 0..m {
 82        let mut xi = [0u8; 16];
 83        xi.clone_from_slice(&x[i * 16..i * 16 + 16]);
 84        let y_xi = xor_blocks(y, xi);
 85        y = gf_2_128_mul(y_xi, h);
 86    }
 87    y
 88}
 89fn xor_blocks(b1: [u8; 16], b2: [u8; 16]) -> [u8; 16] {
 90    let mut r = [0u8; 16];
 91    for i in 0..16 {
 92        r[i] = b1[i] ^ b2[i];
 93    }
 94    r
 95}
 96
 97/// multiplication of blocks (in GF(2^128)), Algorithm 1, Section 6.3
 98pub fn gf_2_128_mul(x: [u8; 16], y: [u8; 16]) -> [u8; 16] {
 99    // R = 11100001 || 0^120
100    const R: [u8; 16] = [
101        0xE1, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
102        0x00,
103    ];
104
105    let mut z = [0u8; 16];
106    let mut v = y;
107    for i in 0..128 {
108        // xi: i-th bit of x
109        let byte_index = i / 8;
110        let bit_index = 7 - (i % 8);
111        let xi = (x[byte_index] >> bit_index) & 1;
112
113        if xi == 1 {
114            z = xor_blocks(z, v);
115        }
116        let lsb = v[15] & 1;
117        right_shift_one(&mut v);
118        if lsb == 1 {
119            v = xor_blocks(v, R);
120        }
121    }
122    z
123}
124fn right_shift_one(block: &mut [u8; 16]) {
125    let mut carry = 0u8;
126
127    for byte in block.iter_mut() {
128        let new_carry = *byte & 1;
129        *byte = (*byte >> 1) | (carry << 7);
130        carry = new_carry;
131    }
132}
133
134/// increment the right-most 32 bits of the given block (128 bits). Section 2.
135pub fn inc32(b: [u8; 16]) -> [u8; 16] {
136    let mut r = b.clone();
137    // counter = last 32bits
138    let counter = u32::from_be_bytes([r[12], r[13], r[14], r[15]]);
139    let incremented = counter.wrapping_add(1);
140    let bytes = incremented.to_be_bytes();
141    r[12..16].copy_from_slice(&bytes);
142    r
143}
144
145/// returns the t left-most (most-significant in BE) bits of the block
146pub fn msb_t(t: usize, block: &[u8]) -> Vec<u8> {
147    assert!(t <= 128, "MSB_t err: max 128 bits");
148
149    let mut out = Vec::with_capacity((t + 7) / 8); // ceiling bytes
150    let mut bits_left = t;
151    for byte in block.iter() {
152        if bits_left == 0 {
153            break;
154        }
155        if bits_left >= 8 {
156            out.push(*byte);
157            bits_left -= 8;
158        } else {
159            // take the top `bits_left` bits of the byte
160            let mask = 0xFFu8 << (8 - bits_left);
161            out.push(byte & mask);
162            break;
163        }
164    }
165    out
166}
167
168#[cfg(test)]
169mod tests {
170    use super::*;
171
172    use aes_gcm::{
173        aead::{Aead, AeadCore, Key, KeyInit, OsRng},
174        Aes128Gcm, Nonce,
175    };
176
177    /// test checking against NIST test vector
178    #[test]
179    fn test_aes_gcm_encrypt_nist_vector() {
180        // https://csrc.nist.gov/CSRC/media/Projects/Cryptographic-Standards-and-Guidelines/documents/examples/AES_GCM.pdf
181        let key: [u8; 16] = [
182            0xFE, 0xFF, 0xE9, 0x92, 0x86, 0x65, 0x73, 0x1C, 0x6D, 0x6A, 0x8F, 0x94, 0x67, 0x30,
183            0x83, 0x08,
184        ];
185        let nonce: [u8; 12] = [
186            0xCA, 0xFE, 0xBA, 0xBE, 0xFA, 0xCE, 0xDB, 0xAD, 0xDE, 0xCA, 0xF8, 0x88,
187        ];
188        let plaintext = [
189            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
190            0x00, 0x00,
191        ];
192        let aad: &[u8] = &[];
193        let expected_ciphertext: &[u8] = &[
194            0x32, 0x47, 0x18, 0x4B, 0x3C, 0x4F, 0x69, 0xA4, 0x4D, 0xBC, 0xD2, 0x28, 0x87, 0xBB,
195            0xB4, 0x18,
196        ];
197
198        todo!();
199    }
200
201    /// test checking against external aes-gcm lib
202    #[test]
203    fn test_with_external_lib() -> anyhow::Result<()> {
204        let key: &[u8; 16] = &[42; 16];
205        let nonce: &[u8; 12] = &[111; 12];
206        let pt = b"plaintext test message";
207
208        let key: &Key<Aes128Gcm> = key.into();
209        let cipher = Aes128Gcm::new(&key);
210        let nonce = Nonce::from_slice(nonce);
211
212        let ciphertext = cipher.encrypt(&nonce, pt.as_ref()).unwrap();
213
214        let plaintext = cipher.decrypt(&nonce, ciphertext.as_ref()).unwrap();
215        assert_eq!(&plaintext, pt);
216        Ok(())
217    }
218}