qmux 0.5.0

use bytes::Bytes;

use crate::Error;

/// Abstracts message I/O over a reliable transport.
///
/// Each `send`/`recv` operates on a single complete message (frame).
/// For WebSocket, this maps to individual WS binary messages.
/// For TCP/TLS byte streams, the transport handles frame delimiting.
///
/// Use [`Stream`] for arbitrary [`AsyncRead`](tokio::io::AsyncRead) +
/// [`AsyncWrite`](tokio::io::AsyncWrite) byte streams. Implement this trait when
/// integrating a transport with its own message boundaries, framing, or
/// lifecycle that the built-in byte-stream and WebSocket adapters do not cover.
///
/// A transport splits into independently-owned send and receive halves so the
/// session can drive them on separate tasks: a write blocked on transport
/// backpressure must never stall reads (and vice versa). This decoupling is what
/// lets the session observe backpressure and shed unreliable datagrams instead
/// of buffering them behind a stalled socket.
pub trait Transport: Send + 'static {
    /// The independently-owned send half.
    type Writer: Writer;
    /// The independently-owned receive half.
    type Reader: Reader;

    /// Split into send and receive halves.
    fn split(self) -> (Self::Writer, Self::Reader);
}

/// The send half of a [`Transport`].
pub trait Writer: Send + 'static {
    /// Send a single complete message.
    fn send(&mut self, data: Bytes) -> impl std::future::Future<Output = Result<(), Error>> + Send;

    /// Gracefully close the transport.
    fn close(&mut self) -> impl std::future::Future<Output = Result<(), Error>> + Send;

    /// Perform any timer-driven background work and resolve once it's done. The
    /// session's writer loop selects on this alongside outbound frames, so a
    /// transport can piggy-back periodic maintenance (e.g. a WebSocket keep-alive
    /// Ping) on the same task that owns the send half. The default never
    /// resolves — transports with nothing to do (TCP, Unix sockets) use it as-is.
    fn maintain(&mut self) -> impl std::future::Future<Output = Result<(), Error>> + Send {
        std::future::pending()
    }
}

/// The receive half of a [`Transport`].
pub trait Reader: Send + 'static {
    /// Receive the next complete message.
    fn recv(&mut self) -> impl std::future::Future<Output = Result<Bytes, Error>> + Send;
}

