tachyon-i2p 0.0.4

//! [`I2pStream`]: a single stream-protocol connection, implementing
//! [`AsyncRead`]/[`AsyncWrite`].

use crate::router::I2pRouter;
use std::future::Future;
use std::io;
use std::pin::Pin;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::task::{Context, Poll};
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
use tokio::task::JoinHandle;

/// Per-call poll interval for the blocking receive worker. Long enough not to busy-loop on an
/// idle connection; closing the stream wakes an in-progress `Receive` through libi2pd's own
/// mechanism (see `shim.h`), so this never delays a shutdown.
const RECEIVE_POLL_SECS: i32 = 10;

/// Upper bound on the scratch buffer handed to one libi2pd receive call.
///
/// A receive is sized to the caller's `ReadBuf`, which hyper and friends routinely make far
/// bigger than any single I2P stream packet; uncapped, every `poll_read` grows and zero-fills a
/// buffer that size for a receive returning a few kilobytes. Overflow is retained as leftover
/// (see [`ReadState::Idle`]), so the cap costs only an extra `poll_read` round on the rare
/// oversized read.
const MAX_RECEIVE_CHUNK: usize = 64 * 1024;

struct StreamHandle {
    ptr: *mut i2pd_sys::I2pdStream,
    /// Set the first time [`close`](StreamHandle::close) runs, so `Stream::Close` is driven at
    /// most once from here.
    ///
    /// Close is called from both `poll_shutdown` and `Drop for I2pStream` (which fires even after
    /// an explicit shutdown), and libi2pd's `Stream::Close` is not an idempotent flag flip: it
    /// mutates the stream's status, send buffer and packet queues on the calling thread, and
    /// re-entering it on an already-closed stream drives it through `Terminate`. `Drop for
    /// StreamHandle` then reaches it once more via `i2pd_destroy_stream`'s own close --
    /// close-then-terminate is libi2pd's teardown order for a stream it is discarding.
    closed: AtomicBool,
}

// SAFETY: every libi2pd entry point this handle's pointer is passed to (`shim.h`) is explicitly
// documented as safe to call concurrently from independent threads -- `Send`/`Receive` operate
// on separate internal buffers and both dispatch through libi2pd's own thread-safe io_service.
unsafe impl Send for StreamHandle {}
unsafe impl Sync for StreamHandle {}

impl StreamHandle {
    /// Signals libi2pd to close the stream. Non-blocking, and wakes any `Receive` blocked on
    /// another thread, which is what keeps dropping an `I2pStream` mid-read responsive instead of
    /// waiting out `RECEIVE_POLL_SECS`.
    fn close(&self) {
        if self.closed.swap(true, Ordering::AcqRel) {
            return; // see `closed`
        }
        // SAFETY: `self.ptr` is a valid handle for as long as this `StreamHandle` exists.
        unsafe { i2pd_sys::i2pd_stream_close(self.ptr) }
    }

    /// Whether libi2pd still reports the stream as open. A status-field read; non-blocking.
    fn is_open(&self) -> bool {
        // SAFETY: `self.ptr` is a valid handle for as long as this `StreamHandle` exists.
        unsafe { i2pd_sys::i2pd_stream_is_open(self.ptr) != 0 }
    }
}

impl Drop for StreamHandle {
    fn drop(&mut self) {
        // SAFETY: `self.ptr` is uniquely owned by this handle; nothing else holds a reference by
        // the time the last `Arc<StreamHandle>` clone drops. `destroy` closes the stream itself,
        // so a path that never called `close` still tears it down inside libi2pd rather than
        // leaking a half-open one onto its destination.
        unsafe { i2pd_sys::i2pd_destroy_stream(self.ptr) }
    }
}

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

/// Carries the borrowed buffer and handle back out of the worker so the next poll can reuse them
/// without reallocating.
type ReadOutcome = (Arc<StreamHandle>, Vec<u8>, io::Result<usize>);

