serde_bolt/

types.rs

use crate::take::Take;
use alloc::boxed::Box;
use alloc::string::{String, ToString};
use alloc::vec::Vec;
use bitcoin::consensus::encode::MAX_VEC_SIZE;
use bitcoin::consensus::{encode::Error, Decodable, Encodable};
use bitcoin::io::{self, Read, Write, ErrorKind};
use chunked_buffer::{GenericChunkedBuffer, IterChunk};
use core::cmp::min;
use core::fmt::{Debug, Formatter};
use core::ops::{Deref, DerefMut};

/// Vec<u8> with nicer debug formatting and u32 big-endian size prefix.
/// Maximum 65535 bytes, for larger data use `LargeOctets`.
#[derive(Clone)]
#[cfg_attr(test, derive(PartialEq))]
pub struct Octets(pub Vec<u8>);

impl Octets {
    /// An empty vector
    pub const EMPTY: Self = Octets(Vec::new());
}

impl Debug for Octets {
    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
        f.write_str(&hex::encode(&self.0))
    }
}

impl Deref for Octets {
    type Target = Vec<u8>;
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl DerefMut for Octets {
    fn deref_mut(&mut self) -> &mut Vec<u8> {
        &mut self.0
    }
}

impl From<Vec<u8>> for Octets {
    fn from(v: Vec<u8>) -> Self {
        Self(v)
    }
}

impl Encodable for Octets {
    fn consensus_encode<W: Write + ?Sized>(&self, writer: &mut W) -> Result<usize, io::Error> {
        let len = self.0.len();
        if len > 0xFFFF {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                "Octets length exceeds 65535",
            ));
        }
        let mut count = 0;
        count += (len as u16).to_be_bytes().consensus_encode(writer)?;
        writer.write_all(&self.0)?;
        count += len;
        Ok(count)
    }
}

impl Decodable for Octets {
    fn consensus_decode<R: Read + ?Sized>(reader: &mut R) -> Result<Self, Error> {
        let len = reader.read_u16_be()?;
        let mut buf = Vec::with_capacity(len as usize);
        buf.resize(len as usize, 0);
        reader.read_exact(&mut buf)?;
        Ok(Octets(buf))
    }
}

/// A Vec that implements `Encodable` and `Decodable` as a length-prefixed array,
/// with a big-endian u16 length prefix.
#[derive(Clone, Debug)]
#[cfg_attr(feature = "test_utils", derive(PartialEq))]
pub struct Array<T: Encodable + Decodable + Debug>(pub Vec<T>);

impl<T: Encodable + Decodable + Debug> Array<T> {
    /// An empty vector
    pub fn new() -> Self {
        Self(Vec::new())
    }
}

