openrtc 1.0.4

//! Length-framed application-crypto adapters for iroh `SendStream` / `RecvStream`.
//!
//! Mirrors the TypeScript `createApplicationCryptoStreamTools` contract: each
//! `write_all` produces one encrypted frame; `read` reassembles framed ciphertext
//! from arbitrary chunk boundaries before returning plaintext.

use std::io;

use iroh::endpoint::{ClosedStream, ReadError, RecvStream, SendStream, WriteError};
use tokio::io::AsyncWriteExt;

use crate::application_crypto::{
    protect_raw_stream_frame, ApplicationCryptoFrameDecoder, ApplicationCryptoStreamError,
    APPLICATION_KEY_BYTES,
};

#[derive(Debug)]
struct ApplicationCryptoStreamIoError(ApplicationCryptoStreamError);

impl std::fmt::Display for ApplicationCryptoStreamIoError {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(formatter, "{:?}", self.0)
    }
}

impl std::error::Error for ApplicationCryptoStreamIoError {}

/// Returns true only when encrypted stream decoding reached AEAD verification
/// and rejected the frame. Transport EOF, connection loss, and incomplete
/// shutdown frames are intentionally distinct so callers do not misreport
/// lifecycle teardown as retired key material.
pub fn is_application_crypto_authentication_failure(error: &io::Error) -> bool {
    error
        .get_ref()
        .and_then(|source| source.downcast_ref::<ApplicationCryptoStreamIoError>())
        .is_some_and(|source| {
            matches!(
                &source.0,
                ApplicationCryptoStreamError::AuthenticationFailed(_)
            )
        })
}

#[derive(Debug)]
pub struct ApplicationCryptoSendStream {
    inner: SendStream,
    key: [u8; APPLICATION_KEY_BYTES],
}

impl ApplicationCryptoSendStream {
    pub fn new(inner: SendStream, key: [u8; APPLICATION_KEY_BYTES]) -> Self {
        Self { inner, key }
    }

    pub async fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
        let frame = protect_raw_stream_frame(&self.key, buf).map_err(crypto_io_error)?;
        self.inner.write_all(&frame).await.map_err(map_write_error)
    }

    pub async fn flush(&mut self) -> io::Result<()> {
        self.inner.flush().await
    }

    pub fn finish(mut self) -> io::Result<()> {
        self.inner.finish().map_err(map_finish_error)
    }

    pub fn into_inner(self) -> SendStream {
        self.inner
    }
}

#[derive(Debug)]
pub struct ApplicationCryptoRecvStream {
    inner: RecvStream,
    decoder: ApplicationCryptoFrameDecoder,
    read_buf: Vec<u8>,
    finished: bool,
}

impl ApplicationCryptoRecvStream {
    pub fn new(inner: RecvStream, key: [u8; APPLICATION_KEY_BYTES]) -> io::Result<Self> {
        Ok(Self {
            inner,
            decoder: ApplicationCryptoFrameDecoder::new(&key).map_err(crypto_io_error)?,
            read_buf: Vec::new(),
            finished: false,
        })
    }

    pub fn new_with_prefix(
        inner: RecvStream,
        key: [u8; APPLICATION_KEY_BYTES],
        prefix: &[u8],
    ) -> io::Result<Self> {
        let mut decoder = ApplicationCryptoFrameDecoder::new(&key).map_err(crypto_io_error)?;
        let mut read_buf = Vec::new();
        if !prefix.is_empty() {
            for frame in decoder.push(prefix).map_err(crypto_stream_io_error)? {
                read_buf.extend_from_slice(&frame);
            }
        }
        Ok(Self {
            inner,
            decoder,
            read_buf,
            finished: false,
        })
    }

    pub async fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        if !self.read_buf.is_empty() {
            let copied = self.read_buf.len().min(buf.len());
            buf[..copied].copy_from_slice(&self.read_buf[..copied]);
            self.read_buf.drain(..copied);
            return Ok(copied);
        }

        if self.finished {
            return Ok(0);
        }

