hatchery_node/
restream_session.rs

1//! Node end of a live restreamer session (A2 / A3).
2//!
3//! The node dials C2's uplink, OPENs a channel, then runs the inner Noise_NK
4//! handshake toward the relayed recipient key as ordinary DATA frames:
5//! msg1 goes up through C2 into the core drain, msg2 comes back down the same
6//! stream. After that every request and reply is Noise transport ciphertext,
7//! wrapped in C2's outer layer on the way up. C2 only ever sees opaque bytes,
8//! and the node never learns what sits behind C2.
9
10use crate::uplink_journal::UplinkJournal;
11use hatchery_core::inner::{decode_inner, encode_inner, InnerKind};
12use hatchery_core::nested::{node_seal_outer, NoiseInitiator, NoiseTransport, PublicKeyBytes};
13use hatchery_core::{ChannelKind, Frame, FrameError, FrameType, OpenPayload};
14use serde_json::{json, Value};
15use std::net::SocketAddr;
16use std::sync::Arc;
17use std::time::Duration;
18use tokio::io::{AsyncReadExt, AsyncWriteExt};
19use tokio::net::TcpStream;
20use tokio::sync::mpsc;
21use tokio::task::JoinHandle;
22
23/// Keys C2 hands a node at enrollment: its own outer key and the upstream recipient key.
24#[derive(Clone, Copy, Debug)]
25pub struct UpstreamKeys {
26    pub c2_public: PublicKeyBytes,
27    pub recipient_public: PublicKeyBytes,
28}
29
30/// One TCP uplink to C2 carrying any number of channels. Downlink DATA frames are
31/// demultiplexed by channel id. C2 keys its downlink route by node, so every channel a
32/// node opens must share this single connection.
33pub struct NodeUplink {
34    out: mpsc::Sender<Frame>,
35    routes: Arc<std::sync::Mutex<std::collections::BTreeMap<u32, mpsc::UnboundedSender<Vec<u8>>>>>,
36    journal: Option<Arc<UplinkJournal>>,
37    tasks: Vec<JoinHandle<()>>,
38}
39
40impl Drop for NodeUplink {
41    fn drop(&mut self) {
42        for t in &self.tasks {
43            t.abort();
44        }
45    }
46}
47
48impl NodeUplink {
49    pub async fn connect(uplink: SocketAddr, journal: Option<Arc<UplinkJournal>>) -> Result<Arc<Self>, String> {
50        let stream = TcpStream::connect(uplink).await.map_err(|e| e.to_string())?;
51        let _ = stream.set_nodelay(true);
52        let (mut rd, mut wr) = stream.into_split();
53        let (out_tx, mut out_rx) = mpsc::channel::<Frame>(64);
54        let routes: Arc<std::sync::Mutex<std::collections::BTreeMap<u32, mpsc::UnboundedSender<Vec<u8>>>>> = Arc::default();
55        let writer = tokio::spawn(async move {
56            while let Some(frame) = out_rx.recv().await {
57                let Ok(bytes) = frame.encode() else { break };
58                if wr.write_all(&bytes).await.is_err() {
59                    break;
60                }
61            }
62        });
63        let reader_out = out_tx.clone();
64        let reader_journal = journal.clone();
65        let reader_routes = Arc::clone(&routes);
66        // The reader never blocks: per-channel queues are unbounded (a stalled reader
67        // stalls C2's downlink write and with it our uplink: a deadlock).
68        let reader = tokio::spawn(async move {
69            let mut buf = vec![0u8; 128 * 1024];
70            let mut filled = 0usize;
71            loop {
72                let Ok(n) = rd.read(&mut buf[filled..]).await else { break };
73                if n == 0 {
74                    break;
75                }
76                filled += n;
77                loop {
78                    match Frame::decode(&buf[..filled]) {
79                        Ok((frame, used)) => {
80                            buf.copy_within(used..filled, 0);
81                            filled -= used;
82                            match frame.frame_type {
83                                FrameType::Data => {
84                                    let tx = reader_routes.lock().expect("routes").get(&frame.channel_id).cloned();
85                                    if let Some(tx) = tx {
86                                        let _ = tx.send(frame.payload);
87                                    }
88                                }
89                                FrameType::JournalReq => {
90                                    let Some(j) = reader_journal.as_ref() else { continue };
91                                    let limit = frame.journal_req_limit().unwrap_or(64) as usize;
92                                    if let Ok(body) = j.journal().reply_body(frame.channel_id, frame.seq, limit) {
93                                        if let Ok(reply) = Frame::journal_reply(frame.channel_id, &body) {
94                                            let _ = reader_out.send(reply).await;
95                                        }
96                                    }
97                                }
98                                FrameType::GoAway | FrameType::Close => {
99                                    reader_routes.lock().expect("routes").clear();
100                                    return;
101                                }
102                                _ => {}
103                            }
104                        }
105                        Err(FrameError::Truncated) => break,
106                        Err(_) => {
107                            reader_routes.lock().expect("routes").clear();
108                            return;
109                        }
110                    }
111                }
112                if filled == buf.len() {
113                    break;
114                }
115            }
116            // Dropping every route sender ends each session's recv with "uplink closed".
117            reader_routes.lock().expect("routes").clear();
118        });
119        Ok(Arc::new(Self { out: out_tx, routes, journal, tasks: vec![writer, reader] }))
120    }
121
122    /// OPEN `channel_id` and complete the Noise_NK handshake on it.
123    pub async fn open(
124        self: &Arc<Self>,
125        node_id: &str,
126        channel_id: u32,
127        kind: ChannelKind,
128        keys: UpstreamKeys,
129        handshake_timeout: Duration,
130    ) -> Result<NodeSession, String> {
131        let (down_tx, mut down_rx) = mpsc::unbounded_channel::<Vec<u8>>();
132        self.routes.lock().expect("routes").insert(channel_id, down_tx);
133        let open = Frame::open(channel_id, &OpenPayload { channel_kind: kind, node_id: node_id.to_owned() }).map_err(|e| e.to_string())?;
134        self.out.send(open).await.map_err(|_| "uplink closed".to_owned())?;
135        let mut initiator = NoiseInitiator::new(keys.recipient_public).map_err(|e| e.to_string())?;
136        let msg1 = initiator.write_message1().map_err(|e| e.to_string())?;
137        let mut tmp = Partial { channel_id, keys, out: self.out.clone(), journal: self.journal.clone() };
138        tmp.send_inner(&encode_inner(InnerKind::Handshake1, &msg1)).await?;
139        let msg2 = tokio::time::timeout(handshake_timeout, async {
140            loop {
141                let Some(bytes) = down_rx.recv().await else { return Err("uplink closed".to_owned()) };
142                if let Some((InnerKind::Handshake2, body)) = decode_inner(&bytes) {
143                    return Ok(body.to_vec());
144                }
145            }
146        })
147        .await
148        .map_err(|_| "handshake timeout".to_owned())??;
149        let transport = initiator.read_message2(&msg2).map_err(|e| e.to_string())?;
150        Ok(NodeSession {
151            channel_id,
152            keys,
153            out: self.out.clone(),
154            down: down_rx,
155            transport,
156            journal: self.journal.clone(),
157            _uplink: Arc::clone(self),
158        })
159    }
160}
161
162pub struct NodeSession {
163    channel_id: u32,
164    keys: UpstreamKeys,
165    out: mpsc::Sender<Frame>,
166    down: mpsc::UnboundedReceiver<Vec<u8>>,
167    transport: NoiseTransport,
168    journal: Option<Arc<UplinkJournal>>,
169    _uplink: Arc<NodeUplink>,
170}
171
172impl NodeSession {
173    /// Dial a dedicated uplink, OPEN one channel, and complete the Noise_NK handshake.
174    pub async fn connect(
175        uplink: SocketAddr,
176        node_id: &str,
177        channel_id: u32,
178        kind: ChannelKind,
179        keys: UpstreamKeys,
180        journal: Option<Arc<UplinkJournal>>,
181        handshake_timeout: Duration,
182    ) -> Result<Self, String> {
183        let link = NodeUplink::connect(uplink, journal).await?;
184        link.open(node_id, channel_id, kind, keys, handshake_timeout).await
185    }
186
187    /// The shared uplink this session rides on (to open sibling channels).
188    pub fn uplink(&self) -> Arc<NodeUplink> {
189        Arc::clone(&self._uplink)
190    }
191
192    /// Encrypt and send one inner message toward core.
193    pub async fn send(&mut self, plaintext: &[u8]) -> Result<(), String> {
194        let ct = self.transport.encrypt(plaintext).map_err(|e| e.to_string())?;
195        let inner = encode_inner(InnerKind::Transport, &ct);
196        Partial { channel_id: self.channel_id, keys: self.keys, out: self.out.clone(), journal: self.journal.clone() }
197            .send_inner(&inner)
198            .await
199    }
200
201    /// Next decrypted message from core.
202    pub async fn recv(&mut self, timeout: Duration) -> Result<Vec<u8>, String> {
203        tokio::time::timeout(timeout, async {
204            loop {
205                let Some(bytes) = self.down.recv().await else { return Err("uplink closed".to_owned()) };
206                if let Some((InnerKind::Transport, ct)) = decode_inner(&bytes) {
207                    return self.transport.decrypt(ct).map_err(|e| e.to_string());
208                }
209            }
210        })
211        .await
212        .map_err(|_| "reply timeout".to_owned())?
213    }
214
215    /// Request/response helper for node calls (`mail.*`, `task.*`). The node id is never
216    /// in the payload: core takes it from the stream the envelope arrived on.
217    pub async fn call(&mut self, op: &str, payload: &Value, timeout: Duration) -> Result<Value, String> {
218        let req = serde_json::to_vec(&json!({ "op": op, "payload": payload })).map_err(|e| e.to_string())?;
219        self.send(&req).await?;
220        let reply = self.recv(timeout).await?;
221        let reply: Value = serde_json::from_slice(&reply).map_err(|e| e.to_string())?;
222        let status = reply.get("status").and_then(Value::as_u64).unwrap_or(500);
223        let body = reply.get("body").cloned().unwrap_or(Value::Null);
224        if status == 200 {
225            Ok(body)
226        } else {
227            let msg = body.get("error").and_then(Value::as_str).unwrap_or("error");
228            Err(format!("{status}: {msg}"))
229        }
230    }
231}
232
233struct Partial {
234    channel_id: u32,
235    keys: UpstreamKeys,
236    out: mpsc::Sender<Frame>,
237    journal: Option<Arc<UplinkJournal>>,
238}
239
240impl Partial {
241    async fn send_inner(&mut self, inner: &[u8]) -> Result<(), String> {
242        // Journal the inner bytes (what C2 forwards after its layer) for gap recovery.
243        if let Some(j) = &self.journal {
244            j.append(self.channel_id, inner)?;
245        }
246        let sealed = node_seal_outer(self.keys.c2_public, inner).map_err(|e| e.to_string())?;
247        let frame = Frame::data(self.channel_id, 0, sealed).map_err(|e| e.to_string())?;
248        self.out.send(frame).await.map_err(|_| "uplink closed".to_owned())
249    }
250}
hatchery_node/restream_session.rs

hatchery_node/
restream_session.rs
