Merge pull request #928 from TheBlueMatt/2021-05-really-tlv-ser

[rust-lightning] / lightning / src / util / ser.rs

diff --git a/lightning/src/util/ser.rs b/lightning/src/util/ser.rs
index c31fcaa1ff6de28eb450f770182414145cda1909..a52ba09c9c94100df8018ef3209f7469d46a5027 100644 (file)
--- a/lightning/src/util/ser.rs
+++ b/lightning/src/util/ser.rs
@@ -10,11 +10,12 @@
 //! A very simple serialization framework which is used to serialize/deserialize messages as well
 //! as ChannelsManagers and ChannelMonitors.
 
+use prelude::*;
 use std::io::{Read, Write};
 use std::collections::HashMap;
-use std::hash::Hash;
+use core::hash::Hash;
 use std::sync::Mutex;
-use std::cmp;
+use core::cmp;
 
 use bitcoin::secp256k1::Signature;
 use bitcoin::secp256k1::key::{PublicKey, SecretKey};
@@ -25,7 +26,7 @@ use bitcoin::consensus;
 use bitcoin::consensus::Encodable;
 use bitcoin::hashes::sha256d::Hash as Sha256dHash;
 use bitcoin::hash_types::{Txid, BlockHash};
-use std::marker::Sized;
+use core::marker::Sized;
 use ln::msgs::DecodeError;
 use ln::{PaymentPreimage, PaymentHash, PaymentSecret};
 use util::byte_utils;
@@ -222,6 +223,40 @@ pub trait MaybeReadable
        fn read<R: Read>(reader: &mut R) -> Result<Option<Self>, DecodeError>;
 }
 
+pub(crate) struct OptionDeserWrapper<T: Readable>(pub Option<T>);
+impl<T: Readable> Readable for OptionDeserWrapper<T> {
+       fn read<R: Read>(reader: &mut R) -> Result<Self, DecodeError> {
+               Ok(Self(Some(Readable::read(reader)?)))
+       }
+}
+
+const MAX_ALLOC_SIZE: u64 = 64*1024;
+
+pub(crate) struct VecWriteWrapper<'a, T: Writeable>(pub &'a Vec<T>);
+impl<'a, T: Writeable> Writeable for VecWriteWrapper<'a, T> {
+       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(())
+       }
+}
+pub(crate) struct VecReadWrapper<T: Readable>(pub Vec<T>);
+impl<T: Readable> Readable for VecReadWrapper<T> {
+       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) {
+                               Ok(v) => { values.push(v); },
+                               Err(e) => return Err(e),
+                       }
+               }
+               Ok(Self(values))
+       }
+}
+
 pub(crate) struct U48(pub u64);
 impl Writeable for U48 {
        #[inline]