narju/

link.rs

use std::io;

use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use tokio::sync::mpsc::{channel, Sender};

use crate::floor::{rc_val, Val};
use crate::sched::{addr, Addr};
use crate::wire;
use crate::world::{ConnId, Event};

const HEADER: usize = 16;

/// A length is the one field a peer can make arbitrarily large without sending
/// anything, so it is checked before it is believed.
const LIMIT: usize = 16 << 20;

/// The task id it is for on the receiving side, and the encoded payload.
#[derive(Debug, PartialEq)]
pub struct Frame {
    pub to: u64,
    pub body: Vec<u8>,
}

/// The frame that is for the node itself rather than for a task in it. A task
/// id indexes a `Vec`, so this is a number that cannot be one. It carries what
/// must outlive the task it names - a watch and the death it asked about
/// ([`crate::world::Event::Node`]) - and so cannot be addressed to that task.
pub const NODE: u64 = u64::MAX;

pub async fn write_frame<W: AsyncWrite + Unpin>(w: &mut W, f: &Frame) -> io::Result<()> {
    let mut head = [0u8; HEADER];
    head[..8].copy_from_slice(&f.to.to_be_bytes());
    head[8..].copy_from_slice(&(f.body.len() as u64).to_be_bytes());
    w.write_all(&head).await?;
    w.write_all(&f.body).await?;
    w.flush().await
}

/// A stream plus whatever of it has arrived but not yet been consumed.
pub struct Wire<R> {
    r: R,
    buf: Vec<u8>,
    filled: usize,
}

impl<R: AsyncRead + Unpin> Wire<R> {
    pub fn new(r: R) -> Wire<R> {
        Wire {
            r,
            buf: Vec::new(),
            filled: 0,
        }
    }

    /// `None` when the peer closed the stream between frames. A stream that ends
    /// inside one is an error: the peer vanished rather than finished.
    pub async fn next(&mut self) -> io::Result<Option<Frame>> {
        loop {
            if let Some(f) = self.take()? {
                return Ok(Some(f));
            }
            if self.fill().await? == 0 {
                return if self.filled == 0 {
                    Ok(None)
                } else {
                    Err(io::Error::new(
                        io::ErrorKind::UnexpectedEof,
                        "part of a frame",
                    ))
                };
            }
        }
    }

    /// No `await` past the point the bytes are removed, so a caller dropped
    /// mid-race has consumed nothing.
    fn take(&mut self) -> io::Result<Option<Frame>> {
        if self.filled < HEADER {
            return Ok(None);
        }
        let to = u64::from_be_bytes(self.buf[..8].try_into().expect("eight bytes"));
        let n = u64::from_be_bytes(self.buf[8..HEADER].try_into().expect("eight bytes"));
        if n > LIMIT as u64 {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "frame too large",
            ));
        }
        let n = n as usize;
        if self.filled < HEADER + n {
            return Ok(None);
        }
        let body = self.buf[HEADER..HEADER + n].to_vec();
        self.buf.copy_within(HEADER + n..self.filled, 0);
        self.filled -= HEADER + n;
        Ok(Some(Frame { to, body }))
    }

    /// Cancel-safe, because `read` is: the bytes are in the buffer or they were
    /// never taken from the stream.
    async fn fill(&mut self) -> io::Result<usize> {
        if self.buf.len() - self.filled < HEADER {
            self.buf.resize(self.filled + 8192, 0);
        }
        let n = self.r.read(&mut self.buf[self.filled..]).await?;
        self.filled += n;
        Ok(n)
    }
}

/// What a payload that arrived over `conn` means on this side. An address is the
/// one value that does not survive a link unchanged: a task the sender called
/// its own is one on the far side here, and a task it reached *through* this
/// link is one of ours, so the two cases swap and the sender's conn - a name in
/// its numbering - is discarded. Exact rather than heuristic only because an
/// address is a host type ([`Addr`]) a program cannot forge.
///
/// A third-party address has no image here; [`departs`] refuses to send one.
pub fn across(v: &Val, conn: ConnId) -> Val {
    enum Step<'a> {
        Down(&'a Val),
        Cons,
    }
    let mut work = vec![Step::Down(v)];
    let mut out: Vec<Val> = Vec::new();
    while let Some(s) = work.pop() {
        match s {
            Step::Down(Val::Pair(a, b)) => {
                work.push(Step::Cons);
                work.push(Step::Down(b));
                work.push(Step::Down(a));
            }
            Step::Down(v) => out.push(match Addr::of(v) {
                Some(Addr::Task(id)) => Addr::Remote(conn, *id).val(),
                Some(Addr::Remote(_, id)) => Addr::Task(*id).val(),
                None => v.clone(),
            }),
            Step::Cons => {
                let b = out.pop().expect("a cdr");
                let a = out.pop().expect("a car");
                out.push(Val::Pair(rc_val(a), rc_val(b)));
            }
        }
    }
    out.pop().expect("one value")
}

/// Whether every address in a payload leaving over `conn` has an image on the
/// far side - that is, none names a third node. Forwarding one would mean
/// proxying for a link the receiver does not hold, which has its own lifetime
/// and failure questions, so it is refused rather than half-built.
pub fn departs(v: &Val, conn: ConnId) -> bool {
    let mut work = vec![v];
    while let Some(v) = work.pop() {
        match v {
            Val::Pair(a, b) => {
                work.push(a);
                work.push(b);
            }
            _ => match Addr::of(v) {
                Some(Addr::Remote(c, _)) if *c != conn => return false,
                _ => {}
            },
        }
    }
    true
}

const DEPTH: usize = 64;