// Stream: message I/O over a byte stream (TCP/TLS/Unix).
// Handles QMux frame delimiting to return complete frames as Bytes.
//
// Cancel safety: a dedicated reader task owns the read half and pushes complete
// frames into an `mpsc` channel. `recv()` is just `rx.recv().await`, which is
// cancel safe — if the future is dropped (e.g. a sibling `tokio::select!` branch
// wins), the buffered frame stays in the channel for the next call. The reader
// task itself never gets cancelled mid-parse, so the multi-step async reads in
// `recv_record`/`recv_qmux00_frame` are safe to keep as-is.
#[cfg(any(feature = "tcp", all(unix, feature = "uds")))]
mod stream_transport {
    use bytes::{BufMut, Bytes, BytesMut};
    use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufReader, BufWriter};
    use tokio::sync::mpsc;
    use tokio::task::JoinHandle;
    use web_transport_proto::VarInt;

    use super::{Reader, Transport, Writer};
    use crate::{Error, Version, MAX_FRAME_SIZE};

    /// Bound on queued frames waiting for the session to drain them. Bytes the
    /// session hasn't picked up yet are also buffered in the OS receive window
    /// once this fills; the channel just gives a small slack so each `recv()`
    /// is a cheap hand-off rather than a syscall.
    const RECV_CHANNEL_CAPACITY: usize = 16;

    /// QMux message I/O over any reliable byte stream (`AsyncRead + AsyncWrite`).
    ///
    /// Handles QMux frame/record delimiting so [`Session`](crate::Session) sees
    /// complete frames. Pair it with [`Session::connect`](crate::Session::connect)
    /// or [`Session::accept`](crate::Session::accept) to run QMux over a transport
    /// the built-in `tcp`/`tls`/`ws` helpers don't cover — a Unix socket, a pipe,
    /// an in-memory duplex, a custom tunnel, etc.:
    ///
    /// ```no_run
    /// # async fn f(stream: tokio::net::TcpStream) -> Result<(), qmux::Error> {
    /// use qmux::transport::Stream;
    /// use qmux::{Config, Session, Version};
    ///
    /// let config = Config::new(Version::QMux01);
    /// let transport = Stream::new(stream, config.version, config.max_record_size);
    /// let session = Session::connect(transport, config).await?;
    /// # let _ = session; Ok(())
    /// # }
    /// ```
    pub struct Stream<T> {
        writer: StreamWriter<T>,
        reader: StreamReader,
    }

    /// The send half of a byte-stream [`Stream`].
    pub struct StreamWriter<T> {
        writer: BufWriter<tokio::io::WriteHalf<T>>,
        version: Version,
    }

    /// The receive half of a byte-stream [`Stream`]. Owns the reader task's
    /// abort handle so the task can't outlive the receive half.
    pub struct StreamReader {
        rx: mpsc::Receiver<Result<Bytes, Error>>,
        /// Aborted on drop so the reader task can't outlive the transport.
        reader_task: JoinHandle<()>,
    }

    impl<T: AsyncRead + AsyncWrite + Send + 'static> Stream<T> {
        /// Wrap a byte stream speaking QMux `version`.
        ///
        /// `our_max_record_size` bounds incoming draft-01 records (use
        /// [`Config::max_record_size`](crate::Config::max_record_size)); it is
        /// ignored for draft-00 and the legacy `webtransport` wire format.
        pub fn new(stream: T, version: Version, our_max_record_size: u64) -> Self {
            let (read, write) = tokio::io::split(stream);
            let (tx, rx) = mpsc::channel(RECV_CHANNEL_CAPACITY);
            let reader_task = tokio::spawn(reader_loop(
                BufReader::new(read),
                version,
                our_max_record_size as usize,
                tx,
            ));
            Self {
                writer: StreamWriter {
                    writer: BufWriter::new(write),
                    version,
                },
                reader: StreamReader { rx, reader_task },
            }
        }
    }

    impl<T: AsyncRead + AsyncWrite + Send + 'static> Transport for Stream<T> {
        type Writer = StreamWriter<T>;
        type Reader = StreamReader;

        fn split(self) -> (StreamWriter<T>, StreamReader) {
            (self.writer, self.reader)
        }
    }

    impl Drop for StreamReader {
        fn drop(&mut self) {
            // Make sure the reader task doesn't outlive the transport; otherwise
            // it would hold the read half open until the connection drops.
            self.reader_task.abort();
        }
    }

    impl<T: AsyncWrite + Send + 'static> Writer for StreamWriter<T> {
        async fn send(&mut self, data: Bytes) -> Result<(), Error> {
            // Record-framed drafts (QMux01+) travel inside size-prefixed records
            // on byte streams. (Records are implicit on WebSocket, where the
            // message boundary delimits them.)
            if self.version.uses_records() {
                let mut size_buf = BytesMut::with_capacity(8);
                VarInt::try_from(data.len())?.encode(&mut size_buf);
                self.writer.write_all(&size_buf).await?;
            }
            self.writer.write_all(&data).await?;
            self.writer.flush().await?;
            Ok(())
        }

        async fn close(&mut self) -> Result<(), Error> {
            self.writer.shutdown().await?;
            Ok(())
        }
    }

    impl Reader for StreamReader {
        async fn recv(&mut self) -> Result<Bytes, Error> {
            // mpsc::Receiver::recv is cancel safe, so dropping this future never
            // loses a buffered frame. `None` means the reader task exited without
            // sending — treat as a clean close.
            self.rx.recv().await.unwrap_or(Err(Error::Closed))
        }
    }

    /// Reader task: pull complete frames off the wire and ship them through `tx`.
    /// On parse error, send the error and exit. If `tx` is closed (the transport
    /// was dropped), exit silently.
    async fn reader_loop<R: AsyncRead + Unpin>(
        mut reader: BufReader<R>,
        version: Version,
        our_max_record_size: usize,
        tx: mpsc::Sender<Result<Bytes, Error>>,
    ) {
        loop {
            let result = match version {
                Version::QMux01 | Version::QMux02 => {
                    recv_record(&mut reader, our_max_record_size).await
                }
                Version::QMux00 | Version::WebTransport => recv_qmux00_frame(&mut reader).await,
            };
            let stop = result.is_err();
            if tx.send(result).await.is_err() {
                return;
            }
            if stop {
                return;
            }
        }
    }

    /// Read a varint from the stream, returning the decoded value.
    /// If `buf` is provided, appends the raw bytes to it.
    async fn read_varint_into<R: AsyncRead + Unpin>(
        reader: &mut R,
        buf: &mut BytesMut,
    ) -> Result<VarInt, Error> {
        let first = reader.read_u8().await?;
        buf.put_u8(first);

        let tag = first >> 6;
        let len = 1usize << tag;

        if len == 1 {
            return Ok(VarInt::try_from((first & 0x3f) as u64).unwrap());
        }

        let start = buf.len();
        buf.resize(start + len - 1, 0);
        reader.read_exact(&mut buf[start..]).await?;

        let mut raw = [0u8; 8];
        raw[0] = first & 0x3f;
        raw[1..len].copy_from_slice(&buf[start..start + len - 1]);

        let value = match len {
            2 => u16::from_be_bytes([raw[0], raw[1]]) as u64,
            4 => u32::from_be_bytes([raw[0], raw[1], raw[2], raw[3]]) as u64,
            8 => u64::from_be_bytes(raw),
            _ => unreachable!(),
        };

        VarInt::try_from(value).map_err(|_| Error::Short)
    }

    /// Read a varint from the stream without collecting raw bytes.
    async fn read_varint<R: AsyncRead + Unpin>(reader: &mut R) -> Result<VarInt, Error> {
        let first = reader.read_u8().await?;
        let tag = first >> 6;
        let len = 1usize << tag;

        if len == 1 {
            return Ok(VarInt::try_from((first & 0x3f) as u64).unwrap());
        }

        let mut raw = [0u8; 8];
        raw[0] = first & 0x3f;
        reader.read_exact(&mut raw[1..len]).await?;

        let value = match len {
            2 => u16::from_be_bytes([raw[0], raw[1]]) as u64,
            4 => u32::from_be_bytes([raw[0], raw[1], raw[2], raw[3]]) as u64,
            8 => u64::from_be_bytes(raw),
            _ => unreachable!(),
        };

        VarInt::try_from(value).map_err(|_| Error::Short)
    }

    /// Read exactly `len` bytes, appending to buf.
    async fn read_bytes<R: AsyncRead + Unpin>(
        reader: &mut R,
        len: usize,
        buf: &mut BytesMut,
    ) -> Result<(), Error> {
        let start = buf.len();
        buf.resize(start + len, 0);
        reader.read_exact(&mut buf[start..]).await?;
        Ok(())
    }

    /// Read one QMux Record from the byte stream (draft-01).
    /// Returns the record payload (frames concatenated).
    async fn recv_record<R: AsyncRead + Unpin>(
        reader: &mut R,
        our_max_record_size: usize,
    ) -> Result<Bytes, Error> {
        let size = read_varint(reader).await?.into_inner() as usize;
        if size > our_max_record_size {
            return Err(Error::FrameTooLarge);
        }
        let mut buf = BytesMut::zeroed(size);
        reader.read_exact(&mut buf).await?;
        Ok(buf.freeze())
    }

    /// Read one complete QMux frame from the byte stream (draft-00), returning raw bytes.
    async fn recv_qmux00_frame<R: AsyncRead + Unpin>(reader: &mut R) -> Result<Bytes, Error> {
        let mut buf = BytesMut::new();
        let frame_type = read_varint_into(reader, &mut buf).await?.into_inner();

        // STREAM frames: 0x08-0x0f
        if (0x08..=0x0f).contains(&frame_type) {
            let has_off = frame_type & 0x04 != 0;
            let has_len = frame_type & 0x02 != 0;

            read_varint_into(reader, &mut buf).await?; // stream id

            if has_off {
                read_varint_into(reader, &mut buf).await?; // offset
            }

            if has_len {
                let len = read_varint_into(reader, &mut buf).await?.into_inner() as usize;
                // draft-00 §5.2: `max_frame_size` bounds the whole frame, so the
                // header bytes read so far count against it. A byte stream cannot
                // resynchronize past an oversized frame, so this stays fatal — as
                // the draft requires (FRAME_ENCODING_ERROR).
                if buf.len() + len > MAX_FRAME_SIZE {
                    return Err(Error::FrameTooLarge);
                }
                read_bytes(reader, len, &mut buf).await?;
            } else {
                return Err(Error::Short);
            }

            return Ok(buf.freeze());
        }

        match frame_type {
            // PADDING
            0x00 => {}
            // RESET_STREAM
            0x04 => {
                read_varint_into(reader, &mut buf).await?; // id
                read_varint_into(reader, &mut buf).await?; // code
                read_varint_into(reader, &mut buf).await?; // final_size
            }
            // STOP_SENDING
            0x05 => {
                read_varint_into(reader, &mut buf).await?; // id
                read_varint_into(reader, &mut buf).await?; // code
            }
            // CONNECTION_CLOSE (0x1c) carries the Frame Type field; APPLICATION_CLOSE
            // (0x1d) omits it (RFC 9000 §19.19).
            0x1c | 0x1d => {
                read_varint_into(reader, &mut buf).await?; // code
                if frame_type == 0x1c {
                    read_varint_into(reader, &mut buf).await?; // frame_type
                }
                let reason_len = read_varint_into(reader, &mut buf).await?.into_inner() as usize;
                if buf.len() + reason_len > MAX_FRAME_SIZE {
                    return Err(Error::FrameTooLarge);
                }
                read_bytes(reader, reason_len, &mut buf).await?;
            }
            // MAX_DATA
            0x10 => {
                read_varint_into(reader, &mut buf).await?;
            }
            // MAX_STREAM_DATA
            0x11 => {
                read_varint_into(reader, &mut buf).await?; // id
                read_varint_into(reader, &mut buf).await?; // max
            }
            // MAX_STREAMS (bidi/uni)
            0x12 | 0x13 => {
                read_varint_into(reader, &mut buf).await?;
            }
            // DATA_BLOCKED
            0x14 => {
                read_varint_into(reader, &mut buf).await?;
            }
            // STREAM_DATA_BLOCKED
            0x15 => {
                read_varint_into(reader, &mut buf).await?; // id
                read_varint_into(reader, &mut buf).await?; // limit
            }
            // STREAMS_BLOCKED (bidi/uni)
            0x16 | 0x17 => {
                read_varint_into(reader, &mut buf).await?;
            }
            // DATAGRAM without length — can't delimit on a byte stream
            0x30 => return Err(Error::InvalidFrameType(frame_type)),
            // DATAGRAM with length
            0x31 => {
                let len = read_varint_into(reader, &mut buf).await?.into_inner() as usize;
                if buf.len() + len > MAX_FRAME_SIZE {
                    return Err(Error::FrameTooLarge);
                }
                read_bytes(reader, len, &mut buf).await?;
            }
            // QX_TRANSPORT_PARAMETERS
            0x3f5153300d0a0d0a => {
                let len = read_varint_into(reader, &mut buf).await?.into_inner() as usize;
                if buf.len() + len > MAX_FRAME_SIZE {
                    return Err(Error::FrameTooLarge);
                }
                read_bytes(reader, len, &mut buf).await?;
            }
            // QX_PING request/response (also valid in draft-00 for forward compat)
            0x348c67529ef8c7bd | 0x348c67529ef8c7be => {
                read_varint_into(reader, &mut buf).await?; // sequence
            }
            _ => return Err(Error::InvalidFrameType(frame_type)),
        }

        Ok(buf.freeze())
    }

    #[cfg(test)]
    mod tests {
        use super::*;
        use crate::transport::{Reader, Transport};
        use tokio::io::AsyncWriteExt;

        // Drip a frame in one byte at a time, racing each `recv` against an
        // immediate yield so the recv future is dropped every iteration. With
        // a cancel-safe `recv`, the final call must still return the whole
        // frame intact.
        #[tokio::test]
        async fn recv_is_cancel_safe_across_partial_writes() {
            let (client, mut server) = tokio::io::duplex(64 * 1024);
            let (_writer, mut transport) = Stream::new(client, Version::QMux00, 16 * 1024).split();

            // STREAM frame, type 0x0a (len bit set), id=4, length=5, payload="hello".
            let mut frame = Vec::new();
            frame.push(0x0a);
            frame.push(0x04);
            frame.push(0x05);
            frame.extend_from_slice(b"hello");

            for chunk in frame.chunks(1).take(frame.len() - 1) {
                server.write_all(chunk).await.unwrap();
                server.flush().await.unwrap();
                tokio::select! {
                    _ = transport.recv() => panic!("recv completed with a partial frame"),
                    _ = tokio::task::yield_now() => {}
                }
            }

            server.write_all(&frame[frame.len() - 1..]).await.unwrap();
            server.flush().await.unwrap();

            let got = transport.recv().await.expect("frame should decode");
            assert_eq!(&got[..], frame.as_slice());
        }

        #[tokio::test]
        async fn recv_qmux01_record_is_cancel_safe() {
            let (client, mut server) = tokio::io::duplex(64 * 1024);
            let (_writer, mut transport) = Stream::new(client, Version::QMux01, 16 * 1024).split();

            // 1-byte varint length (0x08) followed by 8 bytes of payload.
            let mut record = vec![0x08];
            record.extend_from_slice(b"abcdefgh");

            for chunk in record.chunks(1).take(record.len() - 1) {
                server.write_all(chunk).await.unwrap();
                server.flush().await.unwrap();
                tokio::select! {
                    _ = transport.recv() => panic!("recv completed with a partial record"),
                    _ = tokio::task::yield_now() => {}
                }
            }

            server.write_all(&record[record.len() - 1..]).await.unwrap();
            server.flush().await.unwrap();

            let got = transport.recv().await.expect("record should decode");
            assert_eq!(&got[..], b"abcdefgh");
        }

        // Two frames arrive in a single write. Each recv() must return one
        // complete frame, in order. Exercises the channel queue + the reader
        // task looping on a buffer that still has bytes after parsing.
        #[tokio::test]
        async fn recv_returns_consecutive_frames_in_order() {
            let (client, mut server) = tokio::io::duplex(64 * 1024);
            let (_writer, mut transport) = Stream::new(client, Version::QMux00, 16 * 1024).split();

            // Two STREAM frames (type 0x0a) for stream ids 4 and 8.
            let frame_a: Vec<u8> = [0x0a, 0x04, 0x05].into_iter().chain(*b"hello").collect();
            let frame_b: Vec<u8> = [0x0a, 0x08, 0x05].into_iter().chain(*b"world").collect();
            let mut combined = frame_a.clone();
            combined.extend_from_slice(&frame_b);

            server.write_all(&combined).await.unwrap();
            server.flush().await.unwrap();

            let got_a = transport.recv().await.expect("first frame should decode");
            let got_b = transport.recv().await.expect("second frame should decode");
            assert_eq!(&got_a[..], frame_a.as_slice());
            assert_eq!(&got_b[..], frame_b.as_slice());
        }

        // Reader task hits a parse error: `recv()` returns it, and the next
        // `recv()` returns Error::Closed since the task has exited.
        #[tokio::test]
        async fn recv_propagates_parse_error_then_closes() {
            let (client, mut server) = tokio::io::duplex(64 * 1024);
            let (_writer, mut transport) = Stream::new(client, Version::QMux00, 16 * 1024).split();

            // Frame type 0x02 isn't a recognized QMux00 frame type.
            server.write_all(&[0x02]).await.unwrap();
            server.flush().await.unwrap();

            let err = transport.recv().await.expect_err("parse error expected");
            assert!(matches!(err, Error::InvalidFrameType(0x02)), "got {err:?}");

            // Task has exited after sending the error; subsequent recv sees the
            // closed channel and reports Error::Closed.
            let err = transport.recv().await.expect_err("closed expected");
            assert!(matches!(err, Error::Closed), "got {err:?}");
        }

        // A record whose declared size exceeds `our_max_record_size` is
        // rejected with FrameTooLarge before any payload is consumed.
        #[tokio::test]
        async fn recv_record_exceeding_max_returns_frame_too_large() {
            let (client, mut server) = tokio::io::duplex(64 * 1024);
            let (_writer, mut transport) = Stream::new(client, Version::QMux01, 4).split();

            // 1-byte varint length = 5, which exceeds the configured max of 4.
            server.write_all(&[0x05]).await.unwrap();
            server.flush().await.unwrap();

            let err = transport.recv().await.expect_err("FrameTooLarge expected");
            assert!(matches!(err, Error::FrameTooLarge), "got {err:?}");
        }

        /// draft-00 §5.2 bounds the whole frame at `max_frame_size`, header
        /// included: a frame of exactly that size parses, and one byte more is a
        /// frame-size violation rather than something to read into memory.
        #[tokio::test]
        async fn recv_qmux00_bounds_the_whole_frame() {
            use crate::proto::{Frame, Stream as StreamFrame};
            use crate::{StreamDir, StreamId};

            let frame = |payload: usize| {
                Frame::Stream(StreamFrame {
                    id: StreamId::new(0, StreamDir::Uni, true),
                    offset: 0,
                    data: vec![0x5a; payload].into(),
                    fin: false,
                })
                .encode(Version::QMux00)
                .unwrap()
            };

            let at_limit = frame(MAX_FRAME_SIZE - 5);
            assert_eq!(at_limit.len(), MAX_FRAME_SIZE);
            let over_limit = frame(MAX_FRAME_SIZE - 4);
            assert_eq!(over_limit.len(), MAX_FRAME_SIZE + 1);

            let (client, mut server) = tokio::io::duplex(64 * 1024);
            let (_writer, mut transport) = Stream::new(client, Version::QMux00, 0).split();

            server.write_all(&at_limit).await.unwrap();
            server.flush().await.unwrap();
            assert_eq!(transport.recv().await.unwrap(), at_limit);

            server.write_all(&over_limit).await.unwrap();
            server.flush().await.unwrap();
            let err = transport.recv().await.expect_err("FrameTooLarge expected");
            assert!(matches!(err, Error::FrameTooLarge), "got {err:?}");
        }
    }
}

