1use byteorder::{ByteOrder, LittleEndian};
4
5pub const BLOCK_SIZE: usize = 16;
6pub const KEY_SIZE: usize = 16;
7pub const NONCE_SIZE: usize = (BLOCK_SIZE - 6) + 10;
8
9const ROUNDS: usize = 32;
10const ROUNDS_PER_STEP: usize = 4;
11const STEPS: usize = ROUNDS / ROUNDS_PER_STEP;
12
13pub type KeySchedule = [u32; ROUNDS_PER_STEP * (4 * STEPS + 1)];
14
15#[inline]
16fn spec_key(k: &mut u32) {
17    let mut left = *k as u16;
18    let mut right = (*k >> 16) as u16;
19    left = left.rotate_right(7).wrapping_add(right);
20    right = right.rotate_left(2) ^ left;
21    *k = (left as u32) | ((right as u32) << 16)
22}
23
24#[inline]
25fn spec_key_inv(k: &mut u32) {
26    let mut left = *k as u16;
27    let mut right = (*k >> 16) as u16;
28    right = (right ^ left).rotate_right(2);
29    left = left.wrapping_sub(right).rotate_left(7);
30    *k = (left as u32) | ((right as u32) << 16)
31}
32
33#[inline]
34fn key_perm(k: &mut [u32; 4], c: u16) {
35    spec_key(&mut k[0]);
36    k[1] = ((k[1] as u16).wrapping_add(k[0] as u16) as u32)
37        | (((k[1] >> 16) as u16).wrapping_add((k[0] >> 16) as u16) as u32) << 16;
38    spec_key(&mut k[2]);
39    k[3] = ((k[3] as u16).wrapping_add(k[2] as u16) as u32)
40        | (((k[3] >> 16) as u16)
41            .wrapping_add((k[2] >> 16) as u16)
42            .wrapping_add(c) as u32)
43            << 16;
44    let tmp = k[3];
45    k[3] = k[2];
46    k[2] = k[1];
47    k[1] = k[0];
48    k[0] = tmp;
49}
50
51pub fn key_schedule_encrypt(key: &[u8; KEY_SIZE]) -> KeySchedule {
53    let mut ks = [0u32; ROUNDS_PER_STEP * (4 * STEPS + 1)];
54    let mut k = [0u32; 4];
55    k[0] = LittleEndian::read_u32(&key[0 * 4..]);
56    k[1] = LittleEndian::read_u32(&key[1 * 4..]);
57    k[2] = LittleEndian::read_u32(&key[2 * 4..]);
58    k[3] = LittleEndian::read_u32(&key[3 * 4..]);
59    let mut j = 0;
60    for c in 0..(4 * STEPS + 1) {
61        for &ksp in k.iter().take(ROUNDS_PER_STEP) {
62            ks[j] = ksp;
63            j += 1;
64        }
65        key_perm(&mut k, c as u16 + 1);
66    }
67    ks
68}
69
70pub fn key_schedule_decrypt(key: &[u8; KEY_SIZE]) -> KeySchedule {
72    key_schedule_encrypt(key)
73}
74
75pub fn encrypt_block(block: &mut [u8; BLOCK_SIZE], ks: &KeySchedule) {
77    let mut ksi = ks.iter();
78    for _ in 0..STEPS {
79        for b in 0..4 {
80            for _ in 0..ROUNDS_PER_STEP {
81                let mut tmp = LittleEndian::read_u32(&block[4 * b..]) ^ ksi.next().unwrap();
82                spec_key(&mut tmp);
83                LittleEndian::write_u32(&mut block[4 * b..], tmp);
84            }
85        }
86        let x0 = LittleEndian::read_u32(&block[4 * 0..]);
87        let x1 = LittleEndian::read_u32(&block[4 * 1..]);
88        let mut x2 = LittleEndian::read_u32(&block[4 * 2..]);
89        let mut x3 = LittleEndian::read_u32(&block[4 * 3..]);
90        let tmp =
91            ((x0 as u16) ^ ((x0 >> 16) as u16) ^ (x1 as u16) ^ ((x1 >> 16) as u16)).rotate_left(8);
92        let tmp = (tmp as u32) | ((tmp as u32) << 16);
93        x2 = ((((x2 as u16) ^ (x1 as u16)) as u32)
94            | ((((x2 >> 16) as u16) ^ ((x0 >> 16) as u16)) as u32) << 16)
95            ^ tmp;
96        x3 = ((((x3 as u16) ^ (x0 as u16)) as u32)
97            | ((((x3 >> 16) as u16) ^ ((x1 >> 16) as u16)) as u32) << 16)
98            ^ tmp;
99        LittleEndian::write_u32(&mut block[4 * 0..], x2);
100        LittleEndian::write_u32(&mut block[4 * 1..], x3);
101        LittleEndian::write_u32(&mut block[4 * 2..], x0);
102        LittleEndian::write_u32(&mut block[4 * 3..], x1);
103    }
104    for b in 0..4 {
105        let tmp = LittleEndian::read_u32(&block[4 * b..]) ^ ksi.next().unwrap();
106        LittleEndian::write_u32(&mut block[4 * b..], tmp);
107    }
108}
109
110pub fn decrypt_block(block: &mut [u8; BLOCK_SIZE], ks: &KeySchedule) {
112    let mut ksi = ks.iter().rev();
113    for b in (0..4).rev() {
114        let tmp = LittleEndian::read_u32(&block[4 * b..]) ^ ksi.next().unwrap();
115        LittleEndian::write_u32(&mut block[4 * b..], tmp);
116    }
117    for _ in 0..STEPS {
118        let x0 = LittleEndian::read_u32(&block[4 * 2..]);
119        let x1 = LittleEndian::read_u32(&block[4 * 3..]);
120        let mut x2 = LittleEndian::read_u32(&block[4 * 0..]);
121        let mut x3 = LittleEndian::read_u32(&block[4 * 1..]);
122        let tmp =
123            ((x0 as u16) ^ ((x0 >> 16) as u16) ^ (x1 as u16) ^ ((x1 >> 16) as u16)).rotate_left(8);
124        let tmp = (tmp as u32) | ((tmp as u32) << 16);
125        x2 = ((((x2 as u16) ^ (x1 as u16)) as u32)
126            | ((((x2 >> 16) as u16) ^ ((x0 >> 16) as u16)) as u32) << 16)
127            ^ tmp;
128        x3 = ((((x3 as u16) ^ (x0 as u16)) as u32)
129            | ((((x3 >> 16) as u16) ^ ((x1 >> 16) as u16)) as u32) << 16)
130            ^ tmp;
131        LittleEndian::write_u32(&mut block[4 * 0..], x0);
132        LittleEndian::write_u32(&mut block[4 * 1..], x1);
133        LittleEndian::write_u32(&mut block[4 * 2..], x2);
134        LittleEndian::write_u32(&mut block[4 * 3..], x3);
135        for b in (0..4).rev() {
136            for _ in 0..ROUNDS_PER_STEP {
137                let mut tmp = LittleEndian::read_u32(&block[4 * b..]);
138                spec_key_inv(&mut tmp);
139                LittleEndian::write_u32(&mut block[4 * b..], tmp ^ ksi.next().unwrap());
140            }
141        }
142    }
143}
144
145pub fn encrypt_ctr(buf: &mut [u8], nonce: &[u8; NONCE_SIZE], key: &[u8; KEY_SIZE]) {
147    if buf.is_empty() {
148        return;
149    }
150    let mut key2 = [0u8; KEY_SIZE];
151    for (i, &x) in key.iter().enumerate().take(10) {
152        key2[i] = x ^ nonce[BLOCK_SIZE - 6 + i];
153    }
154    let ks = key_schedule_encrypt(&key2);
155    let full_blocks_count = (buf.len() / BLOCK_SIZE) as u64;
156    let mut ib = [0u8; BLOCK_SIZE];
157    let mut nc = [0u8; BLOCK_SIZE];
158    nc[..(BLOCK_SIZE - 6)].clone_from_slice(&nonce[..(BLOCK_SIZE - 6)]);
159    let mut n = 0;
160    for i in 0..full_blocks_count {
161        ib.copy_from_slice(&nc);
162        LittleEndian::write_u32(&mut ib[BLOCK_SIZE - 6..], i as u32);
163        LittleEndian::write_u16(&mut ib[BLOCK_SIZE - 2..], (i >> 32) as u16);
164        encrypt_block(&mut ib, &ks);
165        for (j, &ob) in ib.iter().enumerate() {
166            buf[n + j] ^= ob;
167        }
168        n += BLOCK_SIZE;
169    }
170    let remaining_bytes = buf.len() % BLOCK_SIZE;
171    if remaining_bytes > 0 {
172        ib.copy_from_slice(&nc);
173        LittleEndian::write_u32(&mut ib[BLOCK_SIZE - 6..], full_blocks_count as u32);
174        LittleEndian::write_u16(&mut ib[BLOCK_SIZE - 2..], (full_blocks_count >> 32) as u16);
175        encrypt_block(&mut ib, &ks);
176        for (j, &ob) in ib.iter().enumerate().take(remaining_bytes) {
177            buf[n + j] ^= ob
178        }
179    }
180}
181
182pub fn decrypt_ctr(buf: &mut [u8], nonce: &[u8; NONCE_SIZE], key: &[u8; KEY_SIZE]) {
184    encrypt_ctr(buf, nonce, key)
185}
186
187#[test]
188fn test_vector() {
189    let key: [u8; KEY_SIZE] = [
190        0x11, 0x00, 0x33, 0x22, 0x55, 0x44, 0x77, 0x66, 0x99, 0x88, 0xbb, 0xaa, 0xdd, 0xcc, 0xff,
191        0xee,
192    ];
193    let mut block: [u8; BLOCK_SIZE] = [
194        0x23, 0x01, 0x67, 0x45, 0xab, 0x89, 0xef, 0xcd, 0xdc, 0xfe, 0x98, 0xba, 0x54, 0x76, 0x10,
195        0x32,
196    ];
197    let ks = key_schedule_encrypt(&key);
198    let block2 = block;
199    encrypt_block(&mut block, &ks);
200    assert_eq!(
201        [
202            0xee, 0x1c, 0x40, 0x75, 0xbf, 0x7d, 0xd8, 0x23, 0xee, 0xe0, 0x97, 0x15, 0x28, 0xf4,
203            0xd8, 0x52,
204        ],
205        block
206    );
207    decrypt_block(&mut block, &ks);
208    assert_eq!(block2, block);
209}
210
211#[test]
212fn test_ctr() {
213    let nonce: [u8; NONCE_SIZE] = [
214        1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20,
215    ];
216    let key: [u8; KEY_SIZE] = [
217        0x11, 0x00, 0x33, 0x22, 0x55, 0x44, 0x77, 0x66, 0x99, 0x88, 0xbb, 0xaa, 0xdd, 0xcc, 0xff,
218        0xee,
219    ];
220    let input = b"The quick brown fox jumps over the lazy dog";
221    let mut buf = input.to_vec();
222    let expected: [u8; 43] = [
223        187, 164, 175, 197, 150, 51, 7, 71, 250, 16, 102, 26, 154, 89, 226, 186, 171, 49, 0, 228,
224        255, 249, 53, 223, 223, 97, 25, 144, 13, 185, 170, 216, 79, 219, 40, 137, 95, 164, 73, 201,
225        65, 42, 58,
226    ];
227    encrypt_ctr(&mut buf, &nonce, &key);
228    assert_eq!(buf[..], expected[..]);
229    decrypt_ctr(&mut buf, &nonce, &key);
230    assert_eq!(buf[..], input[..]);
231}