rsh/src/message.rs

//! Messages
use super::*;
use std::marker::PhantomData;
use cryptohelpers::{
    sha256,
    aes,
    rsa,
};
use uuid::Uuid;
use std::borrow::{
    Borrow,
    Cow
};
use std::io;
use std::marker::Unpin;
use tokio::io::{
    AsyncWrite,
    AsyncRead,
};

mod serial;
pub use serial::*;

mod builder;
pub use builder::*;

pub mod value;
pub use value::MessageValue;

/// A `SerializedMessage` whos type has been erased.
pub type UntypedSerializedMessage = SerializedMessage<value::UntypedMessageValue>;

/// Size of encrypted AES key
pub const RSA_BLOCK_SIZE: usize = 512;

/// Max size to pre-allocate when reading a message buffer.
pub const MAX_ALLOC_SIZE: usize = 4096; // 4kb

/// Max size to allow reading for a message buffer component.
///
/// Not including the message body, see `MAX_BODY_SIZE` for body.
pub const MAX_READ_SIZE: usize = 2048 * 1024; // 2MB.

/// Max allowed size of a single message body.
///
/// Set to 0 for unlimited.
pub const MAX_BODY_SIZE: usize = (1024 * 1024) * 1024; // 1GB

/// A message that can send a value into bytes.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct Message<V: ?Sized + MessageValue>
{
    header: SerHeader,
    /// Optional key to use to encrypt the message
    key: Option<aes::AesKey>,
    /// Should the message body be signed?
    sign: bool,

    /// Value to serialise
    value: V,
}

