tachyon-i2p 0.0.2

//! [`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::task::{Context, Poll};
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
use tokio::task::JoinHandle;

/// Per-call poll interval for the blocking receive worker (see [`ReadState`]). Short enough that
/// closing the stream (which wakes any in-progress `Receive` promptly via libi2pd's own
/// mechanism -- see `shim.h`) is never meaningfully delayed by this, long enough to avoid a busy
/// loop while a connection is simply idle.
const RECEIVE_POLL_SECS: i32 = 10;

struct StreamHandle(*mut i2pd_sys::I2pdStream);

// 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 (posts internally); in particular, this
    /// promptly wakes up any `Receive` call currently blocked on another thread, which is what
    /// makes dropping an `I2pStream` responsive even while a read is in flight (see `Drop for
    /// I2pStream` below) rather than waiting out `RECEIVE_POLL_SECS`.
    fn close(&self) {
        // SAFETY: `self.0` is a valid handle for as long as this `StreamHandle` exists.
        unsafe { i2pd_sys::i2pd_stream_close(self.0) }
    }
}

impl Drop for StreamHandle {
    fn drop(&mut self) {
        // SAFETY: `self.0` is uniquely owned by this handle; nothing else can be holding a
        // reference by the time the last `Arc<StreamHandle>` clone is dropped.
        unsafe { i2pd_sys::i2pd_destroy_stream(self.0) }
    }
}

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

/// Outcome of one blocking receive worker invocation, handed back through the `JoinHandle` along
/// with the buffer and handle it borrowed 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..]` (`buf`'s logical length, not its capacity) holds bytes
    /// already pulled from libi2pd but not yet delivered to a caller -- ordinarily empty
    /// (`pos == buf.len()`), but can be non-empty if a receive completed with more bytes than fit
    /// in the caller's buffer at the time (e.g. a read resumed, after being cancelled mid-flight,
    /// with a smaller buffer than the one the abandoned receive was sized for). Delivered before
    /// starting any new receive, so no data is ever silently dropped. `buf`'s spare capacity is
    /// reused as scratch space for the next receive once fully drained.
    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 directly into anything that accepts a
/// generic Tokio-compatible transport (e.g. hyper's connection builders).
///
/// # A note on backpressure
/// libi2pd's underlying `Send` call queues data into its own internal send buffer and returns
/// immediately -- there is no backpressure signal from a slow remote peer surfaced back through
/// this API (the same limitation libi2pd's own SAM/BOB bridges have). Writing far faster than
/// the remote reads will grow that internal buffer; this is a known, upstream characteristic of
/// libi2pd's simple blocking `Send` API, not something this wrapper can currently smooth over.
#[derive(Debug)]
pub struct I2pStream {
    handle: Arc<StreamHandle>,
    read: ReadState,
    // Keeps the whole router alive for at least as long as this stream. Declared *last*: Rust
    // drops struct fields in declaration order, and `handle`'s own `Drop` (tearing down this
    // stream via libi2pd) must run before the router's `Drop` (tearing down libi2pd's global
    // state) -- not after.
    _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)),
            read: ReadState::Idle {
                buf: Vec::with_capacity(8192),
                pos: 0,
            },
            _router: router,
        }
    }

    /// Whether the underlying libi2pd stream still considers itself open. A cheap, non-blocking
    /// check -- does not itself detect a half-closed (write-only or read-only) state.
    #[must_use]
    pub fn is_open(&self) -> bool {
        // SAFETY: `self.handle.0` is valid for the lifetime of `self`.
        unsafe { i2pd_sys::i2pd_stream_is_open(self.handle.0) != 0 }
    }
}