        loop {
            let mut chunk = vec![0u8; 16 * 1024];
            let read_bytes = match self.inner.read(&mut chunk).await {
                Ok(Some(0)) => 0,
                Ok(Some(read_bytes)) => read_bytes,
                Ok(None) => 0,
                Err(error) => return Err(map_read_error(error)),
            };
            if read_bytes == 0 {
                self.finished = true;
                self.decoder.finish().map_err(crypto_stream_io_error)?;
                if self.read_buf.is_empty() {
                    return Ok(0);
                }
                let copied = self.read_buf.len().min(buf.len());
                buf[..copied].copy_from_slice(&self.read_buf[..copied]);
                self.read_buf.drain(..copied);
                return Ok(copied);
            }

            chunk.truncate(read_bytes);
            let opened = self.decoder.push(&chunk).map_err(crypto_stream_io_error)?;
            for frame in opened {
                self.read_buf.extend_from_slice(&frame);
            }

            if !self.read_buf.is_empty() {
                let copied = self.read_buf.len().min(buf.len());
                buf[..copied].copy_from_slice(&self.read_buf[..copied]);
                self.read_buf.drain(..copied);
                return Ok(copied);
            }
        }
    }

    pub fn into_inner(self) -> RecvStream {
        self.inner
    }
}

#[derive(Debug)]
pub enum PeerSendStream {
    Plain(SendStream),
    Encrypted(ApplicationCryptoSendStream),
}

impl PeerSendStream {
    pub fn plain(inner: SendStream) -> Self {
        Self::Plain(inner)
    }

    pub fn encrypted(inner: SendStream, key: [u8; APPLICATION_KEY_BYTES]) -> Self {
        Self::Encrypted(ApplicationCryptoSendStream::new(inner, key))
    }

    pub fn is_encrypted(&self) -> bool {
        matches!(self, Self::Encrypted(_))
    }

    pub async fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
        match self {
            Self::Plain(stream) => stream.write_all(buf).await.map_err(map_write_error),
            Self::Encrypted(stream) => stream.write_all(buf).await,
        }
    }

    pub async fn flush(&mut self) -> io::Result<()> {
        match self {
            Self::Plain(stream) => stream.flush().await,
            Self::Encrypted(stream) => stream.flush().await,
        }
    }

    pub fn finish(self) -> io::Result<()> {
        match self {
            Self::Plain(mut stream) => stream.finish().map_err(map_finish_error),
            Self::Encrypted(stream) => stream.finish(),
        }
    }

    /// Finish a short control response and keep its QUIC send half alive until
    /// the peer has acknowledged the buffered bytes (or the bounded wait
    /// expires). `finish()` alone only queues the FIN; a concurrent physical-leg
    /// replacement can otherwise close the connection before the verdict is
    /// observable by the admission dialer.
    pub async fn finish_and_wait_for_peer(self, timeout: std::time::Duration) -> io::Result<()> {
        async fn finish_and_wait(
            mut stream: SendStream,
            timeout: std::time::Duration,
        ) -> io::Result<()> {
            stream.finish().map_err(map_finish_error)?;
            #[cfg(not(target_arch = "wasm32"))]
            match tokio::time::timeout(timeout, stream.stopped()).await {
                Ok(Ok(_)) => Ok(()),
                Ok(Err(error)) => Err(io::Error::other(format!(
                    "wait for peer acknowledgement: {error}"
                ))),
                Err(_) => Err(io::Error::new(
                    io::ErrorKind::TimedOut,
                    "timed out waiting for peer acknowledgement",
                )),
            }
            #[cfg(target_arch = "wasm32")]
            {
                use futures::FutureExt;

                let stopped = stream.stopped().fuse();
                let timeout = gloo_timers::future::sleep(timeout).fuse();
                futures::pin_mut!(stopped, timeout);
                futures::select! {
                    result = stopped => result.map(|_| ()).map_err(|error| {
                        io::Error::other(format!("wait for peer acknowledgement: {error}"))
                    }),
                    _ = timeout => Err(io::Error::new(
                        io::ErrorKind::TimedOut,
                        "timed out waiting for peer acknowledgement",
                    )),
                }
            }
        }

        match self {
            Self::Plain(stream) => finish_and_wait(stream, timeout).await,
            Self::Encrypted(stream) => finish_and_wait(stream.into_inner(), timeout).await,
        }
    }

    pub fn into_plain(self) -> SendStream {
        match self {
            Self::Plain(stream) => stream,
            Self::Encrypted(stream) => stream.into_inner(),
        }
    }
}

