From b385a40e4aadace8bd6646c42c6c662b8f1901e9 Mon Sep 17 00:00:00 2001
From: Matt Corallo <git@bluematt.me>
Date: Sun, 30 May 2021 23:24:43 +0000
Subject: [PATCH] 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.
---
 lightning/src/ln/msgs.rs  |  2 +-
 lightning/src/util/ser.rs | 19 +++++++++++--------
 2 files changed, 12 insertions(+), 9 deletions(-)

diff --git a/lightning/src/ln/msgs.rs b/lightning/src/ln/msgs.rs
index eec1a053..f02a4783 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 7bc09e7f..a49b033d 100644
--- a/lightning/src/util/ser.rs
+++ b/lightning/src/util/ser.rs
@@ -249,28 +249,31 @@ 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),
 			}
 		}
-- 
2.30.2