#[cfg(any(feature = "tcp", all(unix, feature = "uds")))]
pub use stream_transport::{Stream, StreamReader, StreamWriter};

// Shared plumbing for the byte-stream transports (TCP, Unix sockets).
#[cfg(any(feature = "tcp", all(unix, feature = "uds")))]
mod stream_session {
    use tokio::io::{AsyncRead, AsyncWrite};

    use super::Stream;
    use crate::protocol::validate_protocol;
    use crate::{Config, Error, Protocol, Session};

    /// Wrap a byte stream in a [`Stream`] and start a session, validating any
    /// advertised protocol names first. Used by the `tcp`/`uds` builders.
    pub(crate) async fn build<T: AsyncRead + AsyncWrite + Send + 'static>(
        stream: T,
        config: Config,
        is_server: bool,
    ) -> Result<Session, Error> {
        if let Protocol::Negotiate(protocols) = &config.protocol {
            for protocol in protocols {
                validate_protocol(protocol)?;
            }
        }
        let transport = Stream::new(stream, config.version, config.max_record_size);
        if is_server {
            Session::accept(transport, config).await
        } else {
            Session::connect(transport, config).await
        }
    }
}

#[cfg(any(feature = "tcp", all(unix, feature = "uds")))]
pub(crate) use stream_session::build as build_stream_session;