/// Run a connection as a peer of the node reached through `port`, answering with
/// the way in to it. The outbound half comes back rather than registering
/// itself, so whoever established the connection holds it before anything can
/// name it: registering by event is a race the caller cannot wait out, since a
/// task that already knows the far address could send first and be told there
/// was no such link.
///
/// **Every wait here is on the far side, never on the node.** Room in the node's
/// queue is taken *before* there is anything to put in it - waiting for room
/// while holding an unread outbound queue is a deadlock, not a slow path.
pub fn bridge<S>(conn: ConnId, stream: S, port: Sender<Event>) -> Sender<Frame>
where
    S: AsyncRead + AsyncWrite + Send + Unpin + 'static,
{
    let (tx, mut rx) = channel::<Frame>(DEPTH);
    tokio::spawn(async move {
        let (r, mut w) = tokio::io::split(stream);
        let mut inbound = Wire::new(r);
        let mut room = None;
        // The node being gone is not a reason to stop at once: the outbound
        // queue closing is the edge that says there are no more frames, and
        // breaking here would lose whatever a node queued on its way out.
        let mut deaf = false;
        loop {
            tokio::select! {
                got = port.reserve(), if !deaf && room.is_none() => match got {
                    Ok(p) => room = Some(p),
                    Err(_) => deaf = true,
                },
                got = inbound.next(), if !deaf && room.is_some() => {
                    let Ok(Some(f)) = got else { break };
                    // A frame this end cannot read is the two sides disagreeing
                    // about the protocol, not one message going wrong.
                    let Ok(v) = wire::decode(&f.body) else { break };
                    let v = across(&v, conn);
                    let ev = match f.to {
                        NODE => Event::Node(conn, v),
                        to => Event::Post(addr(to as usize), v),
                    };
                    room.take().expect("guarded").send(ev);
                }
                got = rx.recv() => match got {
                    Some(f) => if write_frame(&mut w, &f).await.is_err() { break },
                    None => break,
                },
            }
        }
        match room {
            Some(p) => p.send(Event::Down(conn)),
            None => drop(port.send(Event::Down(conn)).await),
        }
    });
    tx
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::floor::list;

    #[test]
    fn an_address_swaps_sides_and_takes_this_links_number() {
        assert_eq!(across(&addr(3), 7), Addr::Remote(7, 3).val());
        assert_eq!(across(&Addr::Remote(9, 3).val(), 7), addr(3));
        assert_eq!(across(&across(&addr(3), 7), 7), addr(3));
    }

    #[test]
    fn nothing_but_an_address_is_rewritten() {
        let decoy = list(&[Val::Sym("task".into()), Val::Num(3)]);
        assert_eq!(across(&decoy, 7), decoy);
        let msg = list(&[Val::Sym("req".into()), addr(1), Val::Num(9), decoy.clone()]);
        assert_eq!(
            across(&msg, 7),
            list(&[
                Val::Sym("req".into()),
                Addr::Remote(7, 1).val(),
                Val::Num(9),
                decoy
            ])
        );
    }

    #[test]
    fn only_an_address_for_a_third_node_is_refused() {
        assert!(departs(&addr(1), 0));
        assert!(departs(&list(&[Addr::Remote(0, 1).val(), Val::Num(2)]), 0));
        assert!(!departs(&list(&[Addr::Remote(1, 1).val(), Val::Num(2)]), 0));
    }

    fn frame(to: u64, body: &[u8]) -> Frame {
        Frame {
            to,
            body: body.to_vec(),
        }
    }

    async fn round(frames: Vec<Frame>) -> Vec<Frame> {
        let (mut a, b) = tokio::io::duplex(64);
        let feed = tokio::spawn(async move {
            for f in &frames {
                write_frame(&mut a, f).await.expect("writes");
            }
        });
        let mut w = Wire::new(b);
        let mut out = Vec::new();
        while let Some(f) = w.next().await.expect("reads") {
            out.push(f);
        }
        feed.await.expect("no panic");
        out
    }

    #[tokio::test]
    async fn frames_arrive_whole_and_in_order() {
        let big = vec![7u8; 5000];
        let sent = vec![frame(0, b""), frame(1, b"x"), frame(u64::MAX, &big)];
        let want = vec![frame(0, b""), frame(1, b"x"), frame(u64::MAX, &big)];
        assert_eq!(round(sent).await, want);
    }

    #[tokio::test]
    async fn a_truncated_frame_is_an_error_and_a_clean_end_is_not() {
        let (a, b) = tokio::io::duplex(64);
        drop(a);
        assert_eq!(Wire::new(b).next().await.expect("reads"), None);

        let (mut a, b) = tokio::io::duplex(64);
        let mut w = Wire::new(b);
        let feed = tokio::spawn(async move {
            a.write_all(&[0, 0, 0, 0, 0, 0, 0, 1])
                .await
                .expect("writes");
        });
        assert_eq!(
            w.next().await.expect_err("truncated").kind(),
            io::ErrorKind::UnexpectedEof
        );
        feed.await.expect("no panic");
    }

    #[tokio::test]
    async fn an_absurd_length_is_refused_rather_than_allocated() {
        let (mut a, b) = tokio::io::duplex(64);
        let mut w = Wire::new(b);
        let feed = tokio::spawn(async move {
            let mut head = [0u8; HEADER];
            head[8..].copy_from_slice(&u64::MAX.to_be_bytes());
            let _ = a.write_all(&head).await;
        });
        assert_eq!(
            w.next().await.expect_err("refused").kind(),
            io::ErrorKind::InvalidData
        );
        feed.await.expect("no panic");
    }
}