Enable simultaneous deserialization+decryption of a ChaChaPoly stream

valentinewallace · valentinewallace · commit aeaf125abe52 · 2022-06-15T17:51:29.000-04:00
In the upcoming onion messages PR, this will allow us to avoid decrypting onion
message encrypted data in an intermediate Vec before decoding it. Instead we
decrypt and decode it at the same time using this new ChaChaPolyReadAdapter object.

In doing so, we need to adapt the decode_tlv_stream macro such that it will
decode a LengthReadableArgs, which is a new trait as well. This trait is
necessary because ChaChaPoly needs to know the total length ahead of time to
separate out the tag at the end.
diff --git a/lightning/src/util/chacha20poly1305rfc.rs b/lightning/src/util/chacha20poly1305rfc.rs
@@ -10,8 +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};
+use ln::msgs::DecodeError;
+use util::ser::{FixedLengthReader, LengthRead, LengthReadableArgs, Readable, Writeable, Writer};
+use io::{self, Read, Write};
 
 #[cfg(not(fuzzing))]
 mod real_chachapoly {
@@ -93,22 +94,48 @@ mod real_chachapoly {
 		}
 
 		pub fn decrypt(&mut self, input: &[u8], output: &mut [u8], tag: &[u8]) -> bool {
-			assert!(input.len() == output.len());
-			assert!(self.finished == false);
+			if self.decrypt_inner(input, Some(output), Some(tag)) {
+				self.cipher.process(input, output);
+				return true
+			}
+			false
+		}
 
-			self.finished = true;
+		// Decrypt in place, and check the tag if it's provided. If the tag is not provided, then
+		// `finish_and_check_tag` may be called to check it later.
+		pub fn decrypt_in_place(&mut self, input: &mut [u8], tag: Option<&[u8]>) -> bool {
+			if self.decrypt_inner(input, None, tag) {
+				self.cipher.process_in_place(input);
+				return true
+			}
+			false
+		}
+
+		fn decrypt_inner(&mut self, input: &[u8], output: Option<&mut [u8]>, tag: Option<&[u8]>) -> bool {
+			if let Some(output) = output {
+				assert!(input.len() == output.len());
+			}
+			assert!(self.finished == false);
 
 			self.mac.input(input);
 
 			self.data_len += input.len();
+
+			if let Some(tag) = tag {
+				return self.finish_and_check_tag(tag)
+			}
+			true
+		}
+
+		pub fn finish_and_check_tag(&mut self, tag: &[u8]) -> bool {
+			self.finished = true;
 			ChaCha20Poly1305RFC::pad_mac_16(&mut self.mac, self.data_len);
 			self.mac.input(&self.aad_len.to_le_bytes());
 			self.mac.input(&(self.data_len as u64).to_le_bytes());
 
 			let mut calc_tag =  [0u8; 16];
 			self.mac.raw_result(&mut calc_tag);
 			if fixed_time_eq(&calc_tag, tag) {
-				self.cipher.process(input, output);
 				true
 			} else {
 				false
@@ -119,6 +146,25 @@ mod real_chachapoly {
 #[cfg(not(fuzzing))]
 pub use self::real_chachapoly::ChaCha20Poly1305RFC;
 
+/// Enables simultaneously reading and decrypting a ChaCha20Poly1305RFC stream from a std::io::Read.
+pub(crate) struct ChaChaPolyReader<'a, R: Read> {
+	pub chacha: &'a mut ChaCha20Poly1305RFC,
+	pub read: R,
+}
+
+impl<'a, R: Read> Read for ChaChaPolyReader<'a, R> {
+	// Decrypt bytes from Self::read into `dest`.
+	// `ChaCha20Poly1305RFC::finish_and_check_tag` must be called to check the tag after all reads
+	// complete.
+	fn read(&mut self, dest: &mut [u8]) -> Result<usize, io::Error> {
+		let res = self.read.read(dest)?;
+		if res > 0 {
+			self.chacha.decrypt_in_place(&mut dest[0..res], None);
+		}
+		Ok(res)
+	}
+}
+
 /// Enables simultaneously writing and encrypting a byte stream into a Writer.
 pub(crate) struct ChaChaPolyWriter<'a, W: Writer> {
 	pub chacha: &'a mut ChaCha20Poly1305RFC,
@@ -173,6 +219,36 @@ impl<'a, T: Writeable> Writeable for ChaChaPolyWriteAdapter<'a, T> {
 	}
 }
 
+/// Enables the use of the serialization macros for objects that need to be simultaneously decrypted and
+/// deserialized. This allows us to avoid an intermediate Vec allocation.
+pub(crate) struct ChaChaPolyReadAdapter<R: Readable> {
+	#[allow(unused)] // This will be used soon for onion messages
+	pub readable: R,
+}
+
+impl<T: Readable> LengthReadableArgs<[u8; 32]> for ChaChaPolyReadAdapter<T> {
+	// Simultaneously read and decrypt an object from a LengthRead, storing it in Self::readable.
+	// LengthRead must be used instead of std::io::Read because we need the total length to separate
+	// out the tag at the end.
+	fn read<R: LengthRead>(mut r: &mut R, secret: [u8; 32]) -> Result<Self, DecodeError> {
+		if r.total_bytes() < 16 { return Err(DecodeError::InvalidValue) }
+
+		let mut chacha = ChaCha20Poly1305RFC::new(&secret, &[0; 12], &[]);
+		let decrypted_len = r.total_bytes() - 16;
+		let s = FixedLengthReader::new(&mut r, decrypted_len);
+		let mut chacha_stream = ChaChaPolyReader { chacha: &mut chacha, read: s };
+		let readable: T = Readable::read(&mut chacha_stream)?;
+
+		let mut tag = [0 as u8; 16];
+		r.read_exact(&mut tag)?;
+		if !chacha.finish_and_check_tag(&tag) {
+			return Err(DecodeError::InvalidValue)
+		}
+
+		Ok(Self { readable })
+	}
+}
+
 #[cfg(fuzzing)]
 mod fuzzy_chachapoly {
 	#[derive(Clone, Copy)]
@@ -228,7 +304,111 @@ mod fuzzy_chachapoly {
 			self.finished = true;
 			true
 		}
+
+		pub fn decrypt_in_place(&mut self, _input: &mut [u8], tag: Option<&[u8]>) -> bool {
+			assert!(self.finished == false);
+			if let Some(tag) = tag {
+				if tag[..] != self.tag[..] { return false; }
+			}
+			self.finished = true;
+			true
+		}
+
+		pub fn finish_and_check_tag(&mut self, tag: &[u8]) -> bool {
+			if tag[..] != self.tag[..] { return false; }
+			self.finished = true;
+			true
+		}
 	}
 }
 #[cfg(fuzzing)]
 pub use self::fuzzy_chachapoly::ChaCha20Poly1305RFC;
+
+#[cfg(test)]
+mod tests {
+	use ln::msgs::DecodeError;
+	use super::{ChaChaPolyReadAdapter, ChaChaPolyWriteAdapter};
+	use util::ser::{self, FixedLengthReader, LengthReadableArgs, Writeable};
+
+	// Used for for testing various lengths of serialization.
+	#[derive(Debug, PartialEq)]
+	struct TestWriteable {
+		field1: Vec<u8>,
+		field2: Vec<u8>,
+		field3: Vec<u8>,
+	}
+	impl_writeable_tlv_based!(TestWriteable, {
+		(1, field1, vec_type),
+		(2, field2, vec_type),
+		(3, field3, vec_type),
+	});
+
+	#[test]
+	fn test_chacha_stream_adapters() {
+		// Check that ChaChaPolyReadAdapter and ChaChaPolyWriteAdapter correctly encode and decode an
+		// encrypted object.
+		macro_rules! check_object_read_write {
+			($obj: expr) => {
+				// First, serialize the object, encrypted with ChaCha20Poly1305.
+				let rho = [42; 32];
+				let writeable_len = $obj.serialized_length() as u64 + 16;
+				let write_adapter = ChaChaPolyWriteAdapter::new(rho, &$obj);
+				let encrypted_writeable_bytes = write_adapter.encode();
+				let encrypted_writeable = &encrypted_writeable_bytes[..];
+
+				// Now deserialize the object back and make sure it matches the original.
+				let mut rd = FixedLengthReader::new(encrypted_writeable, writeable_len);
+				let read_adapter = <ChaChaPolyReadAdapter<TestWriteable>>::read(&mut rd, rho).unwrap();
+				assert_eq!($obj, read_adapter.readable);
+			};
+		}
+
+		// Try a big object that will require multiple write buffers.
+		let big_writeable = TestWriteable {
+			field1: vec![43],
+			field2: vec![44; 4192],
+			field3: vec![45; 4192 + 1],
+		};
+		check_object_read_write!(big_writeable);
+
+		// Try a small object that fits into one write buffer.
+		let small_writeable = TestWriteable {
+			field1: vec![43],
+			field2: vec![44],
+			field3: vec![45],
+		};
+		check_object_read_write!(small_writeable);
+	}
+
+	fn do_chacha_stream_adapters_ser_macros() -> Result<(), DecodeError> {
+		let writeable = TestWriteable {
+			field1: vec![43],
+			field2: vec![44; 4192],
+			field3: vec![45; 4192 + 1],
+		};
+
+		// First, serialize the object into a TLV stream, encrypted with ChaCha20Poly1305.
+		let rho = [42; 32];
+		let write_adapter = ChaChaPolyWriteAdapter::new(rho, &writeable);
+		let mut writer = ser::VecWriter(Vec::new());
+		encode_tlv_stream!(&mut writer, {
+			(1, write_adapter, required),
+		});
+
+		// Now deserialize the object back and make sure it matches the original.
+		let mut read_adapter: Option<ChaChaPolyReadAdapter<TestWriteable>> = None;
+		decode_tlv_stream!(&writer.0[..], {
+			(1, read_adapter, (read_arg, rho)),
+		});
+		assert_eq!(writeable, read_adapter.unwrap().readable);
+
+		Ok(())
+	}
+
+	#[test]
+	fn chacha_stream_adapters_ser_macros() {
+		// Test that our stream adapters work as expected with the TLV macros.
+		// This also serves to test the `LengthReadableArgs` variant of the `decode_tlv` ser macro.
+		do_chacha_stream_adapters_ser_macros().unwrap()
+	}
+}
diff --git a/lightning/src/util/ser.rs b/lightning/src/util/ser.rs
@@ -134,6 +134,13 @@ impl<R: Read> Read for FixedLengthReader<R> {
 	}
 }
 