// WsTransport: message I/O over WebSocket.
#[cfg(feature = "ws")]
mod ws_transport {
    use std::pin::Pin;
    use std::time::Duration;

    use bytes::Bytes;
    use futures::stream::SplitSink;
    use tokio::sync::mpsc;
    use tokio::task::JoinHandle;
    use tokio::time::{Instant, Interval, MissedTickBehavior, Sleep};
    use tokio_tungstenite::tungstenite;

    use super::{Reader, Transport, Writer};
    use crate::ws::KeepAlive;
    use crate::{Error, Version, MAX_FRAME_SIZE};

    type Message = tungstenite::Message;

    /// Bound on WS frames the pump task reads ahead of the session. Small: the pump
    /// only stages already-flow-control-covered frames; a full channel makes the
    /// pump park on delivery, which is exactly the backpressure signal the deadline
    /// logic keys off. Mirrors the byte-stream reader's `RECV_CHANNEL_CAPACITY`.
    const WS_RECV_CHANNEL_CAPACITY: usize = 16;

    /// The combined `Stream + Sink` bound every WebSocket half requires.
    pub(crate) trait WsStream:
        futures::Stream<Item = Result<Message, tungstenite::Error>>
        + futures::Sink<Message, Error = tungstenite::Error>
        + Unpin
        + Send
        + 'static
    {
    }
    impl<T> WsStream for T where
        T: futures::Stream<Item = Result<Message, tungstenite::Error>>
            + futures::Sink<Message, Error = tungstenite::Error>
            + Unpin
            + Send
            + 'static
    {
    }

