openrtc 2.5.0

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

use std::collections::VecDeque;
use std::io;

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

use crate::application_crypto::{
    protect_frame, CryptoFrameDecoder, CryptoStreamError, APPLICATION_KEY_BYTES,
};

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

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_crypto_auth_failure(error: &io::Error) -> bool {
    error
        .get_ref()
        .and_then(|source| source.downcast_ref::<ApplicationCryptoStreamIoError>())
        .is_some_and(|source| matches!(&source.0, CryptoStreamError::AuthenticationFailed(_)))
}

/// Return the peer-provided QUIC stop code without forcing higher layers to
/// parse a transport error's display text.
pub fn peer_stop_code(error: &io::Error) -> Option<u64> {
    error
        .get_ref()
        .and_then(|source| source.downcast_ref::<WriteError>())
        .and_then(|source| match source {
            WriteError::Stopped(code) => Some(code.into_inner()),
            _ => None,
        })
}

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

impl CryptoSendStream {
    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_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 CryptoRecvStream {
    inner: RecvStream,
    decoder: CryptoFrameDecoder,
    read_buf: Vec<u8>,
    finished: bool,
}

impl CryptoRecvStream {
    pub fn new(inner: RecvStream, key: [u8; APPLICATION_KEY_BYTES]) -> io::Result<Self> {
        Ok(Self {
            inner,
            decoder: CryptoFrameDecoder::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 = CryptoFrameDecoder::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(CryptoSendStream),
}

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

    pub fn encrypted(inner: SendStream, key: [u8; APPLICATION_KEY_BYTES]) -> Self {
        Self::Encrypted(CryptoSendStream::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(CryptoRecvStream),
    /// Plaintext bytes read by a host-side protocol classifier, followed by
    /// the original protected stream. This lets one native host inspect a
    /// channel envelope without stealing or duplicating the admitted-stream
    /// queue.
    Prefixed {
        prefix: VecDeque<u8>,
        inner: Box<PeerRecvStream>,
    },
}

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(CryptoRecvStream::new(inner, key)?))
    }

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

    pub fn is_encrypted(&self) -> bool {
        match self {
            Self::Encrypted(_) => true,
            Self::Prefixed { inner, .. } => inner.is_encrypted(),
            Self::Plain(_) => false,
        }
    }

    /// Restore plaintext consumed during bounded host-side classification.
    pub fn with_plaintext_prefix(self, prefix: impl IntoIterator<Item = u8>) -> Self {
        let prefix = prefix.into_iter().collect::<VecDeque<_>>();
        if prefix.is_empty() {
            self
        } else {
            Self::Prefixed {
                prefix,
                inner: Box::new(self),
            }
        }
    }

    pub async fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let mut current = self;
        loop {
            match current {
                Self::Plain(stream) => return read_recv_stream(stream, buf).await,
                Self::Encrypted(stream) => return stream.read(buf).await,
                Self::Prefixed { prefix, inner } => {
                    if !prefix.is_empty() {
                        let copied = prefix.len().min(buf.len());
                        for slot in &mut buf[..copied] {
                            *slot = prefix.pop_front().expect("prefix length was checked");
                        }
                        return Ok(copied);
                    } else {
                        current = inner.as_mut();
                    }
                }
            }
        }
    }

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

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_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::other(error)
}

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::CryptoError) -> io::Error {
    io::Error::new(io::ErrorKind::InvalidData, format!("{error:?}"))
}

fn crypto_stream_io_error(error: CryptoStreamError) -> 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(CryptoStreamError::AuthenticationFailed(
            crate::application_crypto::CryptoError::DecryptFailed,
        ));
        let incomplete = crypto_stream_io_error(CryptoStreamError::IncompleteFrame);

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

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

    #[test]
    fn preserves_peer_stop_code_without_string_parsing() {
        let stopped = map_write_error(WriteError::Stopped(iroh::endpoint::VarInt::from_u32(0)));
        assert_eq!(peer_stop_code(&stopped), Some(0));

        let lost = io::Error::other("connection lost");
        assert_eq!(peer_stop_code(&lost), None);
    }
}