Add size_hint in ser and call size_hint in all message serializers

diff --git a/src/ln/msgs.rs b/src/ln/msgs.rs
index a96f34f51a2c489c6c3d8cad094d7544494394d8..2c89acf2387686ac8c74a1197109ef07d1daad2f 100644 (file)
--- a/src/ln/msgs.rs
+++ b/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
 });

diff --git a/src/util/ser.rs b/src/util/ser.rs
index 6c40820905adf7291fb31e60debbbe22278a42f7..fe92332dfe93571636284687d345b6cb4f83f230 100644 (file)
--- a/src/util/ser.rs
+++ b/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

diff --git a/src/util/ser_macros.rs b/src/util/ser_macros.rs
index 2e1386a9c304e34f80beedd4b23a82ffac2c49da..162b18947a37986a0eb53b1c532830dbd297af0a 100644 (file)
--- a/src/util/ser_macros.rs
+++ b/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(())
                        }