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_c2_protocol::inner::{decode_inner, encode_inner, InnerKind};
12use hatchery_c2_protocol::nested::{node_seal_outer, NoiseInitiator, NoiseTransport, PublicKeyBytes};
13use hatchery_c2_protocol::{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
30pub struct NodeSession {
31    channel_id: u32,
32    keys: UpstreamKeys,
33    out: mpsc::Sender<Frame>,
34    down: mpsc::Receiver<Vec<u8>>,
35    transport: NoiseTransport,
36    journal: Option<Arc<UplinkJournal>>,
37    tasks: Vec<JoinHandle<()>>,
38}
39
40impl Drop for NodeSession {
41    fn drop(&mut self) {
42        for t in &self.tasks {
43            t.abort();
44        }
45    }
46}
47
48impl NodeSession {
49    /// Dial, OPEN, and complete the Noise_NK handshake over DATA frames.
50    pub async fn connect(
51        uplink: SocketAddr,
52        node_id: &str,
53        channel_id: u32,
54        kind: ChannelKind,
55        keys: UpstreamKeys,
56        journal: Option<Arc<UplinkJournal>>,
57        handshake_timeout: Duration,
58    ) -> Result<Self, String> {
59        let stream = TcpStream::connect(uplink).await.map_err(|e| e.to_string())?;
60        let (mut rd, mut wr) = stream.into_split();
61        let (out_tx, mut out_rx) = mpsc::channel::<Frame>(64);
62        let (down_tx, down_rx) = mpsc::channel::<Vec<u8>>(64);
63
64        let writer = tokio::spawn(async move {
65            while let Some(frame) = out_rx.recv().await {
66                let Ok(bytes) = frame.encode() else { break };
67                if wr.write_all(&bytes).await.is_err() {
68                    break;
69                }
70            }
71        });
72        let reader_out = out_tx.clone();
73        let reader_journal = journal.clone();
74        let reader = tokio::spawn(async move {
75            let mut buf = vec![0u8; 128 * 1024];
76            let mut filled = 0usize;
77            loop {
78                let Ok(n) = rd.read(&mut buf[filled..]).await else { break };
79                if n == 0 {
80                    break;
81                }
82                filled += n;
83                loop {
84                    match Frame::decode(&buf[..filled]) {
85                        Ok((frame, used)) => {
86                            buf.copy_within(used..filled, 0);
87                            filled -= used;
88                            match frame.frame_type {
89                                FrameType::Data => {
90                                    if down_tx.send(frame.payload).await.is_err() {
91                                        return;
92                                    }
93                                }
94                                FrameType::JournalReq => {
95                                    let Some(j) = reader_journal.as_ref() else { continue };
96                                    let limit = frame.journal_req_limit().unwrap_or(64) as usize;
97                                    if let Ok(body) = j.journal().reply_body(frame.channel_id, frame.seq, limit) {
98                                        if let Ok(reply) = Frame::journal_reply(frame.channel_id, &body) {
99                                            let _ = reader_out.send(reply).await;
100                                        }
101                                    }
102                                }
103                                FrameType::GoAway | FrameType::Close => return,
104                                _ => {}
105                            }
106                        }
107                        Err(FrameError::Truncated) => break,
108                        Err(_) => return,
109                    }
110                }
111                if filled == buf.len() {
112                    return;
113                }
114            }
115        });
116
117        let open = Frame::open(
118            channel_id,
119            &OpenPayload { channel_kind: kind, node_id: node_id.to_owned() },
120        )
121        .map_err(|e| e.to_string())?;
122        out_tx.send(open).await.map_err(|_| "uplink closed".to_owned())?;
123
124        let mut initiator = NoiseInitiator::new(keys.recipient_public).map_err(|e| e.to_string())?;
125        let msg1 = initiator.write_message1().map_err(|e| e.to_string())?;
126        let mut session_tmp = Partial { channel_id, keys, out: out_tx.clone(), journal: journal.clone() };
127        session_tmp.send_inner(&encode_inner(InnerKind::Handshake1, &msg1)).await?;
128
129        let mut down_rx = down_rx;
130        let msg2 = tokio::time::timeout(handshake_timeout, async {
131            loop {
132                let Some(bytes) = down_rx.recv().await else { return Err("uplink closed".to_owned()) };
133                if let Some((InnerKind::Handshake2, body)) = decode_inner(&bytes) {
134                    return Ok(body.to_vec());
135                }
136            }
137        })
138        .await
139        .map_err(|_| "handshake timeout".to_owned())??;
140        let transport = initiator.read_message2(&msg2).map_err(|e| e.to_string())?;
141
142        Ok(Self {
143            channel_id,
144            keys,
145            out: out_tx,
146            down: down_rx,
147            transport,
148            journal,
149            tasks: vec![writer, reader],
150        })
151    }
152
153    /// Encrypt and send one inner message toward core.
154    pub async fn send(&mut self, plaintext: &[u8]) -> Result<(), String> {
155        let ct = self.transport.encrypt(plaintext).map_err(|e| e.to_string())?;
156        let inner = encode_inner(InnerKind::Transport, &ct);
157        Partial { channel_id: self.channel_id, keys: self.keys, out: self.out.clone(), journal: self.journal.clone() }
158            .send_inner(&inner)
159            .await
160    }
161
162    /// Next decrypted message from core.
163    pub async fn recv(&mut self, timeout: Duration) -> Result<Vec<u8>, String> {
164        tokio::time::timeout(timeout, async {
165            loop {
166                let Some(bytes) = self.down.recv().await else { return Err("uplink closed".to_owned()) };
167                if let Some((InnerKind::Transport, ct)) = decode_inner(&bytes) {
168                    return self.transport.decrypt(ct).map_err(|e| e.to_string());
169                }
170            }
171        })
172        .await
173        .map_err(|_| "reply timeout".to_owned())?
174    }
175
176    /// Request/response helper for node calls (`mail.*`, `task.*`). The node id is never
177    /// in the payload: core takes it from the stream the envelope arrived on.
178    pub async fn call(&mut self, op: &str, payload: &Value, timeout: Duration) -> Result<Value, String> {
179        let req = serde_json::to_vec(&json!({ "op": op, "payload": payload })).map_err(|e| e.to_string())?;
180        self.send(&req).await?;
181        let reply = self.recv(timeout).await?;
182        let reply: Value = serde_json::from_slice(&reply).map_err(|e| e.to_string())?;
183        let status = reply.get("status").and_then(Value::as_u64).unwrap_or(500);
184        let body = reply.get("body").cloned().unwrap_or(Value::Null);
185        if status == 200 {
186            Ok(body)
187        } else {
188            let msg = body.get("error").and_then(Value::as_str).unwrap_or("error");
189            Err(format!("{status}: {msg}"))
190        }
191    }
192}
193
194struct Partial {
195    channel_id: u32,
196    keys: UpstreamKeys,
197    out: mpsc::Sender<Frame>,
198    journal: Option<Arc<UplinkJournal>>,
199}
200
201impl Partial {
202    async fn send_inner(&mut self, inner: &[u8]) -> Result<(), String> {
203        // Journal the inner bytes (what C2 forwards after its layer) for gap recovery.
204        if let Some(j) = &self.journal {
205            j.append(self.channel_id, inner)?;
206        }
207        let sealed = node_seal_outer(self.keys.c2_public, inner).map_err(|e| e.to_string())?;
208        let frame = Frame::data(self.channel_id, 0, sealed).map_err(|e| e.to_string())?;
209        self.out.send(frame).await.map_err(|_| "uplink closed".to_owned())
210    }
211}
hatchery_node/restream_session.rs

hatchery_node/
restream_session.rs
