You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
365 lines
11 KiB
365 lines
11 KiB
//! 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
|
|
|
|
}
|