async-encrypted-stream 0.2.1

#![doc = include_str!("../README.md")]
use std::ops::Sub;

pub use aead_stream;
pub use chacha20poly1305;

use aead_stream::{Decryptor, Encryptor, NewStream, NonceSize, StreamPrimitive};
use chacha20poly1305::aead::{
    AeadInOut,
    array::{Array, ArraySize},
};

use tokio::io::{AsyncRead, AsyncWrite};

mod read_half;
pub use read_half::ReadHalf;

mod write_half;
pub use write_half::WriteHalf;

pub const DEFAULT_BUFFER_SIZE: usize = 4096;
pub const DEFAULT_CHUNK_SIZE: usize = 1024;

pub type EncryptedStream<R, W, A, S> =
    (ReadHalf<R, Decryptor<A, S>>, WriteHalf<W, Encryptor<A, S>>);

/// Creates a pair of [ReadHalf] and [WriteHalf] with default buffer size of
/// [self::DEFAULT_BUFFER_SIZE] and chunk size of [self::DEFAULT_CHUNK_SIZE]  
///
/// ```rust
/// use async_encrypted_stream::aead_stream::{DecryptorLE31, EncryptorLE31};
/// use async_encrypted_stream::chacha20poly1305::XChaCha20Poly1305;
///
/// use async_encrypted_stream::{ReadHalf, WriteHalf, encrypted_stream};
///
/// let key = [0u8; 32];
/// let nonce = [0u8; 20];
///
/// let (rx, tx) = tokio::io::duplex(4096);
/// let (mut reader, mut writer): (
///     ReadHalf<_, DecryptorLE31<XChaCha20Poly1305>>,
///     WriteHalf<_, EncryptorLE31<XChaCha20Poly1305>>,
/// ) = encrypted_stream(rx, tx, (&key).into(), (&nonce).into());
/// ````
pub fn encrypted_stream<R: AsyncRead, W: AsyncWrite, A, S>(
    read: R,
    write: W,
    key: &Array<u8, A::KeySize>,
    nonce: &Array<u8, NonceSize<A, S>>,
) -> EncryptedStream<R, W, A, S>
where
    S: StreamPrimitive<A> + NewStream<A>,
    A: AeadInOut + chacha20poly1305::KeyInit,
    A::NonceSize: Sub<<S as StreamPrimitive<A>>::NonceOverhead>,
    NonceSize<A, S>: ArraySize,
{
    encrypted_stream_with_capacity(
        read,
        write,
        key,
        nonce,
        DEFAULT_BUFFER_SIZE,
        DEFAULT_CHUNK_SIZE,
    )
}

/// Creates a pair of [ReadHalf] and [WriteHalf] with default buffer size of
/// `buffer_size` and chunk size of `chunk_size`  
///
/// ```rust
/// use async_encrypted_stream::aead_stream::{DecryptorLE31, EncryptorLE31};
/// use async_encrypted_stream::chacha20poly1305::XChaCha20Poly1305;
///
/// use async_encrypted_stream::{ReadHalf, WriteHalf, encrypted_stream_with_capacity};
///
/// let key = [0u8; 32];
/// let nonce = [0u8; 20];
///
/// let (rx, tx) = tokio::io::duplex(4096);
/// let (mut reader, mut writer): (
///     ReadHalf<_, DecryptorLE31<XChaCha20Poly1305>>,
///     WriteHalf<_, EncryptorLE31<XChaCha20Poly1305>>,
/// ) = encrypted_stream_with_capacity(rx, tx, (&key).into(), (&nonce).into(), 4096,
/// 512);
/// ````
pub fn encrypted_stream_with_capacity<R: AsyncRead, W: AsyncWrite, A, S>(
    read: R,
    write: W,
    key: &Array<u8, A::KeySize>,
    nonce: &Array<u8, NonceSize<A, S>>,
    buffer_size: usize,
    chunk_size: usize,
) -> EncryptedStream<R, W, A, S>
where
    S: StreamPrimitive<A> + NewStream<A>,
    A: AeadInOut + chacha20poly1305::KeyInit,
    A::NonceSize: Sub<<S as StreamPrimitive<A>>::NonceOverhead>,
    NonceSize<A, S>: ArraySize,
{
    let encryptor = Encryptor::new(key, nonce);
    let decryptor = Decryptor::new(key, nonce);

    (
        ReadHalf::with_capacity(read, decryptor, buffer_size),
        WriteHalf::with_capacity(write, encryptor, buffer_size, chunk_size),
    )
}

#[cfg(test)]
fn get_key<const S: usize>(plain_key: &str, salt: &str) -> [u8; S] {
    const ITERATIONS: u32 = 4096;
    pbkdf2::pbkdf2_hmac_array::<sha2::Sha256, S>(plain_key.as_bytes(), salt.as_bytes(), ITERATIONS)
}

#[cfg(test)]
mod tests {
    use std::{assert_eq, time::Duration};

    use aead_stream::{DecryptorLE31, EncryptorLE31};
    use chacha20poly1305::XChaCha20Poly1305;
    use tokio::io::{AsyncReadExt, AsyncWriteExt};

    use super::*;

    #[tokio::test]
    pub async fn test_big_transfer() {
        let key: [u8; 32] = get_key("key", "group");
        let nonce = [0u8; 20];

        let (rx, tx) = tokio::io::duplex(4096);
        let (mut reader, mut writer): (
            ReadHalf<_, DecryptorLE31<XChaCha20Poly1305>>,
            WriteHalf<_, EncryptorLE31<XChaCha20Poly1305>>,
        ) = super::encrypted_stream(rx, tx, (&key).into(), (&nonce).into());

        let size = 1024 * 4;
        tokio::spawn(async move {
            let content = vec![100u8; size];
            let _ = writer.write(&content).await;
            let _ = writer.flush().await;
        });

        tokio::time::sleep(Duration::from_millis(100)).await;

        let mut bytes_expected = size;
        let mut read_buf = vec![0u8; 1024];
        while let Ok(bytes_read) = reader.read(&mut read_buf[..]).await {
            if bytes_read == 0 {
                break;
            }

            assert!(read_buf[..bytes_read].iter().all(|b| *b == 100));
            bytes_expected -= bytes_read;
        }

        assert_eq!(0, bytes_expected);
    }
}