Impl `serialized_length()` without `LengthCalculatingWriter`
[rust-lightning] / lightning / src / util / ser_macros.rs
index 8729018935c944b10a24acefa24e52bf8debbf49..45435ecdd64ac5f5f579e168b011896831550709 100644 (file)
@@ -10,7 +10,7 @@
 macro_rules! encode_tlv {
        ($stream: expr, {$(($type: expr, $field: expr)),*}, {$(($optional_type: expr, $optional_field: expr)),*}) => { {
                #[allow(unused_imports)]
-               use util::ser::{BigSize, LengthCalculatingWriter};
+               use util::ser::BigSize;
                // Fields must be serialized in order, so we have to potentially switch between optional
                // fields and normal fields while serializing. Thus, we end up having to loop over the type
                // counts.
@@ -30,9 +30,7 @@ macro_rules! encode_tlv {
                        $(
                                if i == $type {
                                        BigSize($type).write($stream)?;
-                                       let mut len_calc = LengthCalculatingWriter(0);
-                                       $field.write(&mut len_calc)?;
-                                       BigSize(len_calc.0 as u64).write($stream)?;
+                                       BigSize($field.serialized_length() as u64).write($stream)?;
                                        $field.write($stream)?;
                                }
                        )*
@@ -40,9 +38,7 @@ macro_rules! encode_tlv {
                                if i == $optional_type {
                                        if let Some(ref field) = $optional_field {
                                                BigSize($optional_type).write($stream)?;
-                                               let mut len_calc = LengthCalculatingWriter(0);
-                                               field.write(&mut len_calc)?;
-                                               BigSize(len_calc.0 as u64).write($stream)?;
+                                               BigSize(field.serialized_length() as u64).write($stream)?;
                                                field.write($stream)?;
                                        }
                                }
@@ -51,31 +47,36 @@ macro_rules! encode_tlv {
        } }
 }
 
-macro_rules! encode_varint_length_prefixed_tlv {
-       ($stream: expr, {$(($type: expr, $field: expr)),*}, {$(($optional_type: expr, $optional_field: expr)),*}) => { {
-               use util::ser::{BigSize, LengthCalculatingWriter};
+macro_rules! get_varint_length_prefixed_tlv_length {
+       ({$(($type: expr, $field: expr)),*}, {$(($optional_type: expr, $optional_field: expr)),* $(,)*}) => { {
+               use util::ser::LengthCalculatingWriter;
                #[allow(unused_mut)]
                let mut len = LengthCalculatingWriter(0);
                {
                        $(
-                               BigSize($type).write(&mut len)?;
-                               let mut field_len = LengthCalculatingWriter(0);
-                               $field.write(&mut field_len)?;
-                               BigSize(field_len.0 as u64).write(&mut len)?;
-                               len.0 += field_len.0;
+                               BigSize($type).write(&mut len).expect("No in-memory data may fail to serialize");
+                               let field_len = $field.serialized_length();
+                               BigSize(field_len as u64).write(&mut len).expect("No in-memory data may fail to serialize");
+                               len.0 += field_len;
                        )*
                        $(
                                if let Some(ref field) = $optional_field {
-                                       BigSize($optional_type).write(&mut len)?;
-                                       let mut field_len = LengthCalculatingWriter(0);
-                                       field.write(&mut field_len)?;
-                                       BigSize(field_len.0 as u64).write(&mut len)?;
-                                       len.0 += field_len.0;
+                                       BigSize($optional_type).write(&mut len).expect("No in-memory data may fail to serialize");
+                                       let field_len = field.serialized_length();
+                                       BigSize(field_len as u64).write(&mut len).expect("No in-memory data may fail to serialize");
+                                       len.0 += field_len;
                                }
                        )*
                }
+               len.0
+       } }
+}
 
-               BigSize(len.0 as u64).write($stream)?;
+macro_rules! encode_varint_length_prefixed_tlv {
+       ($stream: expr, {$(($type: expr, $field: expr)),*}, {$(($optional_type: expr, $optional_field: expr)),*}) => { {
+               use util::ser::BigSize;
+               let len = get_varint_length_prefixed_tlv_length!({ $(($type, $field)),* }, { $(($optional_type, $optional_field)),* });
+               BigSize(len as u64).write($stream)?;
                encode_tlv!($stream, { $(($type, $field)),* }, { $(($optional_type, $optional_field)),* });
        } }
 }
@@ -115,8 +116,12 @@ macro_rules! decode_tlv {
                                _ => {},
                        }
                        // As we read types, make sure we hit every required type:
-                       $(if (last_seen_type.is_none() || last_seen_type.unwrap() < $reqtype) && typ.0 > $reqtype {
-                               Err(DecodeError::InvalidValue)?
+                       $({
+                               #[allow(unused_comparisons)] // Note that $reqtype may be 0 making the second comparison always true
+                               let invalid_order = (last_seen_type.is_none() || last_seen_type.unwrap() < $reqtype) && typ.0 > $reqtype;
+                               if invalid_order {
+                                       Err(DecodeError::InvalidValue)?
+                               }
                        })*
                        last_seen_type = Some(typ.0);
 
@@ -146,8 +151,12 @@ macro_rules! decode_tlv {
                        s.eat_remaining()?;
                }
                // Make sure we got to each required type after we've read every TLV:
-               $(if last_seen_type.is_none() || last_seen_type.unwrap() < $reqtype {
-                       Err(DecodeError::InvalidValue)?
+               $({
+                       #[allow(unused_comparisons)] // Note that $reqtype may be 0 making the second comparison always true
+                       let missing_req_type = last_seen_type.is_none() || last_seen_type.unwrap() < $reqtype;
+                       if missing_req_type {
+                               Err(DecodeError::InvalidValue)?
+                       }
                })*
        } }
 }
@@ -165,13 +174,29 @@ macro_rules! impl_writeable {
                                        if $len != 0 {
                                                use util::ser::LengthCalculatingWriter;
                                                let mut len_calc = LengthCalculatingWriter(0);
-                                               $( self.$field.write(&mut len_calc)?; )*
+                                               $( self.$field.write(&mut len_calc).expect("No in-memory data may fail to serialize"); )*
                                                assert_eq!(len_calc.0, $len);
+                                               assert_eq!(self.serialized_length(), $len);
                                        }
                                }
                                $( self.$field.write(w)?; )*
                                Ok(())
                        }
+
+                       #[inline]
+                       fn serialized_length(&self) -> usize {
+                               if $len == 0 || cfg!(any(test, feature = "fuzztarget")) {
+                                       let mut len_calc = 0;
+                                       $( len_calc += self.$field.serialized_length(); )*
+                                       if $len != 0 {
+                                               // In tests, assert that the hard-coded length matches the actual one
+                                               assert_eq!(len_calc, $len);
+                                       } else {
+                                               return len_calc;
+                                       }
+                               }
+                               $len
+                       }
                }
 
                impl ::util::ser::Readable for $st {
@@ -184,7 +209,7 @@ macro_rules! impl_writeable {
        }
 }
 macro_rules! impl_writeable_len_match {
-       ($struct: ident, $cmp: tt, {$({$match: pat, $length: expr}),*}, {$($field:ident),*}) => {
+       ($struct: ident, $cmp: tt, ($calc_len: expr), {$({$match: pat, $length: expr}),*}, {$($field:ident),*}) => {
                impl Writeable for $struct {
                        fn write<W: Writer>(&self, w: &mut W) -> Result<(), ::std::io::Error> {
                                let len = match *self {
@@ -196,12 +221,30 @@ macro_rules! impl_writeable_len_match {
                                        // In tests, assert that the hard-coded length matches the actual one
                                        use util::ser::LengthCalculatingWriter;
                                        let mut len_calc = LengthCalculatingWriter(0);
-                                       $( self.$field.write(&mut len_calc)?; )*
+                                       $( self.$field.write(&mut len_calc).expect("No in-memory data may fail to serialize"); )*
                                        assert!(len_calc.0 $cmp len);
+                                       assert_eq!(len_calc.0, self.serialized_length());
                                }
                                $( self.$field.write(w)?; )*
                                Ok(())
                        }
+
+                       #[inline]
+                       fn serialized_length(&self) -> usize {
+                               if $calc_len || cfg!(any(test, feature = "fuzztarget")) {
+                                       let mut len_calc = 0;
+                                       $( len_calc += self.$field.serialized_length(); )*
+                                       if !$calc_len {
+                                               assert_eq!(len_calc, match *self {
+                                                       $($match => $length,)*
+                                               });
+                                       }
+                                       return len_calc
+                               }
+                               match *self {
+                                       $($match => $length,)*
+                               }
+                       }
                }
 
                impl ::util::ser::Readable for $struct {
@@ -212,8 +255,11 @@ macro_rules! impl_writeable_len_match {
                        }
                }
        };
+       ($struct: ident, $cmp: tt, {$({$match: pat, $length: expr}),*}, {$($field:ident),*}) => {
+               impl_writeable_len_match!($struct, $cmp, (true), { $({ $match, $length }),* }, { $($field),* });
+       };
        ($struct: ident, {$({$match: pat, $length: expr}),*}, {$($field:ident),*}) => {
-               impl_writeable_len_match!($struct, ==, { $({ $match, $length }),* }, { $($field),* });
+               impl_writeable_len_match!($struct, ==, (false), { $({ $match, $length }),* }, { $($field),* });
        }
 }
 
@@ -246,7 +292,7 @@ macro_rules! write_ver_prefix {
 /// This is the preferred method of adding new fields that old nodes can ignore and still function
 /// correctly.
 macro_rules! write_tlv_fields {
-       ($stream: expr, {$(($type: expr, $field: expr)),*}, {$(($optional_type: expr, $optional_field: expr)),*}) => {
+       ($stream: expr, {$(($type: expr, $field: expr)),* $(,)*}, {$(($optional_type: expr, $optional_field: expr)),* $(,)*}) => {
                encode_varint_length_prefixed_tlv!($stream, {$(($type, $field)),*} , {$(($optional_type, $optional_field)),*});
        }
 }
@@ -267,7 +313,7 @@ macro_rules! read_ver_prefix {
 
 /// Reads a suffix added by write_tlv_fields.
 macro_rules! read_tlv_fields {
-       ($stream: expr, {$(($reqtype: expr, $reqfield: ident)),*}, {$(($type: expr, $field: ident)),*}) => { {
+       ($stream: expr, {$(($reqtype: expr, $reqfield: ident)),* $(,)*}, {$(($type: expr, $field: ident)),* $(,)*}) => { {
                let tlv_len = ::util::ser::BigSize::read($stream)?;
                let mut rd = ::util::ser::FixedLengthReader::new($stream, tlv_len.0);
                decode_tlv!(&mut rd, {$(($reqtype, $reqfield)),*}, {$(($type, $field)),*});
@@ -275,6 +321,124 @@ macro_rules! read_tlv_fields {
        } }
 }
 
+// If we naively create a struct in impl_writeable_tlv_based below, we may end up returning
+// `Self { ,,vecfield: vecfield }` which is obviously incorrect. Instead, we have to match here to
+// detect at least one empty field set and skip the potentially-extra comma.
+macro_rules! _init_tlv_based_struct {
+       ({}, {$($field: ident),*}, {$($vecfield: ident),*}) => {
+               Ok(Self {
+                       $($field),*,
+                       $($vecfield: $vecfield.unwrap().0),*
+               })
+       };
+       ({$($reqfield: ident),*}, {}, {$($vecfield: ident),*}) => {
+               Ok(Self {
+                       $($reqfield: $reqfield.0.unwrap()),*,
+                       $($vecfield: $vecfield.unwrap().0),*
+               })
+       };
+       ({$($reqfield: ident),*}, {$($field: ident),*}, {}) => {
+               Ok(Self {
+                       $($reqfield: $reqfield.0.unwrap()),*,
+                       $($field),*
+               })
+       };
+       ({$($reqfield: ident),*}, {$($field: ident),*}, {$($vecfield: ident),*}) => {
+               Ok(Self {
+                       $($reqfield: $reqfield.0.unwrap()),*,
+                       $($field),*,
+                       $($vecfield: $vecfield.unwrap().0),*
+               })
+       }
+}
+
+// If we don't have any optional types below, but do have some vec types, we end up calling
+// `write_tlv_field!($stream, {..}, {, (vec_ty, vec_val)})`, which is obviously broken.
+// Instead, for write and read we match the missing values and skip the extra comma.
+macro_rules! _write_tlv_fields {
+       ($stream: expr, {$(($type: expr, $field: expr)),* $(,)*}, {}, {$(($optional_type: expr, $optional_field: expr)),* $(,)*}) => {
+               write_tlv_fields!($stream, {$(($type, $field)),*} , {$(($optional_type, $optional_field)),*});
+       };
+       ($stream: expr, {$(($type: expr, $field: expr)),* $(,)*}, {$(($optional_type: expr, $optional_field: expr)),* $(,)*}, {$(($optional_type_2: expr, $optional_field_2: expr)),* $(,)*}) => {
+               write_tlv_fields!($stream, {$(($type, $field)),*} , {$(($optional_type, $optional_field)),*, $(($optional_type_2, $optional_field_2)),*});
+       }
+}
+macro_rules! _get_tlv_len {
+       ({$(($type: expr, $field: expr)),* $(,)*}, {}, {$(($optional_type: expr, $optional_field: expr)),* $(,)*}) => {
+               get_varint_length_prefixed_tlv_length!({$(($type, $field)),*} , {$(($optional_type, $optional_field)),*})
+       };
+       ({$(($type: expr, $field: expr)),* $(,)*}, {$(($optional_type: expr, $optional_field: expr)),* $(,)*}, {$(($optional_type_2: expr, $optional_field_2: expr)),* $(,)*}) => {
+               get_varint_length_prefixed_tlv_length!({$(($type, $field)),*} , {$(($optional_type, $optional_field)),*, $(($optional_type_2, $optional_field_2)),*})
+       }
+}
+macro_rules! _read_tlv_fields {
+       ($stream: expr, {$(($reqtype: expr, $reqfield: ident)),* $(,)*}, {}, {$(($type: expr, $field: ident)),* $(,)*}) => {
+               read_tlv_fields!($stream, {$(($reqtype, $reqfield)),*}, {$(($type, $field)),*});
+       };
+       ($stream: expr, {$(($reqtype: expr, $reqfield: ident)),* $(,)*}, {$(($type: expr, $field: ident)),* $(,)*}, {$(($type_2: expr, $field_2: ident)),* $(,)*}) => {
+               read_tlv_fields!($stream, {$(($reqtype, $reqfield)),*}, {$(($type, $field)),*, $(($type_2, $field_2)),*});
+       }
+}
+
+/// Implements Readable/Writeable for a struct storing it as a set of TLVs
+/// First block includes all the required fields including a dummy value which is used during
+/// deserialization but which will never be exposed to other code.
+/// The second block includes optional fields.
+/// The third block includes any Vecs which need to have their individual elements serialized.
+macro_rules! impl_writeable_tlv_based {
+       ($st: ident, {$(($reqtype: expr, $reqfield: ident)),* $(,)*}, {$(($type: expr, $field: ident)),* $(,)*}, {$(($vectype: expr, $vecfield: ident)),* $(,)*}) => {
+               impl ::util::ser::Writeable for $st {
+                       fn write<W: ::util::ser::Writer>(&self, writer: &mut W) -> Result<(), ::std::io::Error> {
+                               _write_tlv_fields!(writer, {
+                                       $(($reqtype, self.$reqfield)),*
+                               }, {
+                                       $(($type, self.$field)),*
+                               }, {
+                                       $(($vectype, Some(::util::ser::VecWriteWrapper(&self.$vecfield)))),*
+                               });
+                               Ok(())
+                       }
+
+                       #[inline]
+                       fn serialized_length(&self) -> usize {
+                               let len = _get_tlv_len!({
+                                       $(($reqtype, self.$reqfield)),*
+                               }, {
+                                       $(($type, self.$field)),*
+                               }, {
+                                       $(($vectype, Some(::util::ser::VecWriteWrapper(&self.$vecfield)))),*
+                               });
+                               use util::ser::{BigSize, LengthCalculatingWriter};
+                               let mut len_calc = LengthCalculatingWriter(0);
+                               BigSize(len as u64).write(&mut len_calc).expect("No in-memory data may fail to serialize");
+                               len + len_calc.0
+                       }
+               }
+
+               impl ::util::ser::Readable for $st {
+                       fn read<R: ::std::io::Read>(reader: &mut R) -> Result<Self, ::ln::msgs::DecodeError> {
+                               $(
+                                       let mut $reqfield = ::util::ser::OptionDeserWrapper(None);
+                               )*
+                               $(
+                                       let mut $field = None;
+                               )*
+                               $(
+                                       let mut $vecfield = Some(::util::ser::VecReadWrapper(Vec::new()));
+                               )*
+                               _read_tlv_fields!(reader, {
+                                       $(($reqtype, $reqfield)),*
+                               }, {
+                                       $(($type, $field)),*
+                               }, {
+                                       $(($vectype, $vecfield)),*
+                               });
+                               _init_tlv_based_struct!({$($reqfield),*}, {$($field),*}, {$($vecfield),*})
+                       }
+               }
+       }
+}
+
 #[cfg(test)]
 mod tests {
        use prelude::*;