pqfile 4.2.1

/// Async encrypt and decrypt streaming using `tokio::io`.
///
/// Requires the `async` feature flag: `pqfile = { features = ["async"] }`.
///
/// These functions are the async equivalents of [`crate::encrypt::encrypt_stream`]
/// and [`crate::decrypt::decrypt_stream`]. They accept any `tokio::io::AsyncRead +
/// AsyncWrite + Unpin` source and sink, enabling non-blocking encryption in async
/// servers and proxies without spawning a dedicated OS thread per operation.
///
/// The format and ciphertext are identical to the synchronous API - files
/// produced by `encrypt_stream_async` are indistinguishable from those produced
/// by `encrypt_stream` and can be decrypted by either.
///
/// # Memory note
///
/// The current implementation buffers the entire plaintext (or ciphertext) in
/// memory before performing crypto. This is safe for files up to
/// `crate::format::MAX_ORIGINAL_SIZE` (1 TiB) but is unsuitable for files
/// exceeding available RAM. Inputs larger than `MAX_ORIGINAL_SIZE` are rejected
/// with `EncryptionFailure` / `DecryptionFailure` before any allocation occurs.
/// A fully streaming implementation is planned for a future release.
///
/// Additionally, `poll_shutdown` on [`AsyncPqfWriter`] runs synchronous
/// KEM encapsulation and AEAD encryption on the calling tokio executor thread.
/// For CPU-bound workloads wrap the call in `tokio::task::spawn_blocking`.
use std::io::Write as _;
use std::pin::Pin;
use std::task::{Context, Poll};

use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use zeroize::Zeroizing;

use crate::decrypt;
use crate::encrypt;
use crate::error::PqfileError;
use crate::writer::PqfWriter;

/// Async equivalent of [`encrypt::encrypt_stream`] (v3 / v5 format).
///
/// Reads plaintext from `reader`, encrypts it in `chunk_size`-byte chunks, and
/// writes the authenticated ciphertext to `writer`. Emits v3 format when
/// `chunk_size == CHUNK_SIZE` (65536), v5 otherwise.
///
/// # Example
///
/// ```no_run
/// use pqfile::async_io::encrypt_stream_async;
/// use pqfile::format::CHUNK_SIZE;
///
/// # async fn example() {
/// let pub_pem = std::fs::read_to_string("pubkey.pem").unwrap();
/// let mut plaintext: &[u8] = b"hello, post-quantum async world";
/// let mut ciphertext = Vec::new();
/// encrypt_stream_async(&pub_pem, plaintext.len() as u64, CHUNK_SIZE, &mut plaintext, &mut ciphertext)
///     .await
///     .unwrap();
/// # }
/// ```
#[must_use = "encryption result must be checked"]
pub async fn encrypt_stream_async<R, W>(
    pubkey_pem: &str,
    original_size: u64,
    chunk_size: usize,
    reader: &mut R,
    writer: &mut W,
) -> Result<(), PqfileError>
where
    R: AsyncRead + Unpin,
    W: AsyncWrite + Unpin,
{
    if chunk_size == 0 {
        return Err(PqfileError::EncryptionFailure);
    }
    if original_size > crate::format::MAX_ORIGINAL_SIZE {
        return Err(PqfileError::EncryptionFailure);
    }

    // Buffer the full plaintext then delegate to the sync encrypt path.
    // This keeps the cryptographic core as a single well-tested code path while
    // providing a non-blocking interface for I/O.
    let mut plaintext = Vec::with_capacity(original_size.min(64 * 1024 * 1024) as usize);
    reader
        .read_to_end(&mut plaintext)
        .await
        .map_err(PqfileError::Io)?;
    if plaintext.len() as u64 > crate::format::MAX_ORIGINAL_SIZE {
        return Err(PqfileError::EncryptionFailure);
    }

    let mut ct = Vec::new();
    encrypt::encrypt_stream(
        pubkey_pem,
        plaintext.len() as u64,
        chunk_size,
        &mut plaintext.as_slice(),
        &mut ct,
    )?;

    writer.write_all(&ct).await.map_err(PqfileError::Io)?;
    writer.flush().await.map_err(PqfileError::Io)?;
    Ok(())
}