/// Copies as much of `buf[*pos..]` into `out` as fits, advancing `*pos`. Returns `true` if there
/// was any leftover data to copy (`*pos < buf.len()` on entry) -- i.e. whether it actually
/// delivered anything, as opposed to there being nothing to drain and the caller needing to start
/// a new receive instead. When it drains everything, resets `buf`/`pos` back to empty so `buf`'s
/// capacity can be reused as scratch space for the next receive.
///
/// Pulled out of `poll_read` as a pure function so the no-overflow property (never calling
/// `out.put_slice` with more than `out.remaining()`, which would panic) can be unit-tested
/// directly, without needing 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();
        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;
                    scratch.clear();
                    let want = out.remaining().max(1);
                    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);
                        (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 now rather than looping back into the `Idle`
                        // arm, where an empty `buf` would be indistinguishable from "no leftover,
                        // go receive more" and spawn another (pointless, EOF-forever) receive.
                        this.read = ReadState::Idle {
                            buf: Vec::new(),
                            pos: 0,
                        };
                        return Poll::Ready(Ok(()));
                    }
                    // Only the first `n` bytes are valid data -- the rest is this scratch
                    // buffer's stale zero-fill from before the receive call.
                    buf.truncate(n);
                    this.read = ReadState::Idle { buf, pos: 0 };
                    // loop back to the `Idle` arm above to deliver (possibly only part of) it
                }
            }
        }
    }
}

/// Runs on a blocking worker thread: repeatedly calls libi2pd's blocking receive with a short
/// per-call timeout until either real data arrives, the stream is genuinely closed, or an error
/// occurs -- see `shim.h`'s note that a `0` return means *either* a timeout *or* EOF, which this
/// distinguishes via `i2pd_stream_is_open` so a merely-idle connection is never mistaken for one
/// that's actually done (an `AsyncRead` returning `Ok(())` without filling `buf` at all means EOF
/// to every caller of this trait).
fn receive_until_data_or_close(handle: &StreamHandle, buf: &mut [u8]) -> io::Result<usize> {
    loop {
        // SAFETY: `handle.0` is valid; `buf` is a valid, writable buffer for its full length.
        let n = unsafe {
            i2pd_sys::i2pd_stream_receive(handle.0, buf.as_mut_ptr(), buf.len(), RECEIVE_POLL_SECS)
        };
        if n < 0 {
            return Err(io::Error::other("i2p stream receive failed"));
        }
        if n > 0 {
            return Ok(usize::try_from(n).unwrap_or(0));
        }
        // SAFETY: `handle.0` is valid.
        let still_open = unsafe { i2pd_sys::i2pd_stream_is_open(handle.0) != 0 };
        if !still_open {
            return Ok(0); // genuine EOF
        }
        // otherwise: just a quiet period, poll again
    }
}

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();
        // SAFETY: `this.handle.0` is valid; `Send` (per `shim.h`) is non-blocking (posts
        // internally and returns immediately), safe to call directly from an async context.
        let n = unsafe { i2pd_sys::i2pd_stream_send(this.handle.0, buf.as_ptr(), buf.len()) };
        if n < 0 {
            Poll::Ready(Err(io::Error::new(
                io::ErrorKind::BrokenPipe,
                "i2p stream send failed",
            )))
        } else {
            Poll::Ready(Ok(usize::try_from(n).unwrap_or(0)))
        }
    }

    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        // Send already hands data to libi2pd's own queue; there is nothing buffered on this side
        // to flush (see the backpressure note on `I2pStream`).
        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) {
        // Closes eagerly (rather than waiting for the last `Arc<StreamHandle>` clone -- which a
        // still-in-flight read worker may be holding -- to drop) so a stream dropped mid-read
        // (e.g. cancellation) doesn't leave its worker thread blocked for up to
        // `RECEIVE_POLL_SECS` for no reason; actual deallocation is deferred to `StreamHandle`'s
        // own `Drop`, once every clone (including any in-flight worker's) is gone.
        self.handle.close();
    }
}

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

    // Regression test for a real bug: a read resumed (after being cancelled mid-flight, e.g. by
    // a per-call `tokio::time::timeout`) with a *smaller* destination buffer than the abandoned
    // read was originally sized for used to panic inside `ReadBuf::put_slice` (which asserts the
    // slice fits in `remaining()`) instead of retaining the overflow for the next call.
    #[test]
    fn leftover_larger_than_destination_does_not_overflow_or_panic() {
        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 the buffer 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);
    }

    #[test]
    fn leftover_exactly_matching_destination_drains_fully_in_one_call() {
        let mut buf = vec![1u8, 2, 3];
        let mut pos = 0usize;
        let mut dst = [0u8; 3];
        let mut out = ReadBuf::new(&mut dst);
        assert!(drain_leftover(&mut buf, &mut pos, &mut out));
        assert_eq!(out.filled(), &[1, 2, 3]);
        assert_eq!(pos, 0);
        assert!(buf.is_empty());
    }
}