+impl<R: Read> LengthRead for FixedLengthReader<R> {
+	#[inline]
+	fn total_bytes(&self) -> u64 {
+		self.total_bytes
+	}
+}
+
 /// A Read which tracks whether any bytes have been read at all. This allows us to distinguish
 /// between "EOF reached before we started" and "EOF reached mid-read".
 pub(crate) struct ReadTrackingReader<R: Read> {
@@ -220,6 +227,26 @@ pub trait ReadableArgs<P>
 	fn read<R: Read>(reader: &mut R, params: P) -> Result<Self, DecodeError>;
 }
 
+/// A std::io::Read that also provides the total bytes available to read.
+///
+/// (C-not exported) as we only export serialization to/from byte arrays instead
+pub trait LengthRead where Self: Read {
+	/// The total number of bytes available to read.
+	fn total_bytes(&self) -> u64;
+}
+
+/// A trait that various higher-level rust-lightning types implement allowing them to be read in
+/// from a Read given some additional set of arguments which is required to deserialize, requiring
+/// the implementer to provide the total length of the read.
+///
+/// (C-not exported) as we only export serialization to/from byte arrays instead
+pub trait LengthReadableArgs<P>
+where Self: Sized
+{
+	/// Reads a Self in from the given LengthRead
+	fn read<R: LengthRead>(reader: &mut R, params: P) -> Result<Self, DecodeError>;
+}
+
 /// A trait that various rust-lightning types implement allowing them to (maybe) be read in from a Read
 ///
 /// (C-not exported) as we only export serialization to/from byte arrays instead
