//! Container for switching between un/encrypted stream
use super ::* ;
use std ::mem ;
use chacha20stream ::{
AsyncSink ,
Key , IV ,
} ;
use tokio ::io ::AsyncWrite ;
use std ::{
pin ::Pin ,
task ::{ Context , Poll } ,
io ,
marker ::Unpin ,
} ;
bool_type ! ( pub Encrypted ; "Is the value encrypted?" ) ;
bool_type ! ( pub Encryption ; "What way are we en/decrypting?" = > Encrypt , Decrypt ) ;
/// A wrapper `AsyncWrite` stream that allows switching between encrypted (chacha20stream) and plain stream writing.
///
/// # Polymorphic dispatching
/// While this type implements `AsyncWrite` itself, it is recommended to use the polymorphic mutable reference returned by `as_dyn_unpin_mut()` (or `as_dyn_mut()` for non-`Unpin` values of `S`) for writing if lots of `AsyncWrite` methods will be dispatched on the instance.
/// This will prevent the need to check the discriminant of the enum each time the `DualStream`'s `AsyncWrite` methods are polled.
/// The type implements `AsRef/Mut` for streams `S` that are both `Unpin` and not `Unpin` for convenience.
#[ derive(Debug) ]
pub enum DualStream < S >
{
/// If there is a panic while switching modes, the stream is left in this invariant state.
///
/// Stream `S` is dropped if the instance is in this state.
Poisoned ,
/// Stream `S` is being written to through a chacha20stream `AsyncSink` stream cipher.
Encrypted ( AsyncSink < S > ) ,
/// Stream `S` is being written to directly.
Plain ( S ) ,
}
// We can use dynamic dispatching to prevent the need to check the enum's discriminant each time.
// We have `AsRef/Mut`s for normal and `Unpin` polymorphic `AsyncWrite`s
impl < ' a , S : AsyncWrite + Unpin > AsRef < dyn AsyncWrite + Unpin + ' a > for DualStream < S >
where S : ' a
{
fn as_ref ( & self ) -> & ( dyn AsyncWrite + Unpin + ' a )
{
self . as_dyn_unpin ( )
}
}
impl < ' a , S : AsyncWrite + Unpin > AsMut < dyn AsyncWrite + Unpin + ' a > for DualStream < S >
where S : ' a
{
fn as_mut ( & mut self ) -> & mut ( dyn AsyncWrite + Unpin + ' a )
{
self . as_dyn_unpin_mut ( )
}
}
impl < ' a , S : AsyncWrite > AsRef < dyn AsyncWrite + ' a > for DualStream < S >
where S : ' a
{
fn as_ref ( & self ) -> & ( dyn AsyncWrite + ' a )
{
self . as_dyn ( )
}
}
impl < ' a , S : AsyncWrite > AsMut < dyn AsyncWrite + ' a > for DualStream < S >
where S : ' a
{
fn as_mut ( & mut self ) -> & mut ( dyn AsyncWrite + ' a )
{
self . as_dyn_mut ( )
}
}
impl < ' a , S : AsyncWrite + Unpin > DualStream < S >
where S : ' a
{
/// Convert this stream into an encrypted one.
///
/// # Notes
/// This method makes sure to `flush()` the encrypted stream before dropping the cipher.
///
/// # Panics
/// If initialising the cipher fails, this method will panic.
/// # Poisons
/// If this method panics, then the inner stream will be dropped and this instance will be set to `Poisoned`. This is (currently) irrecoverable.
pub async fn to_plain ( & mut self ) -> io ::Result < ( ) >
{
if ! self . is_encrypted ( ) {
// No need to do anything
return Ok ( ( ) ) ;
}
use tokio ::prelude ::* ;
self . flush ( ) . await ? ;
self . to_plain_now ( ) ;
Ok ( ( ) )
}
/// Convert this stream into an encrypted one.
///
/// # Notes
/// This method makes sure to `flush()` the encrypted stream before constructing the cipher.
///
/// # Panics
/// If initialising the cipher fails, this method will panic.
/// # Poisons
/// If this method panics, then the inner stream will be dropped and this instance will be set to `Poisoned`. This is (currently) irrecoverable.
pub async fn to_crypt ( & mut self , enc : Encryption , key : Key , iv : IV ) -> io ::Result < ( ) >
{
// NOTE: We can't skip this like `to_plain()` does by checking if the instance is already at `Encrypted` as the key/IV may differ.
use tokio ::prelude ::* ;
self . flush ( ) . await ? ;
self . to_crypt_now ( enc , key , iv ) ;
Ok ( ( ) )
}
}
impl < ' a , S : AsyncWrite > DualStream < S >
where S : Unpin + ' a
{
/// As an immutable dynamic object for `Unpin` streams.
pub fn as_dyn_unpin ( & self ) -> & ( dyn AsyncWrite + Unpin + ' a )
{
match self {
Self ::Plain ( p ) = > p ,
Self ::Encrypted ( e ) = > e ,
_ = > panic! ( "Poisoned" )
}
}
/// As a mutable dynamic object for `Unpin` streams
pub fn as_dyn_unpin_mut ( & mut self ) -> & mut ( dyn AsyncWrite + Unpin + ' a )
{
match self {
Self ::Plain ( p ) = > p ,
Self ::Encrypted ( e ) = > e ,
_ = > panic! ( "Poisoned" )
}
}
}
impl < ' a , S : AsyncWrite > DualStream < S >
where S : ' a
{
/// Create a transparent wrapper
///
/// Identical to constructing the enum variant `Self::Plain` manually.
#[ inline(always) ] pub fn plain ( stream : S ) -> Self
{
Self ::Plain ( stream )
}
/// Construct an encrypting wrapper over this stream
///
/// # Panics
/// If constructing the cipher fails
#[ inline ] pub fn encrypt ( stream : S , key : Key , iv : IV ) -> Self
{
Self ::Encrypted ( AsyncSink ::encrypt ( stream , key , iv ) . expect ( "Initialising cipher failed" ) )
}
/// Construct a decrypting wrapper over this stream
///
/// # Panics
/// If constructing the cipher fails
#[ inline ] pub fn decrypt ( stream : S , key : Key , iv : IV ) -> Self
{
Self ::Encrypted ( AsyncSink ::decrypt ( stream , key , iv ) . expect ( "Initialising cipher failed" ) )
}
/// Construct an encrypting or decrypting wrapper over this stream
///
/// # Panics
/// If constructing the cipher fails
#[ inline(always) ] pub fn crypt ( stream : S , enc : Encryption , key : Key , iv : IV ) -> Self
{
match enc {
Encryption ::Encrypt = > Self ::encrypt ( stream , key , iv ) ,
Encryption ::Decrypt = > Self ::decrypt ( stream , key , iv ) ,
}
}
/// Is this stream set to encrypted?
#[ inline ] pub fn is_encrypted ( & self ) -> bool
{
if let Self ::Encrypted ( _ ) = self {
true
} else {
false
}
}
/// Is this stream in an invalid state?
#[ inline(always) ] pub fn is_poisoned ( & self ) -> bool {
if let Self ::Poisoned = self {
true
} else {
false
}
}
/// Move out of self, and in turn poison this insatance until self is assigned a proper value again.
#[ inline(always) ] fn poison ( & mut self ) -> Self
{
mem ::replace ( self , Self ::Poisoned )
}
/// Immediately convert this stream into an encrypted one.
///
/// # Notes
/// Make sure to `flush()` the stream **before** calling this or data may be lost.
///
/// # Panics
/// If initialising the cipher fails, this method will panic.
/// # Poisons
/// If this method panics, then the inner stream will be dropped and this instance will be set to `Poisoned`. This is (currently) irrecoverable.
//TODO: in chacha20stream: Add `try_encrypt()`, `try_decrypt()` which both return `Result<Self, (S, Error)>`, to not lose the inner stream if the cipher fails to initialise.
#[ inline ] pub fn to_crypt_now ( & mut self , enc : Encryption , key : Key , iv : IV )
{
let inner = match self . poison ( ) {
Self ::Encrypted ( enc ) = > enc . into_inner ( ) ,
Self ::Plain ( inner ) = > inner ,
_ = > panic! ( "Poisoned" ) ,
} ;
* self = Self ::Encrypted ( match enc {
Encryption ::Encrypt = > AsyncSink ::encrypt ( inner , key , iv ) ,
Encryption ::Decrypt = > AsyncSink ::decrypt ( inner , key , iv ) ,
} . expect ( "Initialising cipher failed" ) ) ;
}
/// Immediately convert this stream into a plain one
///
/// # Notes
/// Make sure to `flush()` the stream **before** calling this or encrypted data may be lost.
///
/// # Panics
/// If dropping the cipher fails, this method will panic.
/// # Poisons
/// If this method panics, then the inner stream will be dropped and this instance will be set to `Poisoned`. This is (currently) irrecoverable.
#[ inline ] pub fn to_plain_now ( & mut self )
{
* self = Self ::Plain ( match self . poison ( ) {
Self ::Plain ( p ) = > p ,
Self ::Encrypted ( e ) = > e . into_inner ( ) ,
_ = > panic! ( "Poisoned" ) ,
} ) ;
}
/// A mutable reference to the inner (plain) stream, whether this instance is set to encrypted or not.
pub fn inner_plain_mut ( & mut self ) -> & mut S
{
match self {
Self ::Plain ( p ) = > p ,
Self ::Encrypted ( e ) = > e . inner_mut ( ) ,
_ = > panic! ( "Poisoned" )
}
}
/// A reference to the inner (plain) stream, whether this instance is set to encrypted or not.
pub fn inner_plain ( & self ) -> & S
{
match self {
Self ::Plain ( p ) = > p ,
Self ::Encrypted ( e ) = > e . inner ( ) ,
_ = > panic! ( "Poisoned" )
}
}
/// Consume into the inner (plain) stream
#[ inline ] pub fn into_inner ( self ) -> S
{
match self {
Self ::Plain ( p ) = > p ,
Self ::Encrypted ( e ) = > e . into_inner ( ) ,
_ = > panic! ( "Poisoned" )
}
}
/// As an immutable dynamic object
pub fn as_dyn ( & self ) -> & ( dyn AsyncWrite + ' a )
{
match self {
Self ::Plain ( p ) = > p ,
Self ::Encrypted ( e ) = > e ,
_ = > panic! ( "Poisoned" )
}
}
/// As a mutable dynamic object
pub fn as_dyn_mut ( & mut self ) -> & mut ( dyn AsyncWrite + ' a )
{
match self {
Self ::Plain ( p ) = > p ,
Self ::Encrypted ( e ) = > e ,
_ = > panic! ( "Poisoned" )
}
}
}
impl < S : AsyncWrite > AsyncWrite for DualStream < S >
{
fn poll_write ( self : Pin < & mut Self > , cx : & mut Context < ' _ > , buf : & [ u8 ] ) -> Poll < Result < usize , io ::Error > > {
let obj = unsafe {
self . map_unchecked_mut ( | this | this . as_dyn_mut ( ) )
} ;
obj . poll_write ( cx , buf )
}
fn poll_flush ( self : Pin < & mut Self > , cx : & mut Context < ' _ > ) -> Poll < Result < ( ) , io ::Error > > {
let obj = unsafe {
self . map_unchecked_mut ( | this | this . as_dyn_mut ( ) )
} ;
obj . poll_flush ( cx )
}
fn poll_shutdown ( self : Pin < & mut Self > , cx : & mut Context < ' _ > ) -> Poll < Result < ( ) , io ::Error > > {
let obj = unsafe {
self . map_unchecked_mut ( | this | this . as_dyn_mut ( ) )
} ;
obj . poll_shutdown ( cx )
}
}
#[ cfg(test) ]
mod tests
{
use tokio ::prelude ::* ;
use chacha20stream ::keygen ;
#[ tokio::test ]
/// Write the whole `input` buffer encrypted, switch to plain, then write the first 5 bytes of `input` again.
async fn wrapper_construct ( )
{
let input = "Hello world!" ;
let backing = Vec ::new ( ) ;
let ( key , iv ) = keygen ( ) ;
let written = {
let mut wrapper = super ::DualStream ::Plain ( backing ) ;
// Encrypting
wrapper . to_crypt ( super ::Encryption ::Encrypt , key , iv ) . await . unwrap ( ) ;
wrapper . write_all ( input . as_bytes ( ) ) . await . unwrap ( ) ;
// Plain
wrapper . to_plain ( ) . await . unwrap ( ) ;
wrapper . write_all ( & input . as_bytes ( ) [ .. 5 ] ) . await . unwrap ( ) ;
// Shutdown the stream and consume it.
wrapper . flush ( ) . await . unwrap ( ) ;
wrapper . shutdown ( ) . await . unwrap ( ) ;
wrapper . into_inner ( )
} ;
eprintln! ( "Output bytes: {:?}" , written ) ;
eprintln! ( "Output attempted string: {:?}" , String ::from_utf8_lossy ( & written [ .. ] ) ) ;
}
// TODO: Write a test using Tokio's `duplex` to read and write encrypted bytes on 2 tasks, then compare the output to the input afterwards
}