rsh/src/message/binary.rs

//! Creating binary from messages
//!
//! `SerializedMessage` to `Bytes`.
use super::*;
use bytes::{
    Bytes,
    BytesMut,
    BufMut,
    Buf,
};

impl<V: ?Sized> SerializedMessage<V>
{
    /// Write this message to a buffer
    ///
    /// # Panics
    /// If `buffer` cannot hold enough bytes.
    pub fn into_buffer(self, mut buffer: impl BufMut) -> eyre::Result<usize>
    {
	let mut w=0;
	macro_rules! write {
	    ($bytes:expr) => {
		{
		    let slice: &[u8] = ($bytes).as_ref();
		    buffer.put_slice(slice);
		    w+=slice.len();
		}
	    };
	    (? $o:expr) => {
		{
		    match $o {
			Some(opt) => {
			    buffer.put_u8(1);
			    write!(opt);
			},
			None => {buffer.put_u8(0);},
		    }
		    w+=1;
		}
	    };
	    (: $ser:expr) => {
		{		    
		    let mut v = StackVec::new();
		    serde_cbor::to_writer(&mut v, $ser)?;
		    buffer.put_u64(v.len().try_into()?);
		    write!(&v[..]);
		}
	    };
	}
	
	write!(: &self.header);
	buffer.put_u64(self.data.len().try_into()?);
	write!(self.data);
	write!(self.hash);
	write!(? self.enc_key);
	write!(? self.sig);
	
	Ok(w)
    }
    
    /// Write this message to a new `Bytes`.
    pub fn into_bytes(self) -> eyre::Result<Bytes>
    {
	let mut output = BytesMut::with_capacity(4096); //TODO: Find a better default capacity for this.
	self.into_buffer(&mut output)?;
	Ok(output.freeze())
    }
}

impl<V: ?Sized + MessageValue> SerializedMessage<V>
{
    /// Create from a buffer of bytes.
    ///
    /// # Panics
    /// If `bytes` does not contain enough data to read.
    pub fn from_buffer(mut bytes: impl Buf) -> eyre::Result<Self>
    {
	macro_rules! read {
	    ($bref:expr) => {
		{
		    let by: &mut [u8] = ($bref).as_mut();
		    bytes.copy_to_slice(by);
		}
	    };
	    (? $odef:expr) => {
		{
		    let by = bytes.get_u8();
		    match by {
			0 => None,
			1 => {
			    let mut def = $odef;
			    read!(&mut def);
			    Some(def)
			},
			x => {
			    return Err(eyre!("Invalid optional-set bit: {}", x));
			}
		    }
		}
	    };
	    (: $ser:ty) => {
		{
		    let len = usize::try_from(bytes.get_u64())?;
		    if len > MAX_ALLOC_SIZE {
			return Err(eyre!("Invalid length read: {}", len)
				   .with_section(|| format!("Max length read: {}", MAX_ALLOC_SIZE)))
		    }
		    alloc_local_bytes(len, |de| {
			read!(&mut de[..]);
			serde_cbor::from_slice::<$ser>(&de[..]).wrap_err(eyre!("Failed to deserialise {} from reader", std::any::type_name::<$ser>()))
		    })?
			
		}
	    };
	    ($into:expr, $num:expr) => {
		{
		    let num = $num;
		    let reader = (&mut bytes).reader();
		    copy_buffer($into, reader, num).wrap_err(eyre!("Failed to read {} bytes from reader", num))?
		}
	    }
	}

	let header = read!(: SerHeader);
	let data_len = usize::try_from(bytes.get_u64())?;
	let mut data = Vec::with_capacity(std::cmp::min(data_len, MAX_ALLOC_SIZE)); //XXX: Redesign so we don't allocate OR try to read massive buffers by accident on corrupted/malformed messages
	read!(&mut data, data_len);
	if data.len()!=data_len {
	    return Err(eyre!("Failed to read {} bytes from buffer (got {})", data_len, data.len()));
	}
	let mut hash = sha256::Sha256Hash::default();
	read!(&mut hash);
	let enc_key: Option<[u8; RSA_BLOCK_SIZE]> = read!(? [0u8; RSA_BLOCK_SIZE]);
	let sig: Option<rsa::Signature> = read!(? rsa::Signature::default());

	Ok(Self {
	    header,
	    data,
	    hash,
	    enc_key,
	    sig,

	    _phantom: PhantomData,
	})
    }

    /// Create from a slice of bytes
    #[inline] pub fn from_slice(bytes: impl AsRef<[u8]>) -> eyre::Result<Self>
    {
	Self::from_buffer(bytes.as_ref())
    }

    /// Create from a `Bytes` instance
    #[inline(always)] pub fn from_bytes(bytes: Bytes) -> eyre::Result<Self>
    {
	Self::from_buffer(bytes)
    }
}