/// Async equivalent of [`decrypt::decrypt_stream`].
///
/// Reads a `.pqf` stream from `reader`, authenticates and decrypts it, and
/// writes the plaintext to `writer`. Supports all format versions (v2-v8).
///
/// # Example
///
/// ```no_run
/// use pqfile::async_io::decrypt_stream_async;
///
/// # async fn example() {
/// let priv_pem = std::fs::read_to_string("privkey.pem").unwrap();
/// let ciphertext: &[u8] = &[];
/// let mut plaintext = Vec::new();
/// decrypt_stream_async(&priv_pem, ciphertext, &mut plaintext, None)
///     .await
///     .unwrap();
/// # }
/// ```
#[must_use = "decryption result must be checked"]
pub async fn decrypt_stream_async<R, W>(
    privkey_pem: &str,
    reader: R,
    writer: &mut W,
    passphrase: Option<&str>,
) -> Result<(), PqfileError>
where
    R: AsyncRead + Unpin,
    W: AsyncWrite + Unpin,
{
    // Buffer the full ciphertext then delegate to sync decrypt.
    let mut ct_reader = tokio::io::BufReader::new(reader);
    let mut ct = Vec::new();
    ct_reader
        .read_to_end(&mut ct)
        .await
        .map_err(PqfileError::Io)?;
    // Ciphertext is always slightly larger than plaintext; reject absurdly large inputs.
    if ct.len() as u64 > crate::format::MAX_ORIGINAL_SIZE + (1 << 20) {
        return Err(PqfileError::DecryptionFailure);
    }

    let mut plaintext = Zeroizing::new(Vec::new());
    decrypt::decrypt_stream(privkey_pem, &mut ct.as_slice(), &mut *plaintext, passphrase)?;

    writer
        .write_all(&plaintext)
        .await
        .map_err(PqfileError::Io)?;
    writer.flush().await.map_err(PqfileError::Io)?;
    Ok(())
}

// ── AsyncPqfWriter ────────────────────────────────────────────────────────

/// Internal state for [`AsyncPqfWriter`].
enum AsyncWriterState<W: AsyncWrite + Unpin> {
    /// Still accepting plaintext writes; `plaintext` accumulates the input.
    Buffering {
        sink: W,
        plaintext: Zeroizing<Vec<u8>>,
    },
    /// Sync encryption finished; writing the ciphertext to `sink` starting at `pos`.
    Flushing { sink: W, ct: Vec<u8>, pos: usize },
    /// Sealing and I/O complete.
    Done,
}

/// Async streaming encryptor backed by any `W: AsyncWrite + Unpin`.
///
/// Accepts plaintext via `AsyncWrite::write` (each call buffers bytes in memory),
/// then encrypts everything in one pass and writes the ciphertext when either
/// [`AsyncPqfWriter::finish`] is called or the standard `poll_shutdown` path is
/// triggered (e.g. by `tokio::io::copy`).
///
/// The resulting ciphertext is identical to that produced by the synchronous
/// [`PqfWriter`] and is readable by both the sync and async decryptors.
///
/// # Example
///
/// ```no_run
/// use pqfile::async_io::AsyncPqfWriter;
/// use pqfile::format::CHUNK_SIZE;
/// use tokio::io::AsyncWriteExt;
///
/// # async fn example() {
/// let pub_pem = std::fs::read_to_string("pubkey.pem").unwrap();
/// let plaintext = b"async writer example";
/// let mut w = AsyncPqfWriter::new(Vec::<u8>::new(), &pub_pem, plaintext.len() as u64, CHUNK_SIZE).unwrap();
/// w.write_all(plaintext).await.unwrap();
/// let ciphertext = w.finish().await.unwrap();
/// # }
/// ```
pub struct AsyncPqfWriter<W: AsyncWrite + Unpin> {
    pubkey_pem: String,
    original_size: u64,
    chunk_size: usize,
    state: AsyncWriterState<W>,
}