impl<T: Encodable + Decodable + Debug> Deref for Array<T> {
    type Target = Vec<T>;
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl<T: Encodable + Decodable + Debug> DerefMut for Array<T> {
    fn deref_mut(&mut self) -> &mut Vec<T> {
        &mut self.0
    }
}

impl<T: Encodable + Decodable + Debug> Encodable for Array<T> {
    fn consensus_encode<W: Write + ?Sized>(&self, writer: &mut W) -> Result<usize, io::Error> {
        let mut count = 0;
        count += (self.0.len() as u16)
            .to_be_bytes()
            .consensus_encode(writer)?;
        for item in &self.0 {
            count += item.consensus_encode(writer)?;
        }
        Ok(count)
    }
}

impl<T: Encodable + Decodable + Debug> Decodable for Array<T> {
    fn consensus_decode<R: Read + ?Sized>(reader: &mut R) -> Result<Self, Error> {
        let len = reader.read_u16_be()?;
        let mut buf = Vec::with_capacity(len as usize);
        for _ in 0..len {
            buf.push(Decodable::consensus_decode(reader)?);
        }
        Ok(Array(buf))
    }
}

impl<T: Encodable + Decodable + Debug> From<Vec<T>> for Array<T> {
    fn from(v: Vec<T>) -> Self {
        Self(v)
    }
}

/// A Vec that implements `Encodable` and `Decodable` as a length-prefixed array,
/// with a big-endian u16 length prefix.
/// Unlike `Array`, this type requires `T` to implement `BigEndianEncodable`.
/// Mostly useful for things like `ArrayBE<u16>`, where we want the elements
/// to be encoded as big-endian.
#[derive(Clone, Debug)]
pub struct ArrayBE<T: BigEndianEncodable + Debug>(pub Vec<T>);

impl<T: BigEndianEncodable + Debug> ArrayBE<T> {
    /// An empty vector
    pub fn new() -> Self {
        Self(Vec::new())
    }
}

impl<T: BigEndianEncodable + Debug> Deref for ArrayBE<T> {
    type Target = Vec<T>;
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl<T: BigEndianEncodable + Debug> DerefMut for ArrayBE<T> {
    fn deref_mut(&mut self) -> &mut Vec<T> {
        &mut self.0
    }
}

impl<T: BigEndianEncodable + Debug> Encodable for ArrayBE<T> {
    fn consensus_encode<W: Write + ?Sized>(&self, writer: &mut W) -> Result<usize, io::Error> {
        let mut count = 0;
        count += (self.0.len() as u16)
            .to_be_bytes()
            .consensus_encode(writer)?;
        for item in &self.0 {
            count += item.consensus_encode_be(writer).map_err(|e| io::Error::new(ErrorKind::Other, e.to_string()))?;
        }
        Ok(count)
    }
}

impl<T: BigEndianEncodable + Debug> Decodable for ArrayBE<T> {
    fn consensus_decode<R: Read + ?Sized>(reader: &mut R) -> Result<Self, Error> {
        let len = reader.read_u16_be()?;
        let mut buf = Vec::with_capacity(len as usize);
        for _ in 0..len {
            buf.push(BigEndianEncodable::consensus_decode_be(reader)?);
        }
        Ok(ArrayBE(buf))
    }
}

impl<T: BigEndianEncodable + Debug> From<Vec<T>> for ArrayBE<T> {
    fn from(v: Vec<T>) -> Self {
        Self(v)
    }
}

/// A potentially large vector of bytes, with a u32 big-endian size.
///
/// Not used in BOLT-1, because messages are limited to 64 KB.
#[derive(Clone)]
#[cfg_attr(test, derive(PartialEq))]
pub struct LargeOctets(pub Vec<u8>);

impl LargeOctets {
    /// An empty vector
    pub const EMPTY: Self = LargeOctets(Vec::new());
}

impl Debug for LargeOctets {
    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
        f.write_str(&hex::encode(&self.0))
    }
}

impl Deref for LargeOctets {
    type Target = Vec<u8>;
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl DerefMut for LargeOctets {
    fn deref_mut(&mut self) -> &mut Vec<u8> {
        &mut self.0
    }
}

impl From<Vec<u8>> for LargeOctets {
    fn from(v: Vec<u8>) -> Self {
        Self(v)
    }
}

impl Encodable for LargeOctets {
    fn consensus_encode<W: Write + ?Sized>(&self, writer: &mut W) -> Result<usize, io::Error> {
        let len = self.0.len();
        let mut count = 0;
        count += (len as u32).to_be_bytes().consensus_encode(writer)?;
        writer.write_all(&self.0)?;
        count += len;
        Ok(count)
    }
}

impl Decodable for LargeOctets {
    fn consensus_decode<R: Read + ?Sized>(reader: &mut R) -> Result<Self, Error> {
        let len = reader.read_u32_be()?;
        if len as usize > MAX_VEC_SIZE {
            return Err(Error::OversizedVectorAllocation {
                requested: len as usize,
                max: MAX_VEC_SIZE,
            });
        }
        let mut buf = Vec::with_capacity(len as usize);
        buf.resize(len as usize, 0);
        reader.read_exact(&mut buf)?;
        Ok(LargeOctets(buf))
    }
}

/// A potentially large vector of bytes, with a u32 size, that we ignore on deserialize.
/// On deserialize, the bytes are read and discarded and the inner will be empty.
#[derive(Clone)]
#[cfg_attr(test, derive(PartialEq))]
pub struct IgnoredLargeOctets(pub Vec<u8>);

impl Debug for IgnoredLargeOctets {
    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
        f.write_str(&hex::encode(&self.0))
    }
}

impl Encodable for IgnoredLargeOctets {
    fn consensus_encode<W: Write + ?Sized>(&self, writer: &mut W) -> Result<usize, io::Error> {
        let len = self.0.len();
        let mut count = 0;
        count += (len as u32).to_be_bytes().consensus_encode(writer)?;
        writer.write_all(&self.0)?;
        count += len;
        Ok(count)
    }
}

impl Decodable for IgnoredLargeOctets {
    fn consensus_decode<R: Read + ?Sized>(reader: &mut R) -> Result<Self, Error> {
        let mut len = reader.read_u32_be()?;
        while len > 0 {
            let mut buf = [0; 1024];
            let read = reader.read(&mut buf[..min(len as usize, 1024)])?;
            len -= read as u32;
        }
        Ok(IgnoredLargeOctets(Vec::new()))
    }
}

/// A variable length zero terminated byte string
#[derive(Clone)]
#[cfg_attr(any(test, feature = "test_utils"), derive(PartialEq, Eq))]
pub struct WireString(pub Vec<u8>);

impl Debug for WireString {
    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
        match String::from_utf8(self.0.clone()) {
            Ok(str) => write!(f, "\"{}\"", str),             // utf8
            Err(_) => write!(f, "{}", hex::encode(&self.0)), // non-uf8
        }
    }
}