#[derive(Debug)]
enum ReadState {
    /// No read in flight. `buf[pos..]` holds bytes pulled from libi2pd but not yet delivered --
    /// ordinarily empty, but non-empty when a receive completed with more bytes than fit in the
    /// caller's buffer (e.g. a read cancelled mid-flight and resumed with a smaller buffer than
    /// the abandoned receive was sized for). Delivered before any new receive starts, so no data
    /// is silently dropped, after which the spare capacity is reused as receive scratch.
    Idle {
        buf: Vec<u8>,
        pos: usize,
    },
    Reading(JoinHandle<ReadOutcome>),
}

/// One stream-protocol connection over I2P -- either accepted via
/// [`Destination::accept`](crate::Destination::accept) or opened via
/// [`Destination::connect`](crate::Destination::connect).
///
/// Implements [`AsyncRead`]/[`AsyncWrite`], so it plugs into anything taking a generic
/// Tokio-compatible transport (e.g. hyper's connection builders).
///
/// # Backpressure
/// libi2pd's `Send` queues into its own buffer and returns immediately, so a slow remote peer
/// produces no backpressure signal through this API -- the same limitation libi2pd's SAM/BOB
/// bridges have. Writing faster than the remote reads grows that internal buffer.
#[derive(Debug)]
pub struct I2pStream {
    handle: Arc<StreamHandle>,
    read: ReadState,
    /// Latches once libi2pd has reported this stream as open at least once -- see `poll_write`.
    saw_open: bool,
    // Keeps the router alive at least as long as this stream. Declared *last* because fields drop
    // in declaration order: `handle`'s `Drop` tears this stream down via libi2pd, which must
    // happen before the router's `Drop` tears down libi2pd's globals.
    _router: I2pRouter,
}

impl I2pStream {
    /// # Safety
    /// `ptr` must be a valid, uniquely-owned `I2pdStream` handle (i.e. not already passed to
    /// this function, [`i2pd_sys::i2pd_destroy_stream`], or any other consumer).
    pub(crate) unsafe fn from_raw(router: I2pRouter, ptr: *mut i2pd_sys::I2pdStream) -> Self {
        Self {
            handle: Arc::new(StreamHandle {
                ptr,
                closed: AtomicBool::new(false),
            }),
            read: ReadState::Idle {
                buf: Vec::with_capacity(8192),
                pos: 0,
            },
            saw_open: false,
            _router: router,
        }
    }

    /// Whether the underlying libi2pd stream still considers itself open. Non-blocking, and does
    /// not detect a half-closed (write-only or read-only) state.
    ///
    /// `false` covers both a finished (closed/reset) stream *and* a fresh outbound one whose SYN
    /// hasn't been acknowledged yet, which libi2pd tracks as a distinct state. This answers "open
    /// right now", not "usable".
    #[must_use]
    pub fn is_open(&self) -> bool {
        self.handle.is_open()
    }
}

/// Copies as much of `buf[*pos..]` into `out` as fits, advancing `*pos`. Returns whether there
/// was anything to copy, i.e. whether the caller still needs to start a receive. Draining
/// everything resets `buf`/`pos` so the capacity can be reused as receive scratch.
///
/// Pulled out of `poll_read` as a pure function so the no-overflow property (never passing
/// `out.put_slice` more than `out.remaining()`, which panics) is unit-testable without a live
/// libi2pd stream.
fn drain_leftover(buf: &mut Vec<u8>, pos: &mut usize, out: &mut ReadBuf<'_>) -> bool {
    if *pos >= buf.len() {
        return false;
    }
    let n = (buf.len() - *pos).min(out.remaining());
    out.put_slice(&buf[*pos..*pos + n]);
    *pos += n;
    if *pos == buf.len() {
        buf.clear();
        *pos = 0;
    }
    true
}