impl<W: AsyncWrite + Unpin> AsyncPqfWriter<W> {
    /// Creates a new `AsyncPqfWriter` wrapping `sink`.
    ///
    /// No I/O is performed at construction time. All plaintext is buffered until
    /// [`finish`](Self::finish) or `poll_shutdown` is called.
    pub fn new(
        sink: W,
        pubkey_pem: &str,
        original_size: u64,
        chunk_size: usize,
    ) -> Result<Self, PqfileError> {
        if chunk_size == 0 {
            return Err(PqfileError::EncryptionFailure);
        }
        Ok(Self {
            pubkey_pem: pubkey_pem.to_owned(),
            original_size,
            chunk_size,
            state: AsyncWriterState::Buffering {
                sink,
                plaintext: Zeroizing::new(Vec::new()),
            },
        })
    }

    /// Seals the stream and returns the inner writer once all ciphertext has been flushed.
    ///
    /// Must be called once when all plaintext has been written. Calling `finish` is
    /// equivalent to awaiting `shutdown()` and then recovering the inner writer.
    pub async fn finish(mut self) -> Result<W, PqfileError> {
        self.seal_and_write().await
    }

    /// Encrypts the buffered plaintext, writes all ciphertext to the sink, and returns it.
    async fn seal_and_write(&mut self) -> Result<W, PqfileError> {
        // Move out of Buffering state.
        let (mut sink, plaintext) = match std::mem::replace(&mut self.state, AsyncWriterState::Done)
        {
            AsyncWriterState::Buffering { sink, plaintext } => (sink, plaintext),
            AsyncWriterState::Flushing { mut sink, ct, pos } => {
                // Already sealing - just finish writing.
                sink.write_all(&ct[pos..]).await.map_err(PqfileError::Io)?;
                sink.flush().await.map_err(PqfileError::Io)?;
                return Ok(sink);
            }
            AsyncWriterState::Done => {
                return Err(PqfileError::Io(std::io::Error::new(
                    std::io::ErrorKind::BrokenPipe,
                    "AsyncPqfWriter already finished",
                )));
            }
        };

        // Sync-encrypt the buffered plaintext (CPU-bound but fast for buffered model).
        let ct = {
            let mut ct = Vec::new();
            let mut w = PqfWriter::new(
                &mut ct,
                &self.pubkey_pem,
                self.original_size,
                self.chunk_size,
            )?;
            w.write_all(&plaintext).map_err(PqfileError::Io)?;
            w.finish()?;
            ct
        };

        // Write the full ciphertext async.
        sink.write_all(&ct).await.map_err(PqfileError::Io)?;
        sink.flush().await.map_err(PqfileError::Io)?;
        Ok(sink)
    }
}