impl Encodable for WireString {
    fn consensus_encode<W: Write + ?Sized>(&self, writer: &mut W) -> Result<usize, io::Error> {
        assert!(!self.0.contains(&0), "WireString cannot contain 0");
        let mut count = 0;
        writer.write_all(&self.0)?;
        count += self.0.len();
        writer.write_all(&[0])?;
        count += 1;
        Ok(count)
    }
}

impl Decodable for WireString {
    fn consensus_decode<R: Read + ?Sized>(reader: &mut R) -> Result<Self, Error> {
        let mut buf = Vec::new();
        loop {
            let mut byte = [0; 1];
            reader.read_exact(&mut byte)?;
            if byte[0] == 0 {
                break;
            }
            buf.push(byte[0]);
        }
        Ok(WireString(buf))
    }
}

/// A wrapper around a type that implements `Encodable` and `Decodable` that
/// prefixes the encoded bytes with the length of the encoded bytes as big-endian
/// u32.
#[derive(Clone, Debug)]
pub struct WithSize<T: Encodable + Decodable + Debug>(pub T);

impl<T: Encodable + Decodable + Debug> Encodable for WithSize<T> {
    fn consensus_encode<W: Write + ?Sized>(&self, writer: &mut W) -> Result<usize, io::Error> {
        // compute the size of the encoded bytes by streaming to a sink
        let mut sink = Sink;
        let size = self.0.consensus_encode(&mut sink)? as u32;
        if size > MAX_VEC_SIZE as u32 {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                "Octets length exceeds MAX_VEC_SIZE",
            ));
        }
        let mut count = 0;
        count += size.to_be_bytes().consensus_encode(writer)?;
        count += self.0.consensus_encode(writer)?;
        Ok(count)
    }
}

impl<T: Encodable + Decodable + Debug> Decodable for WithSize<T> {
    fn consensus_decode<R: Read + ?Sized>(reader: &mut R) -> Result<Self, Error> {
        let len = reader.read_u32_be()?;
        if len as usize > MAX_VEC_SIZE {
            return Err(Error::OversizedVectorAllocation {
                requested: len as usize,
                max: MAX_VEC_SIZE,
            });
        }

        let mut take = Take::new(Box::new(reader), len as u64);
        let inner = T::consensus_decode(&mut take)?;
        if !take.is_empty() {
            return Err(Error::ParseFailed("trailing bytes in WithSize"));
        }
        Ok(Self(inner))
    }
}

impl<T: Encodable + Decodable + Debug> From<T> for WithSize<T> {
    fn from(value: T) -> WithSize<T> {
        WithSize(value)
    }
}

impl<T: Encodable + Decodable + Debug> Deref for WithSize<T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl<T: Encodable + Decodable + Debug> DerefMut for WithSize<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

/// Particularly useful for memory constrained environments, this structure avoids
/// large contiguous memory allocations, and incrementally releases memory as it is consumed.
/// See GenericChunkedBuffer for details.
pub struct NonContiguousOctets<const CHUNK_SIZE: usize>(GenericChunkedBuffer<CHUNK_SIZE>);