// Accessors
impl<V: ?Sized + MessageValue> Message<V>
{
    /// Get a reference to this message's header
    #[inline(always)] pub fn header(&self) -> MessageHeader<'_, V>
    {
	MessageHeader(&self.header, PhantomData)
    }
}
macro_rules! accessor {
    ($name:ident, $fn_name:ident, $type:ty $(; $comment:literal)?) => {
	$(#[doc = $comment])?
	    #[inline] pub fn $fn_name(&self) -> &$type
	{
	    &self.0.$name
	}
    };
    ($name:ident, $type:ty $(; $comment:literal)?) => (accessor!($name, $name, $type $(; $comment)?););
}
impl<'a, V: ?Sized> MessageHeader<'a, V>
{    
    accessor!(id, Uuid; "This message's randomly generated ID");
    accessor!(idemp, idempotence, Uuid; "This message's idempotence key");
    accessor!(timestamp, unix_timestamp_utc, u64; "This message's UTX unix timestamp");
    accessor!(responds_to, Option<Uuid>; "The ID this message is responding to (if any)");
}

/// `SerializedMessage` header.
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
struct SerHeader
{    
    /// Message ID
    id: Uuid,
    /// Message idempodence ID
    idemp: Uuid,
    /// Timestamp of when this message was created (Unix TS, UTC).
    timestamp: u64,
    /// `id` of message this one is responding to, if needed.
    responds_to: Option<Uuid>,
    //TODO: Add `flags` bitflags
    //TODO: Add `kind` enum
}

/// A reference to a message's header.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct MessageHeader<'a, V:?Sized>(&'a SerHeader, PhantomData<&'a V>);

impl<V: ?Sized + MessageValue> AsRef<V> for Message<V>
{
    #[inline(always)] fn as_ref(&self) -> &V
    {
	self.value_ref()
    }
}

/// A serialized message that can be sent over a socket.
///
/// Messages of this type are not yet validated, and may be invalid/corrupt. The validation happens when converting back to a `Message<V>` (of the same `V`.)
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SerializedMessage<V: ?Sized>
{
    header: SerHeader,
    /// cbor serialised `V`.
    data: Vec<u8>,
    /// Hash of `data` (after encryption)
    hash: sha256::Sha256Hash,

    /// `key` encrypted with recipient's RSA public key.
    enc_key: Option<[u8; RSA_BLOCK_SIZE]>, // we can't derive Serialize because of this array.. meh..
    /// Signature of hash of un-encrypted `data`.
    sig: Option<rsa::Signature>,

    //TODO: Add a message header checksum.
    _phantom: PhantomData<V>,
}

impl<V: ?Sized + MessageValue> Message<V>
{
    /// A reference to the value itself.
    pub fn value_ref(&self) -> &V
    {
	&self.value
    }
    
    /// A mutable reference to the value itself.
    ///
    /// # Safety
    /// Mutating a value inside a message may cause invalid metadata.
    pub unsafe fn value_mut(&mut self) -> &mut V
    {
	&mut self.value
    }

    /// Consume into just the value
    pub fn into_value(self) -> V
    {
	self.value
    }
    
    /// Serialise this message into one that can be converted to/from bytes.
    ///
    /// # Panics
    ///  * If this message was specified to be encrypted, but `S` doesn't support encryption.
    ///  * If this message was specified to be signed, but `S` doesn't support signing.
    pub fn serialise<S: ?Sized + MessageSender>(&self, send_with: impl Borrow<S>) -> eyre::Result<SerializedMessage<V>>
    {
	let send_with: &S = send_with.borrow();
	let data = serde_cbor::to_vec(&self.value)?;
	if MAX_BODY_SIZE > 0 && data.len() > MAX_BODY_SIZE {
	    return Err(eyre!("Encoded body is too large"))
		.with_section(|| data.len().header("Body size was"))
		.with_section(|| MAX_BODY_SIZE.header("Max size is"));
	}
	let sig = if self.sign {
	    Some(send_with.sign_data(&data[..]).expect("Message expected signing, sender did not support it"))
	} else {
	    None
	};
	
	let (data, enc_key) = if let Some(key) = &self.key {
	    // Encrypt the body
	    let enc_key = send_with.encrypt_key(key).expect("Message expected encryption, sender did not support it");
	    (aes::encrypt_slice_sync(key, data)?, Some(enc_key))
	} else {
	    // Don't encrypt the body
	    (data, None)
	};
	// Compute hash of data
	let hash = sha256::compute_slice(&data);
	
	Ok(SerializedMessage{
	    header: self.header.clone(),
	    data,
	    sig,
	    enc_key,
	    hash,
	    
	    _phantom: PhantomData,
	})
    }

    /// Try to deserialize and validate a received message.
    ///
    /// If a part of the message is invalid, an error is returned.
    pub fn deserialize<'a, S: ?Sized + MessageReceiver>(serial: &'a SerializedMessage<V>, recv_with: impl Borrow<S>) -> eyre::Result<Self>
    {
	let recv_with: &S = recv_with.borrow();
	macro_rules! assert_valid {
	    ($ex:expr, $fmt:literal $($tt:tt)*) => {
		if $ex {
		    Ok(())
		} else {
		    Err(eyre!($fmt $($tt)*))
		}
	    }
	}
	//Validate hashes
	assert_valid!(sha256::compute_slice(&serial.data) == serial.hash, "Non-matching hashes")?;

	// Decrypted data
	let (data, key): (Cow<'a, [u8]>, Option<aes::AesKey>) = if let Some(enc_key) = &serial.enc_key {
	    let key = recv_with.decrypt_key(enc_key).ok_or(eyre!("Message was decrypted, but receiver doesn't support decryption"))?
		.wrap_err(eyre!("Failed to decrypt session key"))?;
	    let mut data = Vec::with_capacity(serial.data.len());
	    aes::decrypt_stream_sync(&key, &mut &serial.data[..], &mut data).wrap_err(eyre!("Failed to decrypt body"))?;
	    (Cow::Owned(data), Some(key))
	} else {
	    (Cow::Borrowed(&serial.data[..]), None)
	};

	let sign = if let Some(sig) = &serial.sig {
	    // Validate signature
	    assert_valid!(recv_with.verify_data(&data[..], sig).ok_or(eyre!("Message was signed, but receiver doesn't support signature verification"))?.wrap_err("Failed to verify signature")?,
			  "Non-matching signature")
		.with_section(move || format!("Embedded sig was: {}", sig))?;
	    true
	} else {
	    false
	};

	// Deserialise value
	Ok(Self {
	    value: serde_cbor::from_slice(&data[..]).wrap_err(eyre!("Failed to deserialise value"))?,
	    header: serial.header.clone(),
	    key,
	    sign
	})
    }
}

mod binary;
pub use binary::*;