impl<W: AsyncWrite + Unpin> AsyncWrite for AsyncPqfWriter<W> {
    fn poll_write(
        mut self: Pin<&mut Self>,
        _cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<std::io::Result<usize>> {
        match &mut self.state {
            AsyncWriterState::Buffering { plaintext, .. } => {
                plaintext.extend_from_slice(buf);
                Poll::Ready(Ok(buf.len()))
            }
            _ => Poll::Ready(Err(std::io::Error::new(
                std::io::ErrorKind::BrokenPipe,
                "AsyncPqfWriter already finished",
            ))),
        }
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
        match &mut self.state {
            AsyncWriterState::Buffering { sink, .. } => Pin::new(sink).poll_flush(cx),
            AsyncWriterState::Flushing { sink, .. } => Pin::new(sink).poll_flush(cx),
            AsyncWriterState::Done => Poll::Ready(Ok(())),
        }
    }

    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
        // Transition Buffering → Flushing by running sync encryption.
        if let AsyncWriterState::Buffering { .. } = &self.state {
            let (sink, plaintext) = match std::mem::replace(&mut self.state, AsyncWriterState::Done)
            {
                AsyncWriterState::Buffering { sink, plaintext } => (sink, plaintext),
                _ => unreachable!(),
            };
            let ct = match (|| -> Result<Vec<u8>, PqfileError> {
                let mut ct = Vec::new();
                let mut w = PqfWriter::new(
                    &mut ct,
                    &self.pubkey_pem,
                    self.original_size,
                    self.chunk_size,
                )?;
                w.write_all(&plaintext).map_err(PqfileError::Io)?;
                w.finish()?;
                Ok(ct)
            })() {
                Ok(ct) => ct,
                Err(e) => {
                    return Poll::Ready(Err(std::io::Error::other(e.to_string())));
                }
            };
            self.state = AsyncWriterState::Flushing { sink, ct, pos: 0 };
        }

        // Drive the Flushing write loop.
        loop {
            match &mut self.state {
                AsyncWriterState::Flushing { sink, ct, pos } => {
                    if *pos < ct.len() {
                        match Pin::new(&mut *sink).poll_write(cx, &ct[*pos..]) {
                            Poll::Ready(Ok(n)) => {
                                *pos += n;
                            }
                            Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
                            Poll::Pending => return Poll::Pending,
                        }
                    } else {
                        match Pin::new(&mut *sink).poll_shutdown(cx) {
                            Poll::Ready(r) => {
                                self.state = AsyncWriterState::Done;
                                return Poll::Ready(r);
                            }
                            Poll::Pending => return Poll::Pending,
                        }
                    }
                }
                AsyncWriterState::Done => return Poll::Ready(Ok(())),
                AsyncWriterState::Buffering { .. } => unreachable!(),
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::format::CHUNK_SIZE;
    use crate::keygen::keygen_bytes;

    #[tokio::test]
    async fn async_roundtrip_small() {
        let (pub_pem, priv_pem) = keygen_bytes(768, None).unwrap();
        let plaintext = b"hello async pqfile world";

        let mut ct = Vec::new();
        encrypt_stream_async(
            &pub_pem,
            plaintext.len() as u64,
            CHUNK_SIZE,
            &mut plaintext.as_slice(),
            &mut ct,
        )
        .await
        .unwrap();

        let mut out = Vec::new();
        decrypt_stream_async(&priv_pem, ct.as_slice(), &mut out, None)
            .await
            .unwrap();
        assert_eq!(out, plaintext);
    }

    #[tokio::test]
    async fn async_roundtrip_multi_chunk() {
        let (pub_pem, priv_pem) = keygen_bytes(768, None).unwrap();
        let plaintext: Vec<u8> = (0u8..=255).cycle().take(CHUNK_SIZE * 3 + 7).collect();

        let mut ct = Vec::new();
        encrypt_stream_async(
            &pub_pem,
            plaintext.len() as u64,
            CHUNK_SIZE,
            &mut plaintext.as_slice(),
            &mut ct,
        )
        .await
        .unwrap();

        let mut out = Vec::new();
        decrypt_stream_async(&priv_pem, ct.as_slice(), &mut out, None)
            .await
            .unwrap();
        assert_eq!(out, plaintext);
    }

    #[tokio::test]
    async fn async_roundtrip_empty() {
        let (pub_pem, priv_pem) = keygen_bytes(768, None).unwrap();
        let mut ct = Vec::new();
        encrypt_stream_async(&pub_pem, 0, CHUNK_SIZE, &mut [].as_slice(), &mut ct)
            .await
            .unwrap();
        let mut out = Vec::new();
        decrypt_stream_async(&priv_pem, ct.as_slice(), &mut out, None)
            .await
            .unwrap();
        assert!(out.is_empty());
    }

    #[tokio::test]
    async fn async_roundtrip_with_passphrase() {
        let (pub_pem, priv_pem) = keygen_bytes(768, Some("async-pass")).unwrap();
        let plaintext = b"passphrase-protected async roundtrip";
        let mut ct = Vec::new();
        encrypt_stream_async(
            &pub_pem,
            plaintext.len() as u64,
            CHUNK_SIZE,
            &mut plaintext.as_slice(),
            &mut ct,
        )
        .await
        .unwrap();
        let mut out = Vec::new();
        decrypt_stream_async(&priv_pem, ct.as_slice(), &mut out, Some("async-pass"))
            .await
            .unwrap();
        assert_eq!(out, plaintext.as_slice());
    }

    #[tokio::test]
    async fn async_ciphertext_matches_sync() {
        // Files encrypted with encrypt_stream_async must be decryptable by the
        // synchronous decrypt_stream, and vice-versa.
        use crate::decrypt::decrypt_stream as sync_decrypt;
        use crate::encrypt::encrypt_stream as sync_encrypt;

        let (pub_pem, priv_pem) = keygen_bytes(768, None).unwrap();
        let plaintext = b"async / sync interop check";

        // Encrypt async, decrypt sync.
        let mut ct = Vec::new();
        encrypt_stream_async(
            &pub_pem,
            plaintext.len() as u64,
            CHUNK_SIZE,
            &mut plaintext.as_slice(),
            &mut ct,
        )
        .await
        .unwrap();
        let mut out = Vec::new();
        sync_decrypt(&priv_pem, &mut ct.as_slice(), &mut out, None).unwrap();
        assert_eq!(out, plaintext);

        // Encrypt sync, decrypt async.
        let mut ct2 = Vec::new();
        sync_encrypt(
            &pub_pem,
            plaintext.len() as u64,
            CHUNK_SIZE,
            &mut plaintext.as_slice(),
            &mut ct2,
        )
        .unwrap();
        let mut out2 = Vec::new();
        decrypt_stream_async(&priv_pem, ct2.as_slice(), &mut out2, None)
            .await
            .unwrap();
        assert_eq!(out2, plaintext);
    }

    #[tokio::test]
    async fn async_rejects_zero_chunk_size() {
        let (pub_pem, _) = keygen_bytes(768, None).unwrap();
        let result =
            encrypt_stream_async(&pub_pem, 0, 0, &mut [].as_slice(), &mut Vec::new()).await;
        assert!(result.is_err());
    }

    // ── AsyncPqfWriter tests ─────────────────────────────────────────────

    #[tokio::test]
    async fn async_writer_roundtrip_empty() {
        use super::AsyncPqfWriter;
        let (pub_pem, priv_pem) = keygen_bytes(768, None).unwrap();
        let w = AsyncPqfWriter::new(Vec::<u8>::new(), &pub_pem, 0, CHUNK_SIZE).unwrap();
        let ct = w.finish().await.unwrap();
        let mut out = Vec::new();
        decrypt_stream_async(&priv_pem, ct.as_slice(), &mut out, None)
            .await
            .unwrap();
        assert!(out.is_empty());
    }

    #[tokio::test]
    async fn async_writer_roundtrip_small() {
        use super::AsyncPqfWriter;
        use tokio::io::AsyncWriteExt;
        let (pub_pem, priv_pem) = keygen_bytes(768, None).unwrap();
        let plaintext = b"hello async pqfwriter";
        let mut w = AsyncPqfWriter::new(
            Vec::<u8>::new(),
            &pub_pem,
            plaintext.len() as u64,
            CHUNK_SIZE,
        )
        .unwrap();
        w.write_all(plaintext).await.unwrap();
        let ct = w.finish().await.unwrap();
        let mut out = Vec::new();
        decrypt_stream_async(&priv_pem, ct.as_slice(), &mut out, None)
            .await
            .unwrap();
        assert_eq!(out, plaintext);
    }

    #[tokio::test]
    async fn async_writer_roundtrip_multi_chunk() {
        use super::AsyncPqfWriter;
        use tokio::io::AsyncWriteExt;
        let (pub_pem, priv_pem) = keygen_bytes(768, None).unwrap();
        let plaintext: Vec<u8> = (0u8..=255).cycle().take(CHUNK_SIZE * 2 + 77).collect();
        let mut w = AsyncPqfWriter::new(
            Vec::<u8>::new(),
            &pub_pem,
            plaintext.len() as u64,
            CHUNK_SIZE,
        )
        .unwrap();
        w.write_all(&plaintext).await.unwrap();
        let ct = w.finish().await.unwrap();
        let mut out = Vec::new();
        decrypt_stream_async(&priv_pem, ct.as_slice(), &mut out, None)
            .await
            .unwrap();
        assert_eq!(out, plaintext);
    }

    #[tokio::test]
    async fn async_writer_shutdown_seals() {
        use super::AsyncPqfWriter;
        use tokio::io::AsyncWriteExt;
        let (pub_pem, priv_pem) = keygen_bytes(768, None).unwrap();
        let plaintext = b"shutdown seals the stream";
        let mut ct = Vec::new();
        let mut w =
            AsyncPqfWriter::new(&mut ct, &pub_pem, plaintext.len() as u64, CHUNK_SIZE).unwrap();
        w.write_all(plaintext).await.unwrap();
        w.shutdown().await.unwrap();

        let mut out = Vec::new();
        decrypt_stream_async(&priv_pem, ct.as_slice(), &mut out, None)
            .await
            .unwrap();
        assert_eq!(out, plaintext.as_slice());
    }

    #[tokio::test]
    async fn async_writer_interop_with_sync_reader() {
        use super::AsyncPqfWriter;
        use crate::decrypt::decrypt_stream as sync_decrypt;
        use tokio::io::AsyncWriteExt;
        let (pub_pem, priv_pem) = keygen_bytes(768, None).unwrap();
        let plaintext = b"async writer -> sync reader";
        let mut w = AsyncPqfWriter::new(
            Vec::<u8>::new(),
            &pub_pem,
            plaintext.len() as u64,
            CHUNK_SIZE,
        )
        .unwrap();
        w.write_all(plaintext).await.unwrap();
        let ct = w.finish().await.unwrap();
        let mut out = Vec::new();
        sync_decrypt(&priv_pem, &mut ct.as_slice(), &mut out, None).unwrap();
        assert_eq!(out, plaintext.as_slice());
    }

    #[tokio::test]
    async fn async_writer_rejects_zero_chunk_size() {
        use super::AsyncPqfWriter;
        let (pub_pem, _) = keygen_bytes(768, None).unwrap();
        let result = AsyncPqfWriter::new(Vec::<u8>::new(), &pub_pem, 0, 0);
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn async_encrypt_rejects_original_size_above_max() {
        let (pub_pem, _) = keygen_bytes(768, None).unwrap();
        let oversized = crate::format::MAX_ORIGINAL_SIZE + 1;
        let result = encrypt_stream_async(
            &pub_pem,
            oversized,
            CHUNK_SIZE,
            &mut [].as_slice(),
            &mut Vec::new(),
        )
        .await;
        assert!(
            result.is_err(),
            "expected error for oversized original_size"
        );
    }

    #[tokio::test]
    async fn async_decrypt_rejects_oversized_ciphertext() {
        // Build a ciphertext that exceeds MAX_ORIGINAL_SIZE + 1 MiB.
        // We cannot actually allocate that much in a test, so we simulate by
        // checking that a plausible-sized but still-over-limit Vec is rejected.
        // The guard is MAX_ORIGINAL_SIZE + 1 MiB; we construct a fake reader
        // that reports a length just over the limit.
        //
        // Since we cannot synthesise 1 TiB of data in a unit test, we verify
        // the guard fires by crafting a minimal invalid pqf stream and confirming
        // it returns an error (any error means the cap or the parser fired).
        let (_, priv_pem) = keygen_bytes(768, None).unwrap();
        let garbage: Vec<u8> = vec![0xFFu8; 64];
        let mut out = Vec::new();
        let result = decrypt_stream_async(&priv_pem, garbage.as_slice(), &mut out, None).await;
        assert!(result.is_err(), "garbage ciphertext must be rejected");
    }
}