impl<const CHUNK_SIZE: usize> NonContiguousOctets<CHUNK_SIZE> {
    /// Creates a new instance of this type, with an empty chunked buffer wrapped inside
    pub fn new() -> Self {
        NonContiguousOctets(GenericChunkedBuffer::<CHUNK_SIZE>::new())
    }
}

impl<const CHUNK_SIZE: usize> Read for NonContiguousOctets<CHUNK_SIZE> {
    fn read(&mut self, buf: &mut [u8]) -> Result<usize, io::Error> {
        Ok(self.0.read(buf))
    }
}

impl<const CHUNK_SIZE: usize> Write for NonContiguousOctets<CHUNK_SIZE> {
    fn write(&mut self, buf: &[u8]) -> Result<usize, io::Error> {
        self.0.write(buf);
        Ok(buf.len())
    }
    fn flush(&mut self) -> Result<(), io::Error> {
        Ok(())
    }
}

impl<const CHUNK_SIZE: usize> Encodable for NonContiguousOctets<CHUNK_SIZE> {
    fn consensus_encode<W: Write + ?Sized>(&self, writer: &mut W) -> Result<usize, io::Error> {
        let len = self.0.len();
        let mut count = 0;
        count += (len as u32).to_be_bytes().consensus_encode(writer)?;
        for slice in self.0.iter_chunks() {
            writer.write_all(slice)?;
        }
        count += len;
        Ok(count)
    }
}

impl<const CHUNK_SIZE: usize> Decodable for NonContiguousOctets<CHUNK_SIZE> {
    fn consensus_decode<R: Read + ?Sized>(reader: &mut R) -> Result<Self, Error> {
        let len = reader.read_u32_be()?;
        if len as usize > MAX_VEC_SIZE {
            return Err(Error::OversizedVectorAllocation {
                requested: len as usize,
                max: MAX_VEC_SIZE,
            });
        }
        let mut reader = reader.take(len as u64);
        let mut chunk = [0u8; CHUNK_SIZE];
        let mut buf = GenericChunkedBuffer::<CHUNK_SIZE>::new();
        let mut nread = 0;
        while nread < len as usize {
            let n = reader.read(&mut chunk)?;
            if n == 0 {
                return Err(Error::Io(ErrorKind::UnexpectedEof.into()));
            }
            buf.write(&chunk[..n]);
            nread += n;
        }
        Ok(NonContiguousOctets(buf))
    }
}

impl<const CHUNK_SIZE: usize> Deref for NonContiguousOctets<CHUNK_SIZE> {
    type Target = GenericChunkedBuffer<CHUNK_SIZE>;
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl<const CHUNK_SIZE: usize> DerefMut for NonContiguousOctets<CHUNK_SIZE> {
    fn deref_mut(&mut self) -> &mut GenericChunkedBuffer<CHUNK_SIZE> {
        &mut self.0
    }
}

impl<const CHUNK_SIZE: usize> Default for NonContiguousOctets<CHUNK_SIZE> {
    fn default() -> Self {
        NonContiguousOctets::<CHUNK_SIZE>::new()
    }
}

impl<const CHUNK_SIZE: usize> From<GenericChunkedBuffer<CHUNK_SIZE>>
    for NonContiguousOctets<CHUNK_SIZE>
{
    fn from(v: GenericChunkedBuffer<CHUNK_SIZE>) -> Self {
        Self(v)
    }
}

/// A reader that reads the bytes of a NonContiguousOctets without mutating it
pub struct NonContiguousOctetsCursor<'a, const CHUNK_SIZE: usize> {
    current_slice: Option<&'a [u8]>,
    chunks_iterator: IterChunk<'a, CHUNK_SIZE>,
}

impl<'a, const CHUNK_SIZE: usize> NonContiguousOctetsCursor<'a, CHUNK_SIZE> {
    /// Creates a new instance of this cursor
    pub fn new(buf: &'a NonContiguousOctets<CHUNK_SIZE>) -> Self {
        let mut chunks_iterator = buf.iter_chunks();
        Self {
            current_slice: chunks_iterator.next(),
            chunks_iterator,
        }
    }
}

