rsh/src/message.rs

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

mod serial;
pub use serial::*;

/// 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

/// A value that can be used for messages.
pub trait MessageValue: Serialize + for<'de> Deserialize<'de>{}

#[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,
}

/// `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>,
}

/// 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 + MessageValue>
{
    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>,
    
    _phantom: PhantomData<V>,
}

impl<V: ?Sized + MessageValue> Message<V>
{
    /// Serialise this message into one that can be converted to/from bytes.
    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)?;
	
	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,
	})
    }
}

impl<V: ?Sized + MessageValue> SerializedMessage<V>
{
    /// 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 v = serde_cbor::to_vec($ser)?;
		    write!(&v[..]);
		}
	    }
	}
	write!(: &self.header);
	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 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)
    }
    /// 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) => {
		read_all($b, &mut reader)?
	    };
	    (? $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) => {
		serde_cbor::from_reader::<$ty, _>(&mut reader)?
	    };
	    (:) => {
		serde_cbor::from_reader(&mut reader)?
	    };
	    ($into:expr, $num:expr) => {
		copy_buffer($into, &mut reader, $num)?
	    }
	}
	let header: SerHeader = read!(:);
	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 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,
	})
    }
    /// 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
    }
    /// 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[..])
    }
}