From 136dc5f1024ee3ebe5cd71ab4ff2817bd465e491 Mon Sep 17 00:00:00 2001 From: Avril Date: Sun, 11 Apr 2021 01:47:44 +0100 Subject: [PATCH] DualStream wrapper seems to work --- Cargo.toml | 2 + src/dual.rs | 308 ++++++++++++++++++++++++++++++++++++++++------------ src/ext.rs | 126 +++++++++++++++++++++ src/lib.rs | 1 + 4 files changed, 368 insertions(+), 69 deletions(-) create mode 100644 src/ext.rs diff --git a/Cargo.toml b/Cargo.toml index 20c5f3d..a45cb0e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -7,4 +7,6 @@ edition = "2018" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [dependencies] +tokio = {version = "0.2", features=["full"]} chacha20stream = {version = "1.0", features=["async"]} +openssl = "0.10.33" diff --git a/src/dual.rs b/src/dual.rs index 945aa9a..8ef20aa 100644 --- a/src/dual.rs +++ b/src/dual.rs @@ -1,119 +1,289 @@ //! Container for switching between un/encrypted stream use super::*; -use std::mem::{self, MaybeUninit, ManuallyDrop}; -use std::ops::{Drop, Deref, DerefMut}; -use std::ptr; -use std::fmt; -use chacha20stream::AsyncSink; +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); #[derive(Debug)] -pub enum DualStreamKind +pub enum DualStream { + /// If there is a panic while switching modes, the stream is left in this invariant state. + Poisoned, + Encrypted(AsyncSink), Plain(S), } -pub struct DualStream(MaybeUninit>>); +//pub type Error = openssl::error::ErrorStack; -impl fmt::Debug for DualStream +impl<'a, S: AsyncWrite + Unpin> DualStream +where S: 'a { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result + /// 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<()> { - fmt::Debug::fmt(self.as_ref(), f) + 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 +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")) + } -impl DualStream -{ - fn as_mut_ref(&mut self) -> &mut Box> + /// 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 { - // SAFETY: It is always initialised except exactly within the swap function - unsafe { - &mut *self.0.as_mut_ptr() + match enc { + Encryption::Encrypt => Self::encrypt(stream, key, iv), + Encryption::Decrypt => Self::decrypt(stream, key, iv), } } - fn as_ref(&self) -> &Box> + + /// Is this stream set to encrypted? + #[inline] pub fn is_encrypted(&self) -> bool { - // SAFETY: It is always initialised except exactly within the swap function - unsafe { - &*self.0.as_ptr() + 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 } } - pub fn + /// 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) + } - /// Create explicit - pub fn new(k: DualStreamKind) -> Self + /// 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`, 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) { - Self(MaybeUninit::new(Box::new(k))) + 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")); } - /// Consume into explicit (non-swappable) dual stream - pub fn into_inner(self) -> Box> + /// 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) { - let mut md = ManuallyDrop::new(self); - unsafe { - // We could just `read()` the pointer, but this is more semantiacally equivalent to moving the value, I think. Idk if it's more or less efficient or whatever. - mem::replace(&mut md.0, MaybeUninit::uninit()).assume_init() - } + *self = Self::Plain(match self.poison() { + Self::Plain(p) => p, + Self::Encrypted(e) => e.into_inner(), + _ => panic!("Poisoned"), + }); } -} -impl Deref for DualStream -{ - type Target = DualStreamKind; - fn deref(&self) -> &Self::Target { - self.as_ref() + /// 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") + } } -} -impl DerefMut for DualStream -{ - fn deref_mut(&mut self) -> &mut Self::Target { - self.as_mut_ref() + /// 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") + } } -} -impl From>> for DualStream -{ - fn from(from: Box>) -> Self + /// As an immutable dynamic object + pub fn as_dyn(&self) -> &(dyn AsyncWrite + 'a) { - Self(MaybeUninit::new(from)) + match self { + Self::Plain(p) => p, + Self::Encrypted(e) => e, + _ => panic!("Poisoned") + } } -} -impl From> for DualStream -{ - fn from(from: DualStreamKind) -> Self + /// As a mutable dynamic object + pub fn as_dyn_mut(&mut self) -> &mut (dyn AsyncWrite + 'a) { - Self::new(from) + match self { + Self::Plain(p) => p, + Self::Encrypted(e) => e, + _ => panic!("Poisoned") + } } -} -impl From> for Box> -{ - fn from(from: DualStream) -> Self + /// Consume into the inner (plain) stream + #[inline] pub fn into_inner(self) -> S { - from.into_inner() + match self { + Self::Plain(p) => p, + Self::Encrypted(e) => e.into_inner(), + _ => panic!("Poisoned") + } } } - -impl From> for DualStreamKind +impl AsyncWrite for DualStream { - fn from(from: DualStream) -> Self - { - *from.into_inner() + fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll> { + 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> { + 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> { + let obj = unsafe { + self.map_unchecked_mut(|this| this.as_dyn_mut()) + }; + obj.poll_shutdown(cx) } } -impl Drop for DualStream +#[cfg(test)] +mod tests { - fn drop(&mut self) { - // SAFETY: Value is always initialised except exactly within swap function - unsafe { - ptr::drop_in_place(self.0.as_mut_ptr()) - } + 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. + // TODO: Read the `written` buffer back in the same way and order, and check for identity. + 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); + + // Encrypted + wrapper.to_crypt(super::Encryption::Encrypt, key, iv).await.unwrap(); + wrapper.write_all(input.as_bytes()).await.unwrap(); + + // Unencrypted + 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[..])); } } - diff --git a/src/ext.rs b/src/ext.rs new file mode 100644 index 0000000..7e7e9dc --- /dev/null +++ b/src/ext.rs @@ -0,0 +1,126 @@ +use super::*; + +#[macro_export] macro_rules! basic_enum { + ($(#[$meta:meta])* $vis:vis $name:ident $(; $tcomment:literal)?: $($var:ident $(=> $comment:literal)?),+ $(,)?) => { + $(#[$meta])* + #[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Copy)] + $(#[doc = $tcomment])? + $vis enum $name { + $( + $(#[doc = $comment])? + $var + ),+ + } + } +} + + +/// Create a `Yes` or `No` enum. +#[macro_export] macro_rules! bool_type { + ($vis:vis $name:ident $(; $comment:literal)? => $yes:ident, $no:ident) => { + basic_enum!(#[repr(u8)] $vis $name $(; $comment)?: $yes => "# First variant\n\nYes/true", $no => "# Second variant\n\nNo/false"); + + impl From for $name + { + #[inline] fn from(from: bool) -> Self + { + if from { + Self::$yes + } else { + Self::$no + } + } + } + + impl From<$name> for bool + { + #[inline] fn from(from: $name) -> Self + { + match from { + $name::$yes => true, + $name::$no => false, + } + } + } + + impl $name + { + /// Create from a bool value. + #[inline] pub const fn new(from: bool) -> Self + { + if from { + Self::$yes + } else { + Self::$no + } + } + + /// Is this false? + #[inline] pub const fn is_no(self) -> bool + { + !self.is_yes() + } + /// Is this true? + #[inline] pub const fn is_yes(self) -> bool + { + match self { + Self::$yes => true, + Self::$no => false, + } + } + + /// Return Some(T) if self is true. + #[inline] pub fn some(self, value: T) -> Option + { + self.and_then(move || value) + } + + /// Map this value + #[inline] pub fn map(self, f: F) -> T + where F: FnOnce(bool) -> T + { + f(self.is_yes()) + } + + /// Run this closure if value is false + #[inline] pub fn or_else(self, f: F) -> Option + where F: FnOnce() -> T + { + if let Self::$no = self { + Some(f()) + } else { + None + } + } + /// Run this closure if value is true + #[inline] pub fn and_then(self, f: F) -> Option + where F: FnOnce() -> T + { + if let Self::$yes = self { + Some(f()) + } else { + None + } + } + + /// Return `yes` if true and `no` if false + #[inline] pub fn either(self, yes: T, no: T) -> T + { + self.and_either(move || yes, move || no) + } + /// Run closure `yes` if value is true, `no` if value is false. + #[inline] pub fn and_either(self, yes: F, no: G) -> T + where F: FnOnce() -> T, + G: FnOnce() -> T, + { + match self { + Self::$yes => yes(), + Self::$no => no(), + } + } + } + }; + ($vis:vis $name:ident $(; $comment:literal)?) => { + $crate::bool_type!($vis $name $(; $comment)? => Yes, No); + } +} diff --git a/src/lib.rs b/src/lib.rs index 22cf027..f6fd2ed 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,4 +1,5 @@ #![allow(dead_code)] +mod ext; #[macro_use] use ext::*; mod dual;