impl<'a, const CHUNK_SIZE: usize> Read for NonContiguousOctetsCursor<'a, CHUNK_SIZE> {
    fn read(&mut self, buf: &mut [u8]) -> Result<usize, io::Error> {
        let mut n = usize::MAX;
        let mut nread = 0;
        while self.current_slice.is_some() && n > 0 {
            let slice = self.current_slice.as_mut().unwrap();
            n = slice.read(&mut buf[nread..])?;
            nread += n;
            if slice.is_empty() {
                self.current_slice = self.chunks_iterator.next();
            }
        }
        Ok(nread)
    }
}

/// A trait for reading big-endian integers from a stream
pub trait ReadBigEndian {
    /// Read a big-endian u16
    fn read_u16_be(&mut self) -> Result<u16, Error>;
    /// Read a big-endian u32
    fn read_u32_be(&mut self) -> Result<u32, Error>;
    /// Read a big-endian u64
    fn read_u64_be(&mut self) -> Result<u64, Error>;
}

impl<R: Read + ?Sized> ReadBigEndian for R {
    fn read_u16_be(&mut self) -> Result<u16, Error> {
        let mut buf = [0; 2];
        self.read_exact(&mut buf)?;
        Ok(u16::from_be_bytes(buf))
    }

    fn read_u32_be(&mut self) -> Result<u32, Error> {
        let mut buf = [0; 4];
        self.read_exact(&mut buf)?;
        Ok(u32::from_be_bytes(buf))
    }

    fn read_u64_be(&mut self) -> Result<u64, Error> {
        let mut buf = [0; 8];
        self.read_exact(&mut buf)?;
        Ok(u64::from_be_bytes(buf))
    }
}

/// A sink that discards all bytes written to it.
/// core2 is missing this.
pub struct Sink;

impl Write for Sink {
    fn write(&mut self, buf: &[u8]) -> Result<usize, io::Error> {
        Ok(buf.len())
    }

    fn flush(&mut self) -> Result<(), io::Error> {
        Ok(())
    }
}

/// A trait for types that are encoded with big endian byte order.
///
/// This is used for things like `Array<u32>`, so that the elements are encoded as big endian
pub trait BigEndianEncodable: Encodable + Decodable {
    /// Encode the object with Big Endian byte order
    fn consensus_encode_be<W: Write + ?Sized>(&self, writer: &mut W) -> Result<usize, Error>;

    /// Decode the object with Big Endian byte order
    fn consensus_decode_be<R: Read + ?Sized>(reader: &mut R) -> Result<Self, Error>;
}

impl BigEndianEncodable for u8 {
    fn consensus_encode_be<W: Write + ?Sized>(&self, writer: &mut W) -> Result<usize, Error> {
        writer.write_all(&[*self])?;
        Ok(1)
    }

    fn consensus_decode_be<R: Read + ?Sized>(reader: &mut R) -> Result<Self, Error> {
        let mut buf = [0; 1];
        reader.read_exact(&mut buf)?;
        Ok(buf[0])
    }
}

impl BigEndianEncodable for u16 {
    fn consensus_encode_be<W: Write + ?Sized>(&self, writer: &mut W) -> Result<usize, Error> {
        writer.write_all(&self.to_be_bytes())?;
        Ok(2)
    }

    fn consensus_decode_be<R: Read + ?Sized>(reader: &mut R) -> Result<Self, Error> {
        let mut buf = [0; 2];
        reader.read_exact(&mut buf)?;
        Ok(u16::from_be_bytes(buf))
    }
}

impl BigEndianEncodable for u32 {
    fn consensus_encode_be<W: Write + ?Sized>(&self, writer: &mut W) -> Result<usize, Error> {
        writer.write_all(&self.to_be_bytes())?;
        Ok(4)
    }

    fn consensus_decode_be<R: Read + ?Sized>(reader: &mut R) -> Result<Self, Error> {
        let mut buf = [0; 4];
        reader.read_exact(&mut buf)?;
        Ok(u32::from_be_bytes(buf))
    }
}

impl BigEndianEncodable for u64 {
    fn consensus_encode_be<W: Write + ?Sized>(&self, writer: &mut W) -> Result<usize, Error> {
        writer.write_all(&self.to_be_bytes())?;
        Ok(8)
    }

    fn consensus_decode_be<R: Read + ?Sized>(reader: &mut R) -> Result<Self, Error> {
        let mut buf = [0; 8];
        reader.read_exact(&mut buf)?;
        Ok(u64::from_be_bytes(buf))
    }
}