openrtc 2.5.2

//! 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 std::sync::{
    atomic::{AtomicBool, Ordering},
    Arc,
};

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 {}

/// Stream-local access fence, supplied only by the Rust admission/key owner.
/// It observes existing authority, never schedules recovery or retires peers.
#[derive(Clone)]
pub(crate) struct ApplicationStreamAccess {
    current: Arc<dyn Fn() -> bool + Send + Sync>,
    denied: Arc<AtomicBool>,
}

impl std::fmt::Debug for ApplicationStreamAccess {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str("ApplicationStreamAccess")
    }
}

impl ApplicationStreamAccess {
    pub(crate) fn new(current: impl Fn() -> bool + Send + Sync + 'static) -> Self {
        Self {
            current: Arc::new(current),
            denied: Arc::new(AtomicBool::new(false)),
        }
    }

    pub(crate) fn check(&self) -> io::Result<()> {
        if !self.denied.load(Ordering::Acquire) && (self.current)() {
            return Ok(());
        }
        // A stream which observed denial cannot resurrect after reauthorization.
        self.denied.store(true, Ordering::Release);
        Err(io::Error::new(
            io::ErrorKind::PermissionDenied,
            "application stream authorization expired or changed",
        ))
    }
}

fn check_access(access: &Option<ApplicationStreamAccess>) -> io::Result<()> {
    access
        .as_ref()
        .map_or(Ok(()), ApplicationStreamAccess::check)
}

/// Recheck before every poll, including when transport backpressure wakes a
/// pending operation. Checking only before `.await` would authorize later I/O.
async fn with_access<T>(
    access: &Option<ApplicationStreamAccess>,
    operation: impl std::future::Future<Output = io::Result<T>>,
) -> io::Result<T> {
    let mut operation = std::pin::pin!(operation);
    std::future::poll_fn(|cx| {
        check_access(access)?;
        let result = operation.as_mut().poll(cx);
        check_access(access)?;
        result
    })
    .await
}

/// 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],
    access: Option<ApplicationStreamAccess>,
}

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

    pub async fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
        check_access(&self.access)?;
        let frame = protect_frame(&self.key, buf).map_err(crypto_io_error)?;
        // Bound each authorization-to-write interval even for a large frame.
        for chunk in frame.chunks(16 * 1024) {
            with_access(&self.access, async {
                self.inner.write_all(chunk).await.map_err(map_write_error)
            })
            .await?;
        }
        Ok(())
    }

    pub async fn flush(&mut self) -> io::Result<()> {
        with_access(&self.access, 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,
    access: Option<ApplicationStreamAccess>,
}

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,
            access: None,
        })
    }

    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,
            access: None,
        })
    }

    pub async fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        check_access(&self.access)?;
        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 with_access(&self.access, async {
                self.inner.read(&mut chunk).await.map_err(map_read_error)
            })
            .await
            {
                Ok(Some(0)) => 0,
                Ok(Some(read_bytes)) => read_bytes,
                Ok(None) => 0,
                Err(error) => return Err(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);
                }
                check_access(&self.access)?;
                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() {
                check_access(&self.access)?;
                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(crate) fn with_access(mut self, access: ApplicationStreamAccess) -> Self {
        if let Self::Encrypted(stream) = &mut self {
            stream.access = Some(access);
        }
        self
    }
    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(crate) fn with_access(mut self, access: ApplicationStreamAccess) -> Self {
        let mut current = &mut self;
        loop {
            match current {
                Self::Encrypted(stream) => {
                    stream.access = Some(access);
                    break;
                }
                Self::Prefixed { inner, .. } => current = inner.as_mut(),
                Self::Plain(_) => break,
            }
        }
        self
    }

    fn check_access(&self) -> io::Result<()> {
        let mut current = self;
        loop {
            match current {
                Self::Encrypted(stream) => return check_access(&stream.access),
                Self::Prefixed { inner, .. } => current = inner.as_ref(),
                Self::Plain(_) => return Ok(()),
            }
        }
    }
    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> {
        // Restored plaintext prefixes obey the same fence as encrypted bytes.
        self.check_access()?;
        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 resumed_backpressure_cannot_poll_io_after_access_is_lost() {
        use std::{
            cell::Cell,
            future::Future,
            task::{Context, Poll, Waker},
        };
        let allowed = Arc::new(AtomicBool::new(true));
        let observed = allowed.clone();
        let access = Some(ApplicationStreamAccess::new(move || {
            observed.load(Ordering::Acquire)
        }));
        let polls = Cell::new(0);
        let pending = std::future::poll_fn(|_| {
            polls.set(polls.get() + 1);
            Poll::<io::Result<()>>::Pending
        });
        let mut guarded = std::pin::pin!(with_access(&access, pending));
        let mut cx = Context::from_waker(Waker::noop());
        assert!(guarded.as_mut().poll(&mut cx).is_pending());
        allowed.store(false, Ordering::Release);
        assert!(
            matches!(guarded.as_mut().poll(&mut cx), Poll::Ready(Err(error)) if error.kind() == io::ErrorKind::PermissionDenied)
        );
        assert_eq!(
            polls.get(),
            1,
            "a resumed write must not enqueue unauthorized bytes"
        );
        allowed.store(true, Ordering::Release);
        assert!(
            access.as_ref().unwrap().check().is_err(),
            "denied streams stay closed after reauthorization"
        );
    }

    #[test]
    fn access_lost_during_io_does_not_return_its_result() {
        use std::{
            future::Future,
            task::{Context, Poll, Waker},
        };
        let allowed = Arc::new(AtomicBool::new(true));
        let observed = allowed.clone();
        let access = Some(ApplicationStreamAccess::new(move || {
            observed.load(Ordering::Acquire)
        }));
        let operation = async {
            allowed.store(false, Ordering::Release);
            Ok(b"private bytes")
        };
        let mut guarded = std::pin::pin!(with_access(&access, operation));
        assert!(
            matches!(guarded.as_mut().poll(&mut Context::from_waker(Waker::noop())), Poll::Ready(Err(error)) if error.kind() == io::ErrorKind::PermissionDenied)
        );
    }

    #[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);
    }
}