From cf430029c48940e9bf428d1336072573f8397368 Mon Sep 17 00:00:00 2001 From: Matt Corallo Date: Tue, 25 May 2021 21:18:30 +0000 Subject: [PATCH] Add a macro which implements Readable/Writeable using TLVs only This also includes a `VecWriteWrapper` and `VecReadWrapper` which implements serialization for any `Readable`/`Writeable` type that is in a Vec. We do this instead of implementing `Readable`/`Writeable` directly as there isn't always a univerally-defined way to serialize a Vec and this makes things more explicit. Finally, this tweaks existing macros (and in the new macros) to support a trailing `,` after a list, eg `write_tlv_fields!(stream, {(0, a),}, {});` whereas previously the trailing `,` after the `(0, a)` would be a compile-error. --- lightning/src/util/ser.rs | 34 +++++++++++ lightning/src/util/ser_macros.rs | 99 +++++++++++++++++++++++++++++++- 2 files changed, 131 insertions(+), 2 deletions(-) diff --git a/lightning/src/util/ser.rs b/lightning/src/util/ser.rs index 66c89aa83..6a76cdfae 100644 --- a/lightning/src/util/ser.rs +++ b/lightning/src/util/ser.rs @@ -222,6 +222,40 @@ pub trait MaybeReadable fn read(reader: &mut R) -> Result, DecodeError>; } +pub(crate) struct OptionDeserWrapper(pub Option); +impl Readable for OptionDeserWrapper { + fn read(reader: &mut R) -> Result { + Ok(Self(Some(Readable::read(reader)?))) + } +} + +const MAX_ALLOC_SIZE: u64 = 64*1024; + +pub(crate) struct VecWriteWrapper<'a, T: Writeable>(pub &'a Vec); +impl<'a, T: Writeable> Writeable for VecWriteWrapper<'a, T> { + fn write(&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(()) + } +} +pub(crate) struct VecReadWrapper(pub Vec); +impl Readable for VecReadWrapper { + fn read(reader: &mut R) -> Result { + let count: u64 = Readable::read(reader)?; + let mut values = Vec::with_capacity(cmp::min(count, MAX_ALLOC_SIZE / (core::mem::size_of::() as u64)) as usize); + for _ in 0..count { + match Readable::read(reader) { + Ok(v) => { values.push(v); }, + Err(e) => return Err(e), + } + } + Ok(Self(values)) + } +} + pub(crate) struct U48(pub u64); impl Writeable for U48 { #[inline] diff --git a/lightning/src/util/ser_macros.rs b/lightning/src/util/ser_macros.rs index 30fa60f97..db970df5f 100644 --- a/lightning/src/util/ser_macros.rs +++ b/lightning/src/util/ser_macros.rs @@ -254,7 +254,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)),*}); } } @@ -275,7 +275,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)),*}); @@ -283,6 +283,101 @@ 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! _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(&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(()) + } + } + + impl ::util::ser::Readable for $st { + fn read(reader: &mut R) -> Result { + $( + 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 std::io::{Cursor, Read}; -- 2.39.5