Skip to main content

banc_host/net/
mod.rs

1//! Network transport for remote rigs: postcard-rpc nodes and the rig lease
2//! over token-authenticated TCP.
3//!
4//! A rig daemon (built by the essai from these library pieces) listens on one
5//! port; every connection starts with a `Hello` frame declaring a role:
6//!
7//! - [`Role::Node`]: the rest of the stream is postcard-rpc, same frames as
8//!   USB — the daemon is a banc node that happens to live across the network.
9//!   Hosts get an ordinary `HostClient` via [`connect_node`].
10//! - [`Role::Lease`]: the rest of the stream is the [`lease`] protocol,
11//!   arbitrating exclusive rig access between runners that may not share a
12//!   filesystem (the flock in `Rig::acquire` cannot reach across machines).
13//!
14//! Framing is a 4-byte little-endian length prefix per frame, both
15//! directions, capped at [`MAX_FRAME`]. The token is a shared secret; the
16//! daemon should sit behind an authenticated tunnel or a firewalled port, the
17//! token is the second factor, not the perimeter.
18
19pub mod lease;
20
21use postcard_rpc::header::VarSeqKind;
22use postcard_rpc::host_client::{HostClient, WireRx, WireSpawn, WireTx};
23use postcard_rpc::standard_icd::{WireError, ERROR_PATH};
24use serde::{Deserialize, Serialize};
25use std::io::{Read, Write};
26use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
27use tokio::net::TcpStream;
28
29/// Bumped on any breaking change to the handshake or framing.
30pub const NET_VERSION: u8 = 0;
31
32/// Upper bound on a single frame, both directions. Generous for RPC frames
33/// and small chunked payloads; a peer exceeding it is broken or hostile.
34pub const MAX_FRAME: usize = 256 * 1024;
35
36/// What a connection wants to be after the handshake.
37#[derive(Serialize, Deserialize, Debug, Clone, Copy, PartialEq, Eq)]
38pub enum Role {
39    Node,
40    Lease,
41}
42
43#[derive(Serialize, Deserialize, Debug)]
44struct Hello {
45    version: u8,
46    role: Role,
47    token: String,
48}
49
50#[derive(Serialize, Deserialize, Debug)]
51enum HelloReply {
52    Ok { rig_name: String },
53    Denied,
54}
55
56// --- framing ---
57
58pub async fn write_frame<W: AsyncWrite + Unpin>(w: &mut W, data: &[u8]) -> std::io::Result<()> {
59    debug_assert!(data.len() <= MAX_FRAME);
60    w.write_all(&(data.len() as u32).to_le_bytes()).await?;
61    w.write_all(data).await?;
62    w.flush().await
63}
64
65pub async fn read_frame<R: AsyncRead + Unpin>(r: &mut R) -> std::io::Result<Vec<u8>> {
66    let mut len = [0u8; 4];
67    r.read_exact(&mut len).await?;
68    let len = u32::from_le_bytes(len) as usize;
69    if len > MAX_FRAME {
70        return Err(std::io::Error::other(format!("frame of {len} bytes exceeds MAX_FRAME")));
71    }
72    let mut buf = vec![0u8; len];
73    r.read_exact(&mut buf).await?;
74    Ok(buf)
75}
76
77pub fn write_frame_sync<W: Write>(w: &mut W, data: &[u8]) -> std::io::Result<()> {
78    debug_assert!(data.len() <= MAX_FRAME);
79    w.write_all(&(data.len() as u32).to_le_bytes())?;
80    w.write_all(data)?;
81    w.flush()
82}
83
84pub fn read_frame_sync<R: Read>(r: &mut R) -> std::io::Result<Vec<u8>> {
85    let mut len = [0u8; 4];
86    r.read_exact(&mut len)?;
87    let len = u32::from_le_bytes(len) as usize;
88    if len > MAX_FRAME {
89        return Err(std::io::Error::other(format!("frame of {len} bytes exceeds MAX_FRAME")));
90    }
91    let mut buf = vec![0u8; len];
92    r.read_exact(&mut buf)?;
93    Ok(buf)
94}
95
96/// Length-then-content comparison without early exit on content, so a wrong
97/// token costs the same as a right one.
98fn token_eq(a: &str, b: &str) -> bool {
99    let (a, b) = (a.as_bytes(), b.as_bytes());
100    if a.len() != b.len() {
101        return false;
102    }
103    a.iter().zip(b).fold(0u8, |acc, (x, y)| acc | (x ^ y)) == 0
104}
105
106// --- client side ---
107
108async fn handshake_client(stream: &mut TcpStream, role: Role, token: &str) -> anyhow::Result<String> {
109    let hello = Hello { version: NET_VERSION, role, token: token.to_owned() };
110    write_frame(stream, &postcard::to_stdvec(&hello)?).await?;
111    let reply: HelloReply = postcard::from_bytes(&read_frame(stream).await?)?;
112    match reply {
113        HelloReply::Ok { rig_name } => Ok(rig_name),
114        HelloReply::Denied => anyhow::bail!("rig daemon denied the handshake (bad token?)"),
115    }
116}
117
118/// Connect to a remote banc node and return the postcard-rpc client, exactly
119/// as `Node::connect` does over USB. Must run on a tokio runtime (the wire
120/// workers are spawned onto it).
121pub async fn connect_node(addr: &str, token: &str) -> anyhow::Result<HostClient<WireError>> {
122    let mut stream = TcpStream::connect(addr)
123        .await
124        .map_err(|e| anyhow::anyhow!("connecting to node at {addr}: {e}"))?;
125    stream.set_nodelay(true)?;
126    handshake_client(&mut stream, Role::Node, token)
127        .await
128        .map_err(|e| anyhow::anyhow!("node handshake with {addr}: {e}"))?;
129    let (rx, tx) = stream.into_split();
130    Ok(HostClient::new_with_wire(
131        TcpWireTx(tx),
132        TcpWireRx(rx),
133        TokioSpawn,
134        VarSeqKind::Seq2,
135        ERROR_PATH,
136        8,
137    ))
138}
139
140struct TcpWireTx(tokio::net::tcp::OwnedWriteHalf);
141
142impl WireTx for TcpWireTx {
143    type Error = std::io::Error;
144    async fn send(&mut self, data: Vec<u8>) -> Result<(), Self::Error> {
145        write_frame(&mut self.0, &data).await
146    }
147}
148
149struct TcpWireRx(tokio::net::tcp::OwnedReadHalf);
150
151impl WireRx for TcpWireRx {
152    type Error = std::io::Error;
153    async fn receive(&mut self) -> Result<Vec<u8>, Self::Error> {
154        read_frame(&mut self.0).await
155    }
156}
157
158struct TokioSpawn;
159
160impl WireSpawn for TokioSpawn {
161    fn spawn(&mut self, fut: impl std::future::Future<Output = ()> + Send + 'static) {
162        tokio::spawn(fut);
163    }
164}
165
166// --- server side (for rig daemons composing these pieces) ---
167
168/// Run the server half of the handshake on a fresh connection. On success
169/// the stream is positioned at the first post-handshake frame and the caller
170/// dispatches on the role; on failure the peer got `Denied` and the
171/// connection should be dropped.
172pub async fn handshake_server(
173    stream: &mut TcpStream,
174    token: &str,
175    rig_name: &str,
176) -> anyhow::Result<Role> {
177    stream.set_nodelay(true)?;
178    let hello: Hello = postcard::from_bytes(&read_frame(stream).await?)?;
179    if hello.version != NET_VERSION || !token_eq(&hello.token, token) {
180        write_frame(stream, &postcard::to_stdvec(&HelloReply::Denied)?).await?;
181        anyhow::bail!(
182            "handshake denied: version {} (want {NET_VERSION}), token {}",
183            hello.version,
184            if token_eq(&hello.token, token) { "ok" } else { "mismatch" },
185        );
186    }
187    let reply = HelloReply::Ok { rig_name: rig_name.to_owned() };
188    write_frame(stream, &postcard::to_stdvec(&reply)?).await?;
189    Ok(hello.role)
190}
191
192/// Sync client half of the handshake, for the lease client's std stream.
193pub(crate) fn handshake_client_sync(
194    stream: &mut std::net::TcpStream,
195    role: Role,
196    token: &str,
197) -> anyhow::Result<()> {
198    let hello = Hello { version: NET_VERSION, role, token: token.to_owned() };
199    write_frame_sync(stream, &postcard::to_stdvec(&hello)?)?;
200    let reply: HelloReply = postcard::from_bytes(&read_frame_sync(stream)?)?;
201    match reply {
202        HelloReply::Ok { .. } => Ok(()),
203        HelloReply::Denied => anyhow::bail!("rig daemon denied the handshake (bad token?)"),
204    }
205}