impl<V: ?Sized> SerializedMessage<V>
{
    /// Get the message header
    #[inline(always)] pub fn header(&self) -> MessageHeader<'_, V>
    {
	MessageHeader(&self.header, PhantomData)
    }

}
/*
    /// Consume into an async writer
    pub async fn into_writer_async<W:AsyncWrite+Unpin>(self, mut writer: W) -> eyre::Result<usize>
    {
	let mut w = 0;
	macro_rules! write {
	    ($b:expr) => {
		w+=write_all_async(&mut writer, $b).await?
	    };
	    (? $o:expr) => {
		match $o {
		    Some(key) => {
			write!([1]);
			write!(key);
		    },
		    None => {
			write!([0]);
		    },
		}
	    };
	    (: $ser:expr) => {
		{
		    let mut v = StackVec::new();
		    serde_cbor::to_writer(&mut v, $ser)?;
		    write!(&v[..]);
		}
	    }
	}
	write!(: &self.header);
	write!(u64::try_from(self.data.len())?.to_be_bytes());
	write!(self.data);
	write!(self.hash);
	write!(? self.enc_key);
	write!(? self.sig);
	
	Ok(w)
    }
    /// Consume into a syncronous writer
    pub fn into_writer(self, mut writer: impl io::Write) -> eyre::Result<usize>
    {
	let mut w = 0;
	macro_rules! write {
	    ($b:expr) => {
		w+=write_all(&mut writer, $b)?
	    };
	    (? $o:expr) => {
		match $o {
		    Some(key) => {
			write!([1]);
			write!(key);
		    },
		    None => {
			write!([0]);
		    },
		}
	    };
	    (: $ser:expr) => {
		{
		    let mut ser = StackVec::new();
		    serde_cbor::to_writer(&mut ser, $ser)?;
		    write!(u64::try_from(ser.len())?.to_be_bytes());
		    write!(ser);
		    /*
		    let mut w2 = WriteCounter(0, &mut writer);
		    serde_cbor::to_writer(&mut w2, $ser)?;
		    w+=w2.0;*/
		}
	    };
	}
	write!(: &self.header);
	write!(u64::try_from(self.data.len())?.to_be_bytes());
	write!(self.data);
	write!(self.hash);
	write!(? self.enc_key);
	write!(? self.sig);
	
	Ok(w)
    }
    /// Consume into `Vec<u8>`.
    pub fn into_bytes(self) -> Vec<u8>
    {
	let mut v = Vec::with_capacity(self.data.len()<<1);
	self.into_writer(&mut v).expect("Failed to write to in-memory buffer");
	v
    }
}
 */
/*
impl<V: ?Sized + MessageValue> SerializedMessage<V>
{
    /// Create from a reader.
    ///
    /// The message may be in an invalid state. It is only possible to extract the value after validating it into a `Message<V>`.
    pub fn from_reader(mut reader: impl io::Read) -> eyre::Result<Self>
    {
	macro_rules! read {
	    
	    ($b:expr; $fmt:literal $($tt:tt)*) => {
		read_all($b, &mut reader).wrap_err(eyre!($fmt $($tt)*))
	    };
	    ($b:expr) => {
		read!($b; "Failed to read from stream")?;
	    };
	    (? $ot:expr) => {
		{
		    let mut b = [0u8; 1];
		    read!(&mut b[..]);
		    match b[0] {
			1 => {
			    let mut o = $ot;
			    read!(&mut o);
			    Some(o)
			},
			0 => {
			    None
			},
			x => {
			    return Err(eyre!("Invalid option state {:?}", x));
			}
		    }
		}
	    };
	    (: $ser:ty) => {
		{
		    let mut len = [0u8; std::mem::size_of::<u64>()];
		    read!(&mut len[..]);
		    let len = usize::try_from(u64::from_be_bytes(len))?;
		    //TODO: Find realistic max size for `$ser`.
		    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;
		    copy_buffer($into, &mut reader, num).wrap_err(eyre!("Failed to read {} bytes from reader", num))?
		}
	    }
	}
	let header = read!(: SerHeader);
	let data_len = {
	    let mut bytes = [0u8; std::mem::size_of::<u64>()];
	    read!(&mut bytes);
	    u64::from_be_bytes(bytes)
	}.try_into()?;
	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 bytes
    #[inline] pub fn from_bytes(bytes: impl AsRef<[u8]>) -> eyre::Result<Self>
    {
	let bytes = bytes.as_ref();
	Self::from_reader(&mut &bytes[..])
    }
}
*/
#[cfg(test)]
mod tests;