    pub(crate) struct WsTransport<T> {
        ws: T,
        keep_alive: Option<KeepAlive>,
        /// Largest inbound message we accept, or `None` when nothing bounds it.
        recv_limit: Option<usize>,
    }

    impl<T> WsTransport<T> {
        pub fn new(ws: T, version: Version, max_record_size: u64) -> Self {
            // A message is a record on the record-framed drafts, bounded by the
            // `max_record_size` we advertised; on draft-00 it is a single frame,
            // bounded by `max_frame_size` (§5.2), which we never advertise above
            // its initial value. The legacy binding negotiates neither.
            let recv_limit = if version.uses_records() {
                Some(usize::try_from(max_record_size).unwrap_or(usize::MAX))
            } else if version.is_qmux() {
                Some(MAX_FRAME_SIZE)
            } else {
                None
            };
            Self {
                ws,
                keep_alive: None,
                recv_limit,
            }
        }

        pub fn with_keep_alive(mut self, keep_alive: KeepAlive) -> Self {
            self.keep_alive = Some(keep_alive);
            self
        }
    }

    /// Writer-side keep-alive: emit a Ping every `interval`.
    struct PingState {
        // Fires on each interval; the writer sends a Ping when it does.
        interval: Interval,
    }

    impl PingState {
        fn new(config: KeepAlive) -> Self {
            // tokio::time::interval panics on a zero Duration; floor to 1ms so a
            // misconfigured KeepAlive degrades into "very chatty" instead of crashing.
            let interval_dur = config.interval.max(Duration::from_millis(1));
            // Skip catch-up bursts after a long pause; we just want one Ping per tick.
            let mut interval = tokio::time::interval(interval_dur);
            interval.set_missed_tick_behavior(MissedTickBehavior::Delay);
            // First tick fires immediately by default; consume it so we don't ping on connect.
            interval.reset();
            Self { interval }
        }
    }

