//! 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 = [ 0 u8 ; 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 = [ 0 u8 ; 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 ! ( ? [ 0 u8 ; 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 [ .. ] )
}
}