impl AsyncRead for I2pStream {
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        out: &mut ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        let this = self.get_mut();
        // A caller with no room left gets an immediate no-op. Falling through would start a
        // receive sized by `out.remaining()`, i.e. a pointless 1-byte one whose result could
        // only be parked as leftover.
        if out.remaining() == 0 {
            return Poll::Ready(Ok(()));
        }
        loop {
            match &mut this.read {
                ReadState::Idle { buf, pos } => {
                    if drain_leftover(buf, pos, out) {
                        return Poll::Ready(Ok(()));
                    }
                    let mut scratch = std::mem::take(buf);
                    *pos = 0;
                    let want = out.remaining().min(MAX_RECEIVE_CHUNK);
                    // Reuses the previous read's allocation, growing only when too short. The
                    // zero-fill is why `want` is capped: it costs at most `MAX_RECEIVE_CHUNK` per
                    // read instead of scaling with a caller buffer that may be megabytes wide.
                    if scratch.len() < want {
                        scratch.resize(want, 0);
                    }
                    let handle = this.handle.clone();
                    this.read = ReadState::Reading(tokio::task::spawn_blocking(move || {
                        let result = receive_until_data_or_close(&handle, &mut scratch[..want]);
                        (handle, scratch, result)
                    }));
                }
                ReadState::Reading(join) => {
                    let (handle, mut buf, result) = match Pin::new(join).poll(cx) {
                        Poll::Pending => return Poll::Pending,
                        Poll::Ready(Ok(outcome)) => outcome,
                        Poll::Ready(Err(_)) => {
                            this.read = ReadState::Idle {
                                buf: Vec::new(),
                                pos: 0,
                            };
                            return Poll::Ready(Err(io::Error::other(
                                "i2p receive worker panicked",
                            )));
                        }
                    };
                    drop(handle); // the extra Arc clone the worker held
                    let n = match result {
                        Ok(n) => n,
                        Err(e) => {
                            this.read = ReadState::Idle {
                                buf: Vec::new(),
                                pos: 0,
                            };
                            return Poll::Ready(Err(e));
                        }
                    };
                    if n == 0 {
                        // Genuine EOF. Signal it here rather than looping into the `Idle` arm,
                        // where an empty `buf` is indistinguishable from "no leftover, receive
                        // more" and would spawn another EOF-forever receive.
                        this.read = ReadState::Idle {
                            buf: Vec::new(),
                            pos: 0,
                        };
                        return Poll::Ready(Ok(()));
                    }
                    // Only the first `n` bytes are data; the rest is the scratch buffer's stale
                    // zero-fill.
                    buf.truncate(n);
                    this.read = ReadState::Idle { buf, pos: 0 };
                }
            }
        }
    }
}

/// Runs on a blocking worker: calls libi2pd's blocking receive with a short per-call timeout
/// until data arrives, the stream closes, or it errors. Per `shim.h` a `0` return means *either*
/// timeout *or* EOF; `i2pd_stream_is_open` distinguishes them, since an `AsyncRead` returning
/// `Ok(())` with nothing filled means EOF to every caller of the trait.
fn receive_until_data_or_close(handle: &StreamHandle, buf: &mut [u8]) -> io::Result<usize> {
    loop {
        // SAFETY: `handle.ptr` is valid; `buf` is a valid, writable buffer for its full length.
        let n = unsafe {
            i2pd_sys::i2pd_stream_receive(
                handle.ptr,
                buf.as_mut_ptr(),
                buf.len(),
                RECEIVE_POLL_SECS,
            )
        };
        if n < 0 {
            return Err(io::Error::other("i2p stream receive failed"));
        }
        if n > 0 {
            // Clamped: a shim reporting more bytes than were asked for would have `poll_read`
            // treat uninitialized tail bytes as received data.
            return Ok(usize::try_from(n).unwrap_or(0).min(buf.len()));
        }
        if !handle.is_open() {
            return Ok(0); // EOF, not a timeout
        }
    }
}