#[derive(Debug)]
pub enum PeerRecvStream {
    Plain(RecvStream),
    Encrypted(ApplicationCryptoRecvStream),
}

impl PeerRecvStream {
    pub fn plain(inner: RecvStream) -> Self {
        Self::Plain(inner)
    }

    pub fn encrypted(inner: RecvStream, key: [u8; APPLICATION_KEY_BYTES]) -> io::Result<Self> {
        Ok(Self::Encrypted(ApplicationCryptoRecvStream::new(
            inner, key,
        )?))
    }

    pub fn encrypted_with_prefix(
        inner: RecvStream,
        key: [u8; APPLICATION_KEY_BYTES],
        prefix: &[u8],
    ) -> io::Result<Self> {
        Ok(Self::Encrypted(
            ApplicationCryptoRecvStream::new_with_prefix(inner, key, prefix)?,
        ))
    }

    pub fn is_encrypted(&self) -> bool {
        matches!(self, Self::Encrypted(_))
    }

    pub async fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        match self {
            Self::Plain(stream) => read_recv_stream(stream, buf).await,
            Self::Encrypted(stream) => stream.read(buf).await,
        }
    }

    pub fn into_plain(self) -> RecvStream {
        match self {
            Self::Plain(stream) => stream,
            Self::Encrypted(stream) => stream.into_inner(),
        }
    }
}

pub fn wrap_peer_streams(
    key: Option<[u8; APPLICATION_KEY_BYTES]>,
    send: SendStream,
    recv: RecvStream,
) -> io::Result<(PeerSendStream, PeerRecvStream)> {
    match key {
        Some(key) => Ok((
            PeerSendStream::encrypted(send, key),
            PeerRecvStream::encrypted(recv, key)?,
        )),
        None => Ok((PeerSendStream::plain(send), PeerRecvStream::plain(recv))),
    }
}

pub fn wrap_peer_send_stream(
    key: Option<[u8; APPLICATION_KEY_BYTES]>,
    send: SendStream,
) -> PeerSendStream {
    match key {
        Some(key) => PeerSendStream::encrypted(send, key),
        None => PeerSendStream::plain(send),
    }
}

async fn read_recv_stream(stream: &mut RecvStream, buf: &mut [u8]) -> io::Result<usize> {
    match stream.read(buf).await {
        Ok(Some(0)) => Ok(0),
        Ok(Some(read_bytes)) => Ok(read_bytes),
        Ok(None) => Ok(0),
        Err(error) => Err(map_read_error(error)),
    }
}

fn map_write_error(error: WriteError) -> io::Error {
    io::Error::new(io::ErrorKind::Other, error.to_string())
}

fn map_read_error(error: ReadError) -> io::Error {
    io::Error::new(io::ErrorKind::Other, error.to_string())
}

fn map_finish_error(error: ClosedStream) -> io::Error {
    io::Error::new(io::ErrorKind::NotConnected, error.to_string())
}

fn crypto_io_error(error: crate::application_crypto::ApplicationCryptoError) -> io::Error {
    io::Error::new(io::ErrorKind::InvalidData, format!("{error:?}"))
}

fn crypto_stream_io_error(error: ApplicationCryptoStreamError) -> io::Error {
    io::Error::new(
        io::ErrorKind::InvalidData,
        ApplicationCryptoStreamIoError(error),
    )
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn distinguishes_authentication_failure_from_incomplete_shutdown_frame() {
        let authentication =
            crypto_stream_io_error(ApplicationCryptoStreamError::AuthenticationFailed(
                crate::application_crypto::ApplicationCryptoError::DecryptFailed,
            ));
        let incomplete = crypto_stream_io_error(ApplicationCryptoStreamError::IncompleteFrame);

        assert!(is_application_crypto_authentication_failure(
            &authentication
        ));
        assert!(!is_application_crypto_authentication_failure(&incomplete));
    }

    #[test]
    fn transport_errors_are_not_crypto_authentication_failures() {
        let transport = io::Error::new(io::ErrorKind::ConnectionReset, "connection lost");
        assert!(!is_application_crypto_authentication_failure(&transport));
    }
}