Add get_single_block to chacha20 module

In the next commit, we'll want to get a single block from a chacha stream that
takes a 16-byte nonce.

Author: valentinewallace

lightning/src/util/chacha20.rs

@@ -151,6 +151,17 @@ mod real_chacha {
 			ChaCha20{ state: ChaCha20::expand(key, nonce), output: [0u8; BLOCK_SIZE], offset: 64 }
 		}
 
+		/// Get one block from a ChaCha stream. Panics if `key`'s len is not 32 bytes or `nonce`'s len
+		/// is not 16 bytes.
+		pub fn get_single_block(key: &[u8], nonce: &[u8]) -> [u8; 32] {
+			assert!(key.len() == 32);
+			assert!(nonce.len() == 16);
+			let mut chacha = ChaCha20 { state: ChaCha20::expand(key, nonce), output: [0u8; BLOCK_SIZE], offset: 64 };
+			let mut chacha_bytes = [0; 32];
+			chacha.process_in_place(&mut chacha_bytes);
+			chacha_bytes
+		}
+
 		fn expand(key: &[u8], nonce: &[u8]) -> ChaChaState {
 			let constant = match key.len() {
 				16 => b"expand 16-byte k",
@@ -256,6 +267,16 @@ mod real_chacha {
 				self.offset += count;
 			}
 		}
+
+		#[cfg(test)]
+		pub fn seek(&mut self, block_offset: u32, offset_in_block: usize) -> Result<(), ()> {
+			let u32x4(_, d2, d3, d4) = self.state.d;
+			self.state.d = u32x4(block_offset, d2, d3, d4);
+			self.update();
+			self.offset = offset_in_block;
+
+			Ok(())
+		}
 	}
 }
 #[cfg(not(feature = "fuzztarget"))]
@@ -272,6 +293,12 @@ mod fuzzy_chacha {
 			Self {}
 		}
 
+		pub fn get_single_block(key: &[u8], nonce: &[u8]) -> [u8; 32] {
+			assert!(key.len() == 32);
+			assert!(nonce.len() == 16);
+			[0; 32]
+		}
+
 		pub fn process(&mut self, input: &[u8], output: &mut [u8]) {
 			output.copy_from_slice(input);
 		}
@@ -302,6 +329,7 @@ mod test {
 	use core::iter::repeat;
 
 	use super::ChaCha20;
+	use std::convert::TryInto;
 
 	#[test]
 	fn test_chacha20_256_tls_vectors() {
@@ -572,4 +600,35 @@ mod test {
 			assert_eq!(output, tv.keystream);
 		}
 	}
+
+	#[test]
+	fn get_single_block() {
+		// Test that `get_single_block` is equivalent to getting a block from a 12-byte nonce from the
+		// counter offset given by the remaining 4 bytes.
+		let key = [
+			0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
+			0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
+			0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17,
+			0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f,
+		];
+		let nonce_12bytes = [
+			0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
+			0x08, 0x09, 0x0a, 0x0b
+		];
+		let counter_pos = [0x00, 0x01, 0x02, 0x03];
+		let nonce_16bytes = [ // nonce_12bytes prefixed by counter_pos
+			0x00, 0x01, 0x02, 0x03,
+			0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
+			0x08, 0x09, 0x0a, 0x0b
+		];
+
+		// Initialize a ChaCha20 instance with its counter starting at 0.
+		let mut chacha20 = ChaCha20::new(&key, &nonce_12bytes);
+		// Seek its counter to counter_pos.
+		chacha20.seek(u32::from_le_bytes(counter_pos).try_into().unwrap(), 0).unwrap();
+		let mut block_bytes = [0; 32];
+		chacha20.process_in_place(&mut block_bytes);
+
+		assert_eq!(ChaCha20::get_single_block(&key, &nonce_16bytes), block_bytes);
+	}
 }