From: Matt Corallo <git@bluematt.me>
Date: Sun, 30 May 2021 23:24:43 +0000 (+0000)
Subject: Use TLV instead of explicit length for VecReadWrapper termination
X-Git-Url: http://git.bitcoin.ninja/index.cgi?a=commitdiff_plain;h=b06a5afd5a07fcb5d446a38f8930ed3258f50c09;p=rust-lightning

Use TLV instead of explicit length for VecReadWrapper termination

VecReadWrapper is only used in TLVs so there is no need to prepend
a length before writing/reading the objects - we can instead simply
read until we reach the end of the TLV stream.
---

diff --git a/lightning/src/ln/msgs.rs b/lightning/src/ln/msgs.rs
index eec1a0538..f02a47835 100644
--- a/lightning/src/ln/msgs.rs
+++ b/lightning/src/ln/msgs.rs
@@ -46,7 +46,7 @@ use ln::{PaymentPreimage, PaymentHash, PaymentSecret};
 pub(crate) const MAX_VALUE_MSAT: u64 = 21_000_000_0000_0000_000;
 
 /// An error in decoding a message or struct.
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, PartialEq)]
 pub enum DecodeError {
 	/// A version byte specified something we don't know how to handle.
 	/// Includes unknown realm byte in an OnionHopData packet
diff --git a/lightning/src/util/ser.rs b/lightning/src/util/ser.rs
index 7bc09e7f8..06578b9fb 100644
--- a/lightning/src/util/ser.rs
+++ b/lightning/src/util/ser.rs
@@ -249,28 +249,30 @@ impl<T: Readable> Readable for OptionDeserWrapper<T> {
 	}
 }
 
-const MAX_ALLOC_SIZE: u64 = 64*1024;
-
+/// Wrapper to write each element of a Vec with no length prefix
 pub(crate) struct VecWriteWrapper<'a, T: Writeable>(pub &'a Vec<T>);
 impl<'a, T: Writeable> Writeable for VecWriteWrapper<'a, T> {
 	#[inline]
 	fn write<W: Writer>(&self, writer: &mut W) -> Result<(), ::std::io::Error> {
-		(self.0.len() as u64).write(writer)?;
 		for ref v in self.0.iter() {
 			v.write(writer)?;
 		}
 		Ok(())
 	}
 }
+/// Wrapper to read elements from a given stream until it reaches the end of the stream.
 pub(crate) struct VecReadWrapper<T: Readable>(pub Vec<T>);
 impl<T: Readable> Readable for VecReadWrapper<T> {
 	#[inline]
-	fn read<R: Read>(reader: &mut R) -> Result<Self, DecodeError> {
-		let count: u64 = Readable::read(reader)?;
-		let mut values = Vec::with_capacity(cmp::min(count, MAX_ALLOC_SIZE / (core::mem::size_of::<T>() as u64)) as usize);
-		for _ in 0..count {
-			match Readable::read(reader) {
+	fn read<R: Read>(mut reader: &mut R) -> Result<Self, DecodeError> {
+		let mut values = Vec::new();
+		loop {
+			let mut track_read = ReadTrackingReader::new(&mut reader);
+			match Readable::read(&mut track_read) {
 				Ok(v) => { values.push(v); },
+				// If we failed to read any bytes at all, we reached the end of our TLV
+				// stream and have simply exhausted all entries.
+				Err(ref e) if e == &DecodeError::ShortRead && !track_read.have_read => break,
 				Err(e) => return Err(e),
 			}
 		}