    /// Reader-side keep-alive: close the session if no frame arrives within `timeout`.
    struct DeadlineState {
        // Resets every time we receive a frame. If it elapses, the peer is gone.
        deadline: Pin<Box<Sleep>>,
        timeout: Duration,
    }

    impl DeadlineState {
        fn new(config: KeepAlive) -> Self {
            let interval_dur = config.interval.max(Duration::from_millis(1));
            // A deadline shorter than the interval would fire before the first ping.
            let timeout = config.timeout.max(interval_dur);
            Self {
                deadline: Box::pin(tokio::time::sleep(timeout)),
                timeout,
            }
        }

        fn observe_recv(&mut self) {
            self.deadline.as_mut().reset(Instant::now() + self.timeout);
        }
    }

    /// The send half of a [`WsTransport`]: owns the sink and drives keep-alive Pings.
    pub(crate) struct WsWriter<T: WsStream> {
        sink: SplitSink<T, Message>,
        ping: Option<PingState>,
    }

    /// The receive half of a [`WsTransport`].
    ///
    /// A dedicated pump task owns the socket and the idle deadline; `recv` is just a
    /// hand-off from its channel. Reading on its own task is what makes the keep-alive
    /// deadline robust: it advances (and resets on the peer's Pong replies) regardless
    /// of where the *session* is parked — in particular while the session is wedged
    /// handing a stream to a slow `accept_*`, which would otherwise starve a deadline
    /// polled inline in `recv` and fire a spurious timeout. Mirrors [`StreamReader`].
    pub(crate) struct WsReader {
        rx: mpsc::Receiver<Result<Bytes, Error>>,
        /// Aborted on drop so the pump can't outlive the transport.
        pump: JoinHandle<()>,
    }

    impl Drop for WsReader {
        fn drop(&mut self) {
            self.pump.abort();
        }
    }

    impl<T: WsStream> Transport for WsTransport<T> {
        type Writer = WsWriter<T>;
        type Reader = WsReader;

        fn split(self) -> (WsWriter<T>, WsReader) {
            use futures::StreamExt;
            // BiLock-backed halves: the sink and stream can be polled concurrently
            // on separate tasks, briefly serializing on the shared socket.
            let (sink, stream) = self.ws.split();
            let (ping, deadline) = match self.keep_alive {
                Some(ka) => (Some(PingState::new(ka)), Some(DeadlineState::new(ka))),
                None => (None, None),
            };
            let (tx, rx) = mpsc::channel(WS_RECV_CHANNEL_CAPACITY);
            let pump = tokio::spawn(ws_pump(stream, deadline, self.recv_limit, tx));
            (WsWriter { sink, ping }, WsReader { rx, pump })
        }
    }

