chacha20poly1305: enable simultaneous writing+encryption

In the upcoming onion messages PR, this will allow us to avoid encoding onion
message encrypted data into an intermediate Vec before encrypting it.  Instead
we encode and encrypt at the same time using this new ChaChaPolyWriteAdapter object.

Author: valentinewallace

lightning/src/util/chacha20poly1305rfc.rs

@@ -10,6 +10,9 @@
 // This is a port of Andrew Moons poly1305-donna
 // https://github.com/floodyberry/poly1305-donna
 
+use util::ser::{Writeable, Writer};
+use io::{self, Write};
+
 #[cfg(not(fuzzing))]
 mod real_chachapoly {
 	use util::chacha20::ChaCha20;
@@ -58,11 +61,32 @@ mod real_chachapoly {
 		}
 
 		pub fn encrypt(&mut self, input: &[u8], output: &mut [u8], out_tag: &mut [u8]) {
-			assert!(input.len() == output.len());
-			assert!(self.finished == false);
 			self.cipher.process(input, output);
+			self.encrypt_inner(input, Some(output), Some(out_tag));
+		}
+
+		pub fn encrypt_in_place(&mut self, input_output: &mut [u8], out_tag: Option<&mut [u8]>) {
+			self.cipher.process_in_place(input_output);
+			self.encrypt_inner(input_output, None, out_tag);
+		}
+
+		// Encrypt in place, and fill in the tag if it's provided. If the tag is not provided, then
+		// `finish_and_get_tag` may be called to check it later.
+		fn encrypt_inner(&mut self, input: &[u8], output: Option<&mut [u8]>, out_tag: Option<&mut [u8]>) {
+			assert!(self.finished == false);
+			if let Some(output) = output {
+				assert!(input.len() == output.len());
+				self.mac.input(output);
+			} else {
+				self.mac.input(input);
+			}
 			self.data_len += input.len();
-			self.mac.input(output);
+			if let Some(tag) = out_tag {
+				self.finish_and_get_tag(tag);
+			}
+		}
+
+		pub fn finish_and_get_tag(&mut self, out_tag: &mut [u8]) {
 			ChaCha20Poly1305RFC::pad_mac_16(&mut self.mac, self.data_len);
 			self.finished = true;
 			self.mac.input(&self.aad_len.to_le_bytes());
@@ -97,6 +121,56 @@ mod real_chachapoly {
 #[cfg(not(fuzzing))]
 pub use self::real_chachapoly::ChaCha20Poly1305RFC;
 
+/// Enables simultaneously writing and encrypting a byte stream into a Writer.
+pub(crate) struct ChaChaPolyWriter<'a, W: Writer> {
+	pub chacha: &'a mut ChaCha20Poly1305RFC,
+	pub write: &'a mut W,
+}
+
+impl<'a, W: Writer> Writer for ChaChaPolyWriter<'a, W> {
+	// Encrypt then write bytes from `src` into Self::write.
+	// `ChaCha20Poly1305RFC::finish_and_get_tag` can be called to retrieve the tag after all writes
+	// complete.
+	fn write_all(&mut self, src: &[u8]) -> Result<(), io::Error> {
+		let num_writes = (src.len() + (8192 - 1)) / 8192;
+		for i in 0..num_writes {
+			let mut write_buffer = [0; 8192];
+			let bytes_written = (&mut write_buffer[..]).write(&src[i * 8192..])?;
+			self.chacha.encrypt_in_place(&mut write_buffer[..bytes_written], None);
+			self.write.write_all(&write_buffer[..bytes_written])?;
+		}
+		Ok(())
+	}
+}
+
+/// Enables the use of the serialization macros for objects that need to be simultaneously encrypted and
+/// serialized. This allows us to avoid an intermediate Vec allocation.
+pub(crate) struct ChaChaPolyWriteAdapter<'a, W: Writeable> {
+	pub rho: [u8; 32],
+	pub writeable: &'a W,
+}
+
+impl<'a, W: Writeable> ChaChaPolyWriteAdapter<'a, W> {
+	#[allow(unused)] // This will be used for onion messages soon
+	pub fn new(rho: [u8; 32], writeable: &'a W) -> ChaChaPolyWriteAdapter<'a, W> {
+		Self { rho, writeable }
+	}
+}
+
+impl<'a, T: Writeable> Writeable for ChaChaPolyWriteAdapter<'a, T> {
+	// Simultaneously write and encrypt Self::writeable.
+	fn write<W: Writer>(&self, w: &mut W) -> Result<(), io::Error> {
+		let mut chacha = ChaCha20Poly1305RFC::new(&self.rho, &[0; 12], &[]);
+		let mut chacha_stream = ChaChaPolyWriter { chacha: &mut chacha, write: w };
+		self.writeable.write(&mut chacha_stream)?;
+		let mut tag = [0 as u8; 16];
+		chacha.finish_and_get_tag(&mut tag);
+		tag.write(w)?;
+
+		Ok(())
+	}
+}
+
 #[cfg(fuzzing)]
 mod fuzzy_chachapoly {
 	#[derive(Clone, Copy)]
@@ -130,6 +204,19 @@ mod fuzzy_chachapoly {
 			self.finished = true;
 		}
 
+		pub fn encrypt_in_place(&mut self, _input_output: &mut [u8], out_tag: Option<&mut [u8]>) {
+			assert!(self.finished == false);
+			if let Some(tag) = out_tag {
+				tag.copy_from_slice(&self.tag);
+				self.finished = true;
+			}
+		}
+
+		pub fn finish_and_get_tag(&mut self, out_tag: &mut [u8]) {
+			out_tag.copy_from_slice(&self.tag);
+			self.finished = true;
+		}
+
 		pub fn decrypt(&mut self, input: &[u8], output: &mut [u8], tag: &[u8]) -> bool {
 			assert!(input.len() == output.len());
 			assert!(self.finished == false);