parent
d418cf876b
commit
136dc5f102
GPG Key ID:
284488987C31F630
@ -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<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
|
||||
unsafe {
|
||||
&mut *self.0.as_mut_ptr()
|
||||
}
|
||||
Self::Plain(stream)
|
||||
}
|
||||
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
|
||||
unsafe {
|
||||
&*self.0.as_ptr()
|
||||
}
|
||||
Self::Encrypted(AsyncSink::encrypt(stream, key, iv).expect("Initialising cipher failed"))
|
||||
}
|
||||
|
||||
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
|
||||
pub fn new(k: DualStreamKind<S>) -> Self
|
||||
/// 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
|
||||
{
|
||||
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
|
||||
pub fn into_inner(self) -> Box<DualStreamKind<S>>
|
||||
/// Is this stream set to encrypted?
|
||||
#[inline] pub fn is_encrypted(&self) -> bool
|
||||
{
|
||||
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()
|
||||
if let Self::Encrypted(_) = self {
|
||||
true
|
||||
} else {
|
||||
false
|
||||
}
|
||||
}
|
||||
}
|
||||
/// Is this stream in an invalid state?
|
||||
#[inline(always)] pub fn is_poisoned(&self) -> bool {
|
||||
|
||||
impl<S> Deref for DualStream<S>
|
||||
{
|
||||
type Target = DualStreamKind<S>;
|
||||
fn deref(&self) -> &Self::Target {
|
||||
self.as_ref()
|
||||
if let Self::Poisoned = self {
|
||||
true
|
||||
} else {
|
||||
false
|
||||
}
|
||||
}
|
||||
}
|
||||
impl<S> DerefMut for DualStream<S>
|
||||
{
|
||||
fn deref_mut(&mut self) -> &mut Self::Target {
|
||||
self.as_mut_ref()
|
||||
|
||||
/// 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)
|
||||
}
|
||||
}
|
||||
|
||||
impl<S> From<Box<DualStreamKind<S>>> for DualStream<S>
|
||||
{
|
||||
fn from(from: Box<DualStreamKind<S>>) -> 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<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(MaybeUninit::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<DualStreamKind<S>> for DualStream<S>
|
||||
{
|
||||
fn from(from: DualStreamKind<S>) -> Self
|
||||
/// 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::new(from)
|
||||
*self = Self::Plain(match self.poison() {
|
||||
Self::Plain(p) => p,
|
||||
Self::Encrypted(e) => e.into_inner(),
|
||||
_ => panic!("Poisoned"),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
impl<S> From<DualStream<S>> for Box<DualStreamKind<S>>
|
||||
{
|
||||
fn from(from: DualStream<S>) -> Self
|
||||
/// 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
|
||||
{
|
||||
from.into_inner()
|
||||
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")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 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> From<DualStream<S>> for DualStreamKind<S>
|
||||
{
|
||||
fn from(from: DualStream<S>) -> 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<S> Drop for DualStream<S>
|
||||
impl<S: AsyncWrite> AsyncWrite for DualStream<S>
|
||||
{
|
||||
fn drop(&mut self) {
|
||||
// SAFETY: Value is always initialised except exactly within swap function
|
||||
unsafe {
|
||||
ptr::drop_in_place(self.0.as_mut_ptr())
|
||||
}
|
||||
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.
|
||||
// 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[..]));
|
||||
}
|
||||
}
|
||||
|
||||
@ -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);
|
||||
}
|
||||
}
|
||||
@ -1,4 +1,5 @@
|
||||
|
||||
#![allow(dead_code)]
|
||||
|
||||
mod ext; #[macro_use] use ext::*;
|
||||
mod dual;
|
||||
|
||||
Loading…
Reference in new issue