diff --git a/lightning/src/util/ser_macros.rs b/lightning/src/util/ser_macros.rs
@@ -118,6 +118,9 @@ macro_rules! check_tlv_order {
 	($last_seen_type: expr, $typ: expr, $type: expr, $field: ident, ignorable) => {{
 		// no-op
 	}};
+	($last_seen_type: expr, $typ: expr, $type: expr, $field: ident, (read_arg, $read_arg: expr)) => {{
+		// no-op
+	}};
 }
 
 macro_rules! check_missing_tlv {
@@ -144,6 +147,9 @@ macro_rules! check_missing_tlv {
 	($last_seen_type: expr, $type: expr, $field: ident, ignorable) => {{
 		// no-op
 	}};
+	($last_seen_type: expr, $type: expr, $field: ident, (read_arg, $read_arg: expr)) => {{
+		// no-op
+	}};
 }
 
 macro_rules! decode_tlv {
@@ -163,6 +169,10 @@ macro_rules! decode_tlv {
 	($reader: expr, $field: ident, ignorable) => {{
 		$field = ser::MaybeReadable::read(&mut $reader)?;
 	}};
+	// Read a ser::LengthReadableArgs
+	($reader: expr, $field: ident, (read_arg, $read_arg: expr)) => {{
+		$field = Some(ser::LengthReadableArgs::read(&mut $reader, $read_arg)?);
+	}};
 }
 
 macro_rules! decode_tlv_stream {
