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;

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

#[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!(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
    }
}