impl AsyncWrite for I2pStream {
    fn poll_write(
        self: Pin<&mut Self>,
        _cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<io::Result<usize>> {
        let this = self.get_mut();

        // libi2pd's `Send` queues onto its io_service and reports the full length written
        // whatever the stream's state, so on a remote-closed stream it returns success forever
        // and `write_all` discards every byte into a dead connection.
        //
        // `is_open` alone can't stand in for that check: it is also false for a fresh outbound
        // stream whose SYN hasn't been acknowledged, and rejecting writes there would break the
        // ordinary connect-then-write sequence. Latching on the first observed open state
        // separates the two -- not-open-yet only ever precedes the latch, closed only follows it.
        if this.handle.is_open() {
            this.saw_open = true;
        } else if this.saw_open {
            return Poll::Ready(Err(io::Error::new(
                io::ErrorKind::BrokenPipe,
                "i2p stream is closed",
            )));
        }

        // SAFETY: `this.handle.ptr` is valid; `Send` (per libi2pd's `Streaming.cpp`) posts to
        // the router's io_service and returns immediately without blocking, so it is safe to
        // call directly from an async context.
        let n = unsafe { i2pd_sys::i2pd_stream_send(this.handle.ptr, buf.as_ptr(), buf.len()) };
        if n < 0 {
            return Poll::Ready(Err(io::Error::new(
                io::ErrorKind::BrokenPipe,
                "i2p stream send failed",
            )));
        }
        let n = usize::try_from(n).unwrap_or(0).min(buf.len());
        if n == 0 && !buf.is_empty() {
            // `AsyncWrite` forbids a zero-length write for a non-empty buffer: `write_all` reads
            // it as `WriteZero`. Nothing accepted here means the send failed.
            return Poll::Ready(Err(io::Error::new(
                io::ErrorKind::BrokenPipe,
                "i2p stream accepted no bytes",
            )));
        }
        Poll::Ready(Ok(n))
    }

    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        // Send hands data to libi2pd's own queue; nothing is buffered on this side to flush.
        Poll::Ready(Ok(()))
    }

    fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        self.handle.close();
        Poll::Ready(Ok(()))
    }
}

impl Drop for I2pStream {
    fn drop(&mut self) {
        // Eager, rather than waiting on the last `Arc<StreamHandle>` clone (which an in-flight
        // read worker may hold), so a stream dropped mid-read doesn't leave that worker blocked
        // for up to `RECEIVE_POLL_SECS`. Deallocation still waits for `StreamHandle`'s `Drop`.
        self.handle.close();
    }
}

#[cfg(test)]
mod tests {
    use super::drain_leftover;
    use tokio::io::ReadBuf;

    // Regression: a read cancelled mid-flight (e.g. by a per-call `tokio::time::timeout`) and
    // resumed with a *smaller* destination buffer used to panic inside `ReadBuf::put_slice`,
    // which asserts the slice fits in `remaining()`, instead of retaining the overflow.
    #[test]
    fn leftover_larger_than_destination_is_retained() {
        let mut leftover = vec![7u8; 50];
        let mut pos = 0usize;

        let mut dst = [0u8; 10];
        let mut out = ReadBuf::new(&mut dst);
        assert!(drain_leftover(&mut leftover, &mut pos, &mut out));
        assert_eq!(out.filled().len(), 10);
        assert_eq!(out.filled(), &[7u8; 10]);
        assert_eq!(pos, 10);
        assert_eq!(
            leftover.len(),
            50,
            "buffer must be retained until fully drained"
        );

        // A second, larger call drains the rest and resets for scratch reuse.
        let mut dst2 = [0u8; 100];
        let mut out2 = ReadBuf::new(&mut dst2);
        assert!(drain_leftover(&mut leftover, &mut pos, &mut out2));
        assert_eq!(out2.filled().len(), 40);
        assert_eq!(pos, 0);
        assert!(
            leftover.is_empty(),
            "fully-drained buffer must reset to empty"
        );
    }

    #[test]
    fn no_leftover_reports_nothing_to_drain() {
        let mut buf = Vec::new();
        let mut pos = 0usize;
        let mut dst = [0u8; 10];
        let mut out = ReadBuf::new(&mut dst);
        assert!(!drain_leftover(&mut buf, &mut pos, &mut out));
        assert_eq!(out.filled().len(), 0);
    }
}