rsh/src/message/binary.rs

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

macro_rules! try_from {
    (ref $into:ty, $from:expr $(; $fmt:literal)?) => {
	{
	    let from = &$from;
	    #[inline] fn _type_name_of_val<T: ?Sized>(_: &T) -> &'static str
	    {
		::std::any::type_name::<T>()
	    }
	    <$into>::try_from(*from).wrap_err(eyre!("Failed to convert type"))
		.with_section(|| ::std::any::type_name::<$into>().header("New type"))
		.with_section(|| _type_name_of_val(from).header("Old type"))
		.with_section(|| from.to_string().header("Value was"))
	}
    };
}

/// Check bit written/read from binary streams to check basic integrity.
pub const MESSAGE_HEADER_CHECK: [u8; 4] = (0xc0ffee00u32).to_be_bytes();

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();
		    #[inline] fn _type_name_of_val<T: ?Sized>(_: &T) -> &'static str
		    {
			::std::any::type_name::<T>()
		    }
		    let ser = $ser;
		    serde_cbor::to_writer(&mut v, $ser)
			.wrap_err(eyre!("Failed to serialise value to temporary buffer"))
			.with_section(|| _type_name_of_val(ser).header("Type was"))?;
		    buffer.put_u64(try_from!(ref u64, v.len())?);
		    write!(&v[..]);
		}
	    };
	}
	write!(MESSAGE_HEADER_CHECK);
	write!(: &self.header);
	buffer.put_u64(try_from!(ref u64, self.data.len())?);
	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 (should be 0 or 1)").with_section(|| x.header("Value was")));
			}
		    }
		}
	    };
	    (: $ser:ty) => {
		{
		    let len = try_from!(ref usize, bytes.get_u64())?;
		    if len > MAX_READ_SIZE {
			return Err(eyre!("Invalid length read: {}", len)
				   .with_section(|| MAX_READ_SIZE.header("Max length read")));
		    }
		    alloc_local_bytes(len, |de| {
			read!(&mut de[..]);
			serde_cbor::from_slice::<$ser>(&de[..]).wrap_err(eyre!("Failed to deserialise CBOR from reader")).with_section(|| std::any::type_name::<$ser>().header("Type to deserialise was"))
		    })?
			
		}
	    };
	    ($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 mut check = [0u8; MESSAGE_HEADER_CHECK.len()];
	read!(&mut check[..]);
	if check != MESSAGE_HEADER_CHECK {
	    return Err(eyre!("Invalid check bit for message header"))
		.with_section(|| u32::from_be_bytes(check).header("Expected"))
		.with_section(|| u32::from_be_bytes(MESSAGE_HEADER_CHECK).header("Got"));
	}
	
	let header = read!(: SerHeader);
	let data_len = try_from!(ref usize, bytes.get_u64())?;
	if MAX_BODY_SIZE > 0 && data_len > MAX_BODY_SIZE {
	    return Err(eyre!("Body size too large"))
		.with_section(|| data_len.header("Encoded size was"))
		.with_section(|| MAX_BODY_SIZE.header("Max size is"));
	}
	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 body bytes from buffer"))
		.with_section(|| format!("{} bytes", data_len).header("Tried to read"))
		.with_section(|| format!("{} bytes", data.len()).header("Read only"));
	}
	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)
    }
}