1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148

const CHACHA_TAU: &[u8] = b"expand 32-byte k";

fn chacha_quarter_round(state: &mut [u32; 16], a: usize, b: usize, c: usize, d: usize) {
	state[a] = state[a].wrapping_add(state[b]);
	state[d] ^= state[a];
	state[d] = state[d].rotate_left(16);

	state[c] = state[c].wrapping_add(state[d]);
	state[b] ^= state[c];
	state[b] = state[b].rotate_left(12);

	state[a] = state[a].wrapping_add(state[b]);
	state[d] ^= state[a];
	state[d] = state[d].rotate_left(8);

	state[c] = state[c].wrapping_add(state[d]);
	state[b] ^= state[c];
	state[b] = state[b].rotate_left(7);
}

const fn chacha_pack(unpacked: &[u8], idx: usize) -> u32 {
	(unpacked[idx] as u32)
		| ((unpacked[idx + 1] as u32) << 8)
		| ((unpacked[idx + 2] as u32) << 16)
		| ((unpacked[idx + 3] as u32) << 24)
}

/// Do one ChaCha round on the input data.
pub fn chacha_block<const ROUNDS: u8>(input: [u32; 16]) -> [u32; 16] {
	let mut x = input;
	assert_eq!(ROUNDS % 2, 0, "ChaCha rounds must be divisble by 2!");
	for _ in (0..ROUNDS).step_by(2) {
		// Odd rounds
		chacha_quarter_round(&mut x, 0, 4, 8, 12);
		chacha_quarter_round(&mut x, 1, 5, 9, 13);
		chacha_quarter_round(&mut x, 2, 6, 10, 14);
		chacha_quarter_round(&mut x, 3, 7, 11, 15);
		// Even rounds
		chacha_quarter_round(&mut x, 0, 5, 10, 15);
		chacha_quarter_round(&mut x, 1, 6, 11, 12);
		chacha_quarter_round(&mut x, 2, 7, 8, 13);
		chacha_quarter_round(&mut x, 3, 4, 9, 14);
	}
	x.iter_mut()
		.zip(input.iter())
		.for_each(|(l, r)| *l = l.wrapping_add(*r));
	x
}

/// Initialize the ChaCha internal state, with a 256-bit key and 64-bit nonce.
pub const fn chacha_init(key: [u8; 32], nonce: [u8; 8]) -> [u32; 16] {
	let mut state = [0u32; 16];
	state[0] = chacha_pack(CHACHA_TAU, 0);
	state[1] = chacha_pack(CHACHA_TAU, 4);
	state[2] = chacha_pack(CHACHA_TAU, 8);
	state[3] = chacha_pack(CHACHA_TAU, 12);

	state[4] = chacha_pack(&key, 0);
	state[5] = chacha_pack(&key, 4);
	state[6] = chacha_pack(&key, 8);
	state[7] = chacha_pack(&key, 12);
	state[8] = chacha_pack(&key, 16);
	state[9] = chacha_pack(&key, 20);
	state[10] = chacha_pack(&key, 24);
	state[11] = chacha_pack(&key, 28);

	// 64-bit counter
	state[12] = 0;
	state[13] = 0;
	// Nonce
	state[14] = chacha_pack(&nonce, 0);
	state[15] = chacha_pack(&nonce, 4);
	state
}

/// Increment the 64-bit counter of the internal ChaCha20 state by 1.
/// Returns `false` if it overflows, `true` otherwise.
pub fn chacha_increment_counter(state: &mut [u32; 16]) -> bool {
	let counter = ((state[13] as u64) << 32) | (state[12] as u64);
	match counter.checked_add(1) {
		Some(new_counter) => {
			state[12] = (new_counter & 0xFFFFFFFF) as u32;
			state[13] = ((counter >> 32) & 0xFFFFFFFF) as u32;
			true
		}
		None => false,
	}
}

#[cfg(test)]
mod tests {
	use super::*;
	use std::convert::TryInto;

	macro_rules! ietf_test_vector {
		($key_hex: tt, $nonce_hex: tt, $keystream_hex: tt) => {
			let key: [u8; 32] = hex::decode($key_hex).unwrap().try_into().unwrap();
			let nonce: [u8; 8] = hex::decode($nonce_hex).unwrap().try_into().unwrap();
			let expected_keystream: Vec<u8> = hex::decode($keystream_hex).unwrap();

			let mut state = chacha_init(key, nonce);
			let mut keystream: Vec<u8> = Vec::with_capacity(expected_keystream.len());

			while expected_keystream.len() > keystream.len() {
				chacha_block::<20>(state)
					.iter()
					.for_each(|packed| keystream.extend_from_slice(&packed.to_le_bytes()));
				chacha_increment_counter(&mut state);
			}
			keystream.resize(expected_keystream.len(), 0);

			assert_eq!(keystream, expected_keystream);
		};
	}

	#[test]
	fn test_ietf_chacha20_test_vectors() {
		ietf_test_vector!(
			"0000000000000000000000000000000000000000000000000000000000000000",
			"0000000000000000",
			"76b8e0ada0f13d90405d6ae55386bd28bdd219b8a08ded1aa836efcc8b770dc7da41597c5157488d7724e03fb8d84a376a43b8f41518a11cc387b669b2ee6586"
		);

		ietf_test_vector!(
			"0000000000000000000000000000000000000000000000000000000000000001",
			"0000000000000000",
			"4540f05a9f1fb296d7736e7b208e3c96eb4fe1834688d2604f450952ed432d41bbe2a0b6ea7566d2a5d1e7e20d42af2c53d792b1c43fea817e9ad275ae546963"
		);

		ietf_test_vector!(
			"0000000000000000000000000000000000000000000000000000000000000000",
			"0000000000000001",
			"de9cba7bf3d69ef5e786dc63973f653a0b49e015adbff7134fcb7df137821031e85a050278a7084527214f73efc7fa5b5277062eb7a0433e445f41e3"
		);

		ietf_test_vector!(
			"0000000000000000000000000000000000000000000000000000000000000000",
			"0100000000000000",
			"ef3fdfd6c61578fbf5cf35bd3dd33b8009631634d21e42ac33960bd138e50d32111e4caf237ee53ca8ad6426194a88545ddc497a0b466e7d6bbdb0041b2f586b"
		);

		ietf_test_vector!(
			"000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f",
			"0001020304050607",
			"f798a189f195e66982105ffb640bb7757f579da31602fc93ec01ac56f85ac3c134a4547b733b46413042c9440049176905d3be59ea1c53f15916155c2be8241a38008b9a26bc35941e2444177c8ade6689de95264986d95889fb60e84629c9bd9a5acb1cc118be563eb9b3a4a472f82e09a7e778492b562ef7130e88dfe031c79db9d4f7c7a899151b9a475032b63fc385245fe054e3dd5a97a5f576fe064025d3ce042c566ab2c507b138db853e3d6959660996546cc9c4a6eafdc777c040d70eaf46f76dad3979e5c5360c3317166a1c894c94a371876a94df7628fe4eaaf2ccb27d5aaae0ad7ad0f9d4b6ad3b54098746d4524d38407a6deb3ab78fab78c9"
		);
	}
}