DualStream wrapper seems to work

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"

src/dual.rs

@@ -1,119 +1,289 @@
 //! Container for switching between un/encrypted stream 
 use super::*;
-use std::mem::{self, MaybeUninit, ManuallyDrop};
+use std::mem;
-use std::ops::{Drop, Deref, DerefMut};
+use chacha20stream::{
-use std::ptr;
+    AsyncSink,
-use std::fmt;
+    Key, IV,
-use chacha20stream::AsyncSink;
+};
+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<S>
+pub enum DualStream<S>
 {
+    /// If there is a panic while switching modes, the stream is left in this invariant state.
+    Poisoned,
+    
     Encrypted(AsyncSink<S>),
     Plain(S),
 }
 
-pub struct DualStream<S>(MaybeUninit<Box<DualStreamKind<S>>>);
+//pub type Error = openssl::error::ErrorStack;
 
-impl<S: fmt::Debug> fmt::Debug for DualStream<S>
+impl<'a, S: AsyncWrite + Unpin> DualStream<S>
+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.
 	
-impl<S> DualStream<S>
+	use tokio::prelude::*;
+	self.flush().await?;
+	self.to_crypt_now(enc, key, iv);
+	Ok(())
+    }
+}
+impl<'a, S: AsyncWrite> DualStream<S>
+where S: 'a
 {
-    fn as_mut_ref(&mut self) -> &mut Box<DualStreamKind<S>>
+    /// Create a transparent wrapper
+    ///
+    /// Identical to constructing the enum variant `Self::Plain` manually.
+    #[inline(always)] pub fn plain(stream: S) -> Self
     {
-	// SAFETY: It is always initialised except exactly within the swap function
+	Self::Plain(stream)
-	unsafe {
-	    &mut *self.0.as_mut_ptr() 
     }
-    }
+
-    fn as_ref(&self) -> &Box<DualStreamKind<S>>
+    /// Construct an encrypting wrapper over this stream
+    ///
+    /// # Panics
+    /// If constructing the cipher fails
+    #[inline] pub fn encrypt(stream: S, key: Key, iv: IV) -> Self
     {
-	// SAFETY: It is always initialised except exactly within the swap function
+	Self::Encrypted(AsyncSink::encrypt(stream, key, iv).expect("Initialising cipher failed"))
-	unsafe {
-	    &*self.0.as_ptr() 
-	}
     }
 
-    pub fn 
+    /// 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"))
+    }
 
-    /// Create explicit
+    /// Construct an encrypting or decrypting wrapper over this stream
-    pub fn new(k: DualStreamKind<S>) -> Self
+    ///
+    /// # Panics
+    /// If constructing the cipher fails
+    #[inline(always)] pub fn crypt(stream: S, enc: Encryption, key: Key, iv: IV) -> Self
     {
-	Self(MaybeUninit::new(Box::new(k)))
+	match enc {
+	    Encryption::Encrypt => Self::encrypt(stream, key, iv),
+	    Encryption::Decrypt => Self::decrypt(stream, key, iv),
+	}
     }
     
-    /// Consume into explicit (non-swappable) dual stream
+    /// Is this stream set to encrypted?
-    pub fn into_inner(self) -> Box<DualStreamKind<S>>
+    #[inline] pub fn is_encrypted(&self) -> bool
     {
-	let mut md = ManuallyDrop::new(self);
+	if let Self::Encrypted(_) = self {
-	unsafe {
+	    true
-	    // 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.
+	} else {
-	    mem::replace(&mut md.0, MaybeUninit::uninit()).assume_init()
+	    false
 	}
     }
-}
+    /// Is this stream in an invalid state?
+    #[inline(always)] pub fn is_poisoned(&self) -> bool {
 	
-impl<S> Deref for DualStream<S>
+	if let Self::Poisoned = self {
-{
+	    true
-    type Target = DualStreamKind<S>;
+	} else {
-    fn deref(&self) -> &Self::Target {
+	    false
-	self.as_ref()
 	}
-}
-impl<S> DerefMut for DualStream<S>
-{
-    fn deref_mut(&mut self) -> &mut Self::Target {
-	self.as_mut_ref()
     }
-}
 
-impl<S> From<Box<DualStreamKind<S>>> for DualStream<S>
+    /// 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
-    fn from(from: Box<DualStreamKind<S>>) -> Self
     {
-	Self(MaybeUninit::new(from))
+	mem::replace(self, Self::Poisoned)
     }
-}
     
-impl<S> From<DualStreamKind<S>> for DualStream<S>
+    /// Immediately convert this stream into an encrypted one.
-{
+    ///
-    fn from(from: DualStreamKind<S>) -> Self
+    /// # 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)
     {
-	Self::new(from)
+	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"));
     }
-}
 
-impl<S> From<DualStream<S>> for Box<DualStreamKind<S>>
+    /// Immediately convert this stream into a plain one
-{
+    ///
-    fn from(from: DualStream<S>) -> Self
+    /// # 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)
     {
-	from.into_inner()
+	*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")
+	}
     }
-}
 
+    /// 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")
+	}
+    }
 
-impl<S> From<DualStream<S>> for DualStreamKind<S>
+    /// As a mutable dynamic object 
-{
+    pub fn as_dyn_mut(&mut self) -> &mut (dyn AsyncWrite + 'a)
-    fn from(from: DualStream<S>) -> Self
+    {
+	match self {
+	    Self::Plain(p) => p,
+	    Self::Encrypted(e) => e,
+	    _ => panic!("Poisoned")
+	}
+    }
+
+    /// 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<S> Drop for DualStream<S>
+impl<S: AsyncWrite> AsyncWrite for DualStream<S>
 {
-    fn drop(&mut self) {
+    fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<Result<usize, io::Error>> {
-	// SAFETY: Value is always initialised except exactly within swap function
+	let obj = unsafe {
-	unsafe {
+	    self.map_unchecked_mut(|this| this.as_dyn_mut())		
-	    ptr::drop_in_place(self.0.as_mut_ptr())
+	};
+	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.
+    // 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[..]));
+    }
+}

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<bool> 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<T>(self, value: T) -> Option<T>
+	    {
+		self.and_then(move || value)
+	    }
+	    
+	    /// Map this value
+	    #[inline] pub fn map<F, T>(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<F, T>(self, f: F) -> Option<T>
+	    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<F, T>(self, f: F) -> Option<T>
+	    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<T>(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<F, G, T>(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);
+    }
+}

src/lib.rs

@@ -1,4 +1,5 @@
 
 #![allow(dead_code)]
 
+mod ext; #[macro_use] use ext::*;
 mod dual;