Add size_hint in ser and call size_hint in all message serializers

Author: TheBlueMatt

src/ln/msgs.rs

@@ -1704,7 +1704,10 @@ impl MsgDecodable for ErrorMessage {
 	}
 }
 
-impl_writeable!(AcceptChannel, {
+impl_writeable_len_match!(AcceptChannel, {
+		{AcceptChannel{ shutdown_scriptpubkey: Some(ref script), ..}, 270 + 2 + script.len()},
+		{_, 270}
+	}, {
 	temporary_channel_id,
 	dust_limit_satoshis,
 	max_htlc_value_in_flight_msat,
@@ -1722,7 +1725,7 @@ impl_writeable!(AcceptChannel, {
 	shutdown_scriptpubkey
 });
 
-impl_writeable!(AnnouncementSignatures, {
+impl_writeable!(AnnouncementSignatures, 32+8+64*2, {
 	channel_id,
 	short_channel_id,
 	node_signature,
@@ -1731,6 +1734,7 @@ impl_writeable!(AnnouncementSignatures, {
 
 impl<W: Writer> Writeable<W> for ChannelReestablish {
 	fn write(&self, w: &mut W) -> Result<(), DecodeError> {
+		w.size_hint(if self.data_loss_protect.is_some() { 32+2*8+33+32 } else { 32+2*8 });
 		self.channel_id.write(w)?;
 		self.next_local_commitment_number.write(w)?;
 		self.next_remote_commitment_number.write(w)?;
@@ -1763,55 +1767,68 @@ impl<R: Read> Readable<R> for ChannelReestablish{
 	}
 }
 
-impl_writeable!(ClosingSigned, {
+impl_writeable!(ClosingSigned, 32+8+64, {
 	channel_id,
 	fee_satoshis,
 	signature
 });
 
-impl_writeable!(CommitmentSigned, {
+impl_writeable_len_match!(CommitmentSigned, {
+		{ CommitmentSigned { ref htlc_signatures, .. }, 32+64+2+htlc_signatures.len()*64 }
+	}, {
 	channel_id,
 	signature,
 	htlc_signatures
 });
 
-impl_writeable!(DecodedOnionErrorPacket, {
+impl_writeable_len_match!(DecodedOnionErrorPacket, {
+		{ DecodedOnionErrorPacket { ref failuremsg, ref pad, .. }, 32 + 4 + failuremsg.len() + pad.len() }
+	}, {
 	hmac,
 	failuremsg,
 	pad
 });
 
-impl_writeable!(FundingCreated, {
+impl_writeable!(FundingCreated, 32+32+2+64, {
 	temporary_channel_id,
 	funding_txid,
 	funding_output_index,
 	signature
 });
 
-impl_writeable!(FundingSigned, {
+impl_writeable!(FundingSigned, 32+64, {
 	channel_id,
 	signature
 });
 
-impl_writeable!(FundingLocked, {
+impl_writeable!(FundingLocked, 32+33, {
 	channel_id,
 	next_per_commitment_point
 });
 
-impl_writeable!(GlobalFeatures, {
+impl_writeable_len_match!(GlobalFeatures, {
+		{ GlobalFeatures { ref flags }, flags.len() + 2 }
+	}, {
 	flags
 });
 
-impl_writeable!(LocalFeatures, {
+impl_writeable_len_match!(LocalFeatures, {
+		{ LocalFeatures { ref flags }, flags.len() + 2 }
+	}, {
 	flags
 });
 
-impl_writeable!(Init, {
+impl_writeable_len_match!(Init, {
+		{ Init { ref global_features, ref local_features }, global_features.flags.len() + local_features.flags.len() + 4 }
+	}, {
 	global_features,
 	local_features
 });
 
-impl_writeable!(OpenChannel, {
+impl_writeable_len_match!(OpenChannel, {
+		{ OpenChannel { shutdown_scriptpubkey: Some(ref script), .. }, 319 + 2 + script.len() },
+		{ OpenChannel { shutdown_scriptpubkey: None, .. }, 319 }
+	}, {
 	chain_hash,
 	temporary_channel_id,
 	funding_satoshis,
@@ -1833,47 +1850,54 @@ impl_writeable!(OpenChannel, {
 	shutdown_scriptpubkey
 });
 
-impl_writeable!(RevokeAndACK, {
+impl_writeable!(RevokeAndACK, 32+32+33, {
 	channel_id,
 	per_commitment_secret,
 	next_per_commitment_point
 });
 
-impl_writeable!(Shutdown, {
+impl_writeable_len_match!(Shutdown, {
+		{ Shutdown { ref scriptpubkey, .. }, 32 + 2 + scriptpubkey.len() }
+	}, {
 	channel_id,
 	scriptpubkey
 });
 
-impl_writeable!(UpdateFailHTLC, {
+impl_writeable_len_match!(UpdateFailHTLC, {
+		{ UpdateFailHTLC { ref reason, .. }, 32 + 10 + reason.data.len() }
+	}, {
 	channel_id,
 	htlc_id,
 	reason
 });
 
-impl_writeable!(UpdateFailMalformedHTLC, {
+impl_writeable!(UpdateFailMalformedHTLC, 32+8+32+2, {
 	channel_id,
 	htlc_id,
 	sha256_of_onion,
 	failure_code
 });
 
-impl_writeable!(UpdateFee, {
+impl_writeable!(UpdateFee, 32+4, {
 	channel_id,
 	feerate_per_kw
 });
 
-impl_writeable!(UpdateFulfillHTLC, {
+impl_writeable!(UpdateFulfillHTLC, 32+8+32, {
 	channel_id,
 	htlc_id,
 	payment_preimage
 });
 
-impl_writeable!(OnionErrorPacket, {
+impl_writeable_len_match!(OnionErrorPacket, {
+		{ OnionErrorPacket { ref data, .. }, 2 + data.len() }
+	}, {
 	data
 });
 
 impl<W: Writer> Writeable<W> for OnionPacket {
 	fn write(&self, w: &mut W) -> Result<(), DecodeError> {
+		w.size_hint(1 + 33 + 20*65 + 32);
 		self.version.write(w)?;
 		match self.public_key {
 			Ok(pubkey) => pubkey.write(w)?,
@@ -1900,7 +1924,7 @@ impl<R: Read> Readable<R> for OnionPacket {
 	}
 }
 
-impl_writeable!(UpdateAddHTLC, {
+impl_writeable!(UpdateAddHTLC, 32+8+8+32+4+1366, {
 	channel_id,
 	htlc_id,
 	amount_msat,
@@ -1911,6 +1935,7 @@ impl_writeable!(UpdateAddHTLC, {
 
 impl<W: Writer> Writeable<W> for OnionRealm0HopData {
 	fn write(&self, w: &mut W) -> Result<(), DecodeError> {
+		w.size_hint(32);
 		self.short_channel_id.write(w)?;
 		self.amt_to_forward.write(w)?;
 		self.outgoing_cltv_value.write(w)?;
@@ -1935,6 +1960,7 @@ impl<R: Read> Readable<R> for OnionRealm0HopData {
 
 impl<W: Writer> Writeable<W> for OnionHopData {
 	fn write(&self, w: &mut W) -> Result<(), DecodeError> {
+		w.size_hint(65);
 		self.realm.write(w)?;
 		self.data.write(w)?;
 		self.hmac.write(w)?;
@@ -1960,6 +1986,7 @@ impl<R: Read> Readable<R> for OnionHopData {
 
 impl<W: Writer> Writeable<W> for Ping {
 	fn write(&self, w: &mut W) -> Result<(), DecodeError> {
+		w.size_hint(self.byteslen as usize + 4);
 		self.ponglen.write(w)?;
 		vec![0u8; self.byteslen as usize].write(w)?; // size-unchecked write
 		Ok(())
@@ -1981,6 +2008,7 @@ impl<R: Read> Readable<R> for Ping {
 
 impl<W: Writer> Writeable<W> for Pong {
 	fn write(&self, w: &mut W) -> Result<(), DecodeError> {
+		w.size_hint(self.byteslen as usize + 2);
 		vec![0u8; self.byteslen as usize].write(w)?; // size-unchecked write
 		Ok(())
 	}
@@ -2000,6 +2028,7 @@ impl<R: Read> Readable<R> for Pong {
 
 impl<W: Writer> Writeable<W> for UnsignedChannelAnnouncement {
 	fn write(&self, w: &mut W) -> Result<(), DecodeError> {
+		w.size_hint(2 + 2*32 + 4*33 + self.features.flags.len() + self.excess_data.len());
 		self.features.write(w)?;
 		self.chain_hash.write(w)?;
 		self.short_channel_id.write(w)?;
@@ -2037,7 +2066,10 @@ impl<R: Read> Readable<R> for UnsignedChannelAnnouncement {
 	}
 }
 
-impl_writeable!(ChannelAnnouncement,{
+impl_writeable_len_match!(ChannelAnnouncement, {
+		{ ChannelAnnouncement { contents: UnsignedChannelAnnouncement {ref features, ref excess_data, ..}, .. },
+			2 + 2*32 + 4*33 + features.flags.len() + excess_data.len() + 4*64 }
+	}, {
 	node_signature_1,
 	node_signature_2,
 	bitcoin_signature_1,
@@ -2047,6 +2079,7 @@ impl_writeable!(ChannelAnnouncement,{
 
 impl<W: Writer> Writeable<W> for UnsignedChannelUpdate {
 	fn write(&self, w: &mut W) -> Result<(), DecodeError> {
+		w.size_hint(64 + self.excess_data.len());
 		self.chain_hash.write(w)?;
 		self.short_channel_id.write(w)?;
 		self.timestamp.write(w)?;
@@ -2080,15 +2113,20 @@ impl<R: Read> Readable<R> for UnsignedChannelUpdate {
 	}
 }
 
-impl_writeable!(ChannelUpdate, {
+impl_writeable_len_match!(ChannelUpdate, {
+		{ ChannelUpdate { contents: UnsignedChannelUpdate {ref excess_data, ..}, .. },
+			64 + excess_data.len() + 64 }
+	}, {
 	signature,
 	contents
 });
 
 impl<W: Writer> Writeable<W> for ErrorMessage {
 	fn write(&self, w: &mut W) -> Result<(), DecodeError> {
+		w.size_hint(32 + 2 + self.data.len());
 		self.channel_id.write(w)?;
-		self.data.as_bytes().to_vec().write(w)?; // write with size prefix
+		(self.data.len() as u16).write(w)?;
+		w.write_all(self.data.as_bytes())?;
 		Ok(())
 	}
 }
@@ -2113,6 +2151,7 @@ impl<R: Read> Readable<R> for ErrorMessage {
 
 impl<W: Writer> Writeable<W> for UnsignedNodeAnnouncement {
 	fn write(&self, w: &mut W) -> Result<(), DecodeError> {
+		w.size_hint(64 + 76 + self.features.flags.len() + self.addresses.len()*38 + self.excess_address_data.len() + self.excess_data.len());
 		self.features.write(w)?;
 		self.timestamp.write(w)?;
 		self.node_id.write(w)?;
@@ -2279,7 +2318,10 @@ impl<R: Read> Readable<R> for UnsignedNodeAnnouncement {
 	}
 }
 
-impl_writeable!(NodeAnnouncement, {
+impl_writeable_len_match!(NodeAnnouncement, {
+		{ NodeAnnouncement { contents: UnsignedNodeAnnouncement { ref features, ref addresses, ref excess_address_data, ref excess_data, ..}, .. },
+			64 + 76 + features.flags.len() + addresses.len()*38 + excess_address_data.len() + excess_data.len() }
+	}, {
 	signature,
 	contents
 });

src/util/ser.rs

@@ -14,17 +14,25 @@ use util::byte_utils::{be64_to_array, be32_to_array, be16_to_array, slice_to_be1
 
 const MAX_BUF_SIZE: usize = 64 * 1024;
 
-/// A trait that is similar to std::io::Write.
-/// An impl is provided for any type that also impls std::io::Write.
+/// A trait that is similar to std::io::Write but has one extra function which can be used to size
+/// buffers being written into.
+/// An impl is provided for any type that also impls std::io::Write which simply ignores size
+/// hints.
 pub trait Writer {
 	/// Writes the given buf out. See std::io::Write::write_all for more
 	fn write_all(&mut self, buf: &[u8]) -> Result<(), ::std::io::Error>;
+	/// Hints that data of the given size is about the be written. This may not always be called
+	/// prior to data being written and may be safely ignored.
+	fn size_hint(&mut self, size: usize);
 }
 
 impl<W: ::std::io::Write> Writer for W {
+	#[inline]
 	fn write_all(&mut self, buf: &[u8]) -> Result<(), ::std::io::Error> {
 		<Self as ::std::io::Write>::write_all(self, buf)
 	}
+	#[inline]
+	fn size_hint(&mut self, _size: usize) { }
 }
 
 /// A trait that various rust-lightning types implement allowing them to be written out to a Writer

src/util/ser_macros.rs

@@ -1,7 +1,29 @@
 macro_rules! impl_writeable {
-	($st:ident, {$($field:ident),*}) => {
+	($st:ident, $len: expr, {$($field:ident),*}) => {
 		impl<W: Writer> Writeable<W> for $st {
 			fn write(&self, w: &mut W) -> Result<(), DecodeError> {
+				w.size_hint($len);
+				$( self.$field.write(w)?; )*
+				Ok(())
+			}
+		}
+
+		impl<R: Read> Readable<R> for $st {
+			fn read(r: &mut R) -> Result<Self, DecodeError> {
+				Ok(Self {
+					$($field: Readable::read(r)?),*
+				})
+			}
+		}
+	}
+}
+macro_rules! impl_writeable_len_match {
+	($st:ident, {$({$m: pat, $l: expr}),*}, {$($field:ident),*}) => {
+		impl<W: Writer> Writeable<W> for $st {
+			fn write(&self, w: &mut W) -> Result<(), DecodeError> {
+				w.size_hint(match *self {
+					$($m => $l,)*
+				});
 				$( self.$field.write(w)?; )*
 				Ok(())
 			}