    impl<T: WsStream> Writer for WsWriter<T> {
        async fn send(&mut self, data: Bytes) -> Result<(), Error> {
            use futures::SinkExt;
            self.sink
                .send(Message::Binary(data))
                .await
                .map_err(|_| Error::Closed)?;
            Ok(())
        }

        async fn close(&mut self) -> Result<(), Error> {
            use futures::SinkExt;
            self.sink.close().await.map_err(|_| Error::Closed)?;
            Ok(())
        }

        async fn maintain(&mut self) -> Result<(), Error> {
            use futures::SinkExt;
            match &mut self.ping {
                Some(ping) => {
                    // Wait for the next interval, then send one keep-alive Ping. The
                    // session's writer loop re-invokes this each time it resolves, so
                    // pings keep flowing without a dedicated task. tungstenite's
                    // auto-queued Pong replies (from the reader) also flush here.
                    ping.interval.tick().await;
                    self.sink
                        .send(Message::Ping(Bytes::new()))
                        .await
                        .map_err(|_| Error::Closed)?;
                    Ok(())
                }
                // No keep-alive configured: never resolves, so the writer loop's
                // select simply ignores this branch.
                None => std::future::pending().await,
            }
        }
    }

    impl Reader for WsReader {
        async fn recv(&mut self) -> Result<Bytes, Error> {
            // The pump task pushes complete Binary frames (or a terminal error)
            // here; `None` means it exited without sending — treat as a clean close.
            self.rx.recv().await.unwrap_or(Err(Error::Closed))
        }
    }

    /// Pump task: read WS messages off the socket independently of the session,
    /// resetting the keep-alive deadline on every message (including the peer's Pong
    /// replies) and shipping Binary frames to the session over `tx`. On a timeout,
    /// clean close, or socket error it sends a terminal `Err` and exits; if `tx` is
    /// closed (the transport was dropped) it exits silently.
    ///
    /// Why a task rather than an inline deadline in `recv`: the deadline must keep
    /// advancing even while the session is parked (e.g. handing a stream to a slow
    /// `accept_*`), and the peer's liveness must not be judged while we're merely
    /// backpressured. So the deadline is reset both when a message arrives *and*
    /// after each delivery to the session — a long delivery park (session
    /// backpressure) therefore isn't charged against the peer.
    async fn ws_pump<S>(
        mut stream: S,
        mut deadline: Option<DeadlineState>,
        recv_limit: Option<usize>,
        tx: mpsc::Sender<Result<Bytes, Error>>,
    ) where
        S: futures::Stream<Item = Result<Message, tungstenite::Error>> + Unpin + Send + 'static,
    {
        use futures::StreamExt;

        loop {
            // Prefer a real message over the deadline (`biased`): if both are ready,
            // the peer is alive, so don't spuriously time out.
            let message = match &mut deadline {
                Some(d) => tokio::select! {
                    biased;
                    msg = stream.next() => msg,
                    _ = d.deadline.as_mut() => {
                        tracing::debug!("websocket keep_alive timeout");
                        let _ = tx.send(Err(Error::Closed)).await;
                        return;
                    }
                },
                None => stream.next().await,
            };

            // Any read — data or control — proves the peer is alive; reset the deadline.
            if let Some(d) = deadline.as_mut() {
                d.observe_recv();
            }

            let message = match message {
                Some(Ok(message)) => message,
                // The socket errored mid-stream: forward the specific WebSocket
                // error (matching `StreamReader`, and preserving what the old inline
                // `recv()` surfaced via `?`) rather than flattening it to a generic
                // close, so callers can still distinguish e.g. a protocol fault.
                Some(Err(err)) => {
                    let _ = tx.send(Err(err.into())).await;
                    return;
                }
                // Clean end of stream.
                None => {
                    let _ = tx.send(Err(Error::Closed)).await;
                    return;
                }
            };

            match message {
                Message::Binary(data) => {
                    if recv_limit.is_some_and(|limit| data.len() > limit) {
                        // Release the oversized allocation before waiting for room to
                        // report the terminal error. In particular, never let it enter
                        // the bounded record queue while the session is backpressured.
                        drop(data);
                        let _ = tx.send(Err(Error::FrameTooLarge)).await;
                        return;
                    }
                    if tx.send(Ok(data)).await.is_err() {
                        return; // session gone
                    }
                    // Delivered. Reset the deadline again so a long delivery park
                    // (the session backpressured on `accept_*`) isn't charged
                    // against the peer — we only judge liveness while actively
                    // reading, not while wedged handing data to the app.
                    if let Some(d) = deadline.as_mut() {
                        d.observe_recv();
                    }
                }
                Message::Close(_) => {
                    let _ = tx.send(Err(Error::Closed)).await;
                    return;
                }
                Message::Ping(_) | Message::Pong(_) | Message::Text(_) | Message::Frame(_) => {
                    // tungstenite auto-queues a Pong reply when it reads a Ping; the
                    // writer half flushes it on its next send/ping. The pump owns no
                    // sink, so there's nothing to reply with here.
                    continue;
                }
            }
        }
    }

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

        use bytes::Bytes;
        use tokio::sync::mpsc;

        use super::{ws_pump, DeadlineState, Message};
        use crate::ws::KeepAlive;
        use crate::Error;

        type WsResult = Result<Message, tokio_tungstenite::tungstenite::Error>;

        /// A `Stream` of WS messages backed by a `futures` channel. Holding the
        /// returned sender keeps the stream "open but silent" (`next()` pends) when
        /// empty; dropping it ends the stream. Standing in for a live peer.
        fn mock_stream() -> (
            futures::channel::mpsc::UnboundedSender<WsResult>,
            futures::channel::mpsc::UnboundedReceiver<WsResult>,
        ) {
            futures::channel::mpsc::unbounded()
        }

        /// The pump must not spuriously time out a *healthy* peer just because the
        /// session is too busy to drain it — the exact failure the split reader had,
        /// where the keep-alive deadline was only polled inside `recv` and went stale
        /// while the session parked handing a stream to a full `accept_*` channel.
        ///
        /// Here a capacity-1 channel stands in for that full `accept_*`: the pump
        /// parks delivering the second frame while we sit idle well past the
        /// keep-alive `timeout`, then we drain — and both frames must arrive intact,
        /// with no `Error::Closed` from a phantom timeout.
        #[tokio::test]
        async fn survives_a_slow_consumer() {
            let (feed, stream) = mock_stream();
            feed.unbounded_send(Ok(Message::Binary(Bytes::from_static(b"one"))))
                .unwrap();
            feed.unbounded_send(Ok(Message::Binary(Bytes::from_static(b"two"))))
                .unwrap();
            // Keep `feed` alive: the peer stays connected and silent after these two.
            let _feed = feed;

            let ka = KeepAlive::new(Duration::from_millis(10), Duration::from_millis(50));
            let (tx, mut rx) = mpsc::channel(1);
            let pump = tokio::spawn(ws_pump(stream, Some(DeadlineState::new(ka)), None, tx));

            // Model a session wedged on a full accept channel: don't read for well
            // over the 50ms keep-alive timeout while the pump is parked on delivery.
            tokio::time::sleep(Duration::from_millis(150)).await;

            assert_eq!(
                rx.recv().await.expect("channel open").expect("no timeout"),
                Bytes::from_static(b"one"),
            );
            assert_eq!(
                rx.recv().await.expect("channel open").expect("no timeout"),
                Bytes::from_static(b"two"),
            );

            pump.abort();
        }

        /// The deadline must still fire for a genuinely silent peer: once the pump is
        /// caught up and reading (not parked on delivery), `timeout` of silence is a
        /// dead peer, surfaced as `Error::Closed`.
        #[tokio::test]
        async fn times_out_a_silent_peer() {
            // Held open but never fed — `stream.next()` pends forever.
            let (_feed, stream) = mock_stream();

            let ka = KeepAlive::new(Duration::from_millis(10), Duration::from_millis(50));
            let (tx, mut rx) = mpsc::channel(4);
            let pump = tokio::spawn(ws_pump(stream, Some(DeadlineState::new(ka)), None, tx));

            let result = tokio::time::timeout(Duration::from_secs(1), rx.recv())
                .await
                .expect("keep-alive deadline should fire on a silent peer");
            assert!(
                matches!(result, Some(Err(Error::Closed))),
                "expected a keep-alive timeout, got {result:?}"
            );

            pump.abort();
        }

        /// With no keep-alive configured the pump never times out: a silent-but-open
        /// peer just leaves `recv` pending rather than closing.
        #[tokio::test]
        async fn no_keep_alive_never_times_out() {
            let (feed, stream) = mock_stream();
            let _feed = feed;

            let (tx, mut rx) = mpsc::channel(4);
            let pump = tokio::spawn(ws_pump(stream, None, None, tx));

            // No deadline, so recv stays pending well past any keep-alive window.
            let pending = tokio::time::timeout(Duration::from_millis(100), rx.recv()).await;
            assert!(pending.is_err(), "recv must not resolve without activity");

            pump.abort();
        }

        #[tokio::test]
        async fn rejects_record_exceeding_max_before_queueing() {
            let (feed, stream) = mock_stream();
            feed.unbounded_send(Ok(Message::Binary(Bytes::from_static(b"oversized"))))
                .unwrap();

            let (tx, mut rx) = mpsc::channel(1);
            let pump = tokio::spawn(ws_pump(stream, None, Some(4), tx));

            assert!(matches!(rx.recv().await, Some(Err(Error::FrameTooLarge))));
            assert!(
                rx.recv().await.is_none(),
                "pump must stop after the violation"
            );
            pump.await.unwrap();
        }
    }
}

#[cfg(feature = "ws")]
pub(crate) use ws_transport::WsTransport;