hatchery_server/
node_controller.rs

1//! Core-side node controller (A2): g4a parsing lives here, not on C2.
2//!
3//! C2's restreamer forwards opaque multiplex DATA. Core drains the stream,
4//! strips the inner Noise layer (caller-supplied transport), then this module
5//! decodes g4a frames (`RoutedNodeEvent` / `RoutedNodeResponse` / `C2ServerFrame`).
6
7use hatchery_core::{
8    BrowserClientOp, BrowserOpReply, BrowserStreamEvent, ChannelKind, DrainEnvelope,
9    RoutedNodeEvent, RoutedNodeRequest, RoutedNodeResponse, C2ServerFrame,
10};
11use serde_json::Value;
12
13#[derive(Debug, thiserror::Error, PartialEq, Eq)]
14pub enum NodeControllerError {
15    #[error("payload is not valid JSON")]
16    InvalidJson,
17    #[error("payload is not a known g4a frame")]
18    UnknownFrame,
19    #[error("channel kind does not carry g4a")]
20    NotG4a,
21}
22
23/// Decoded g4a traffic from one drain envelope (after outer/inner crypto).
24#[derive(Debug)]
25pub enum G4aMessage {
26    ServerFrame(C2ServerFrame),
27    Event(RoutedNodeEvent),
28    Response(RoutedNodeResponse),
29    Request(RoutedNodeRequest),
30    /// Opaque JSON that is not yet a typed frame (forward-compatible).
31    Json(Value),
32}
33
34/// Returns true when the channel kind is one of the g4a.* multiplex kinds.
35pub fn is_g4a_channel(kind: ChannelKind) -> bool {
36    matches!(
37        kind,
38        ChannelKind::G4aControl | ChannelKind::G4aPty | ChannelKind::G4aEvents
39    )
40}
41
42/// Parse opaque payload bytes as a g4a message. C2 must not call this.
43pub fn parse_g4a_payload(payload: &[u8]) -> Result<G4aMessage, NodeControllerError> {
44    let value: Value =
45        serde_json::from_slice(payload).map_err(|_| NodeControllerError::InvalidJson)?;
46    if let Ok(frame) = serde_json::from_value::<C2ServerFrame>(value.clone()) {
47        return Ok(G4aMessage::ServerFrame(frame));
48    }
49    if let Ok(event) = serde_json::from_value::<RoutedNodeEvent>(value.clone()) {
50        return Ok(G4aMessage::Event(event));
51    }
52    if let Ok(response) = serde_json::from_value::<RoutedNodeResponse>(value.clone()) {
53        return Ok(G4aMessage::Response(response));
54    }
55    if let Ok(request) = serde_json::from_value::<RoutedNodeRequest>(value.clone()) {
56        return Ok(G4aMessage::Request(request));
57    }
58    Ok(G4aMessage::Json(value))
59}
60
61/// Handle one drain envelope: only g4a channels are parsed; others stay opaque.
62pub fn handle_drain_envelope(
63    envelope: &DrainEnvelope,
64    channel_kind: Option<ChannelKind>,
65) -> Result<Option<G4aMessage>, NodeControllerError> {
66    if let Some(kind) = channel_kind {
67        if !is_g4a_channel(kind) {
68            return Err(NodeControllerError::NotG4a);
69        }
70    }
71    // DATA frames carry the payload; non-DATA is ignored here.
72    if envelope.frame.payload.is_empty() {
73        return Ok(None);
74    }
75    Ok(Some(parse_g4a_payload(&envelope.frame.payload)?))
76}
77
78/// Drain worker: pull envelopes from C2 and parse g4a on core.
79pub struct DrainWorker {
80    pub client: hatchery_c2_client::DrainClient,
81}
82
83impl DrainWorker {
84    pub fn new(client: hatchery_c2_client::DrainClient) -> Self {
85        Self { client }
86    }
87
88    pub async fn connect_and_open(
89        &self,
90        node_id: &str,
91        channel_id: u32,
92        kind: ChannelKind,
93        after: u64,
94    ) -> Result<hatchery_c2_client::DrainSession, std::io::Error> {
95        let mut session = self
96            .client
97            .connect()
98            .await
99            .map_err(|e| std::io::Error::other(e.to_string()))?;
100        session
101            .open(node_id, channel_id, kind)
102            .await
103            .map_err(|e| std::io::Error::other(e.to_string()))?;
104        session
105            .resume(after)
106            .await
107            .map_err(|e| std::io::Error::other(e.to_string()))?;
108        Ok(session)
109    }
110}
111
112/// Operator control routed through core (A2): core sends g4a requests down a node's
113/// control stream and correlates replies. C2 relays the bytes without parsing them.
114///
115/// Wire (inner plaintext): `{"op_id": u64, "g4a": <frame>}` both ways. Messages without
116/// `op_id` are node-initiated events and go to the events channel.
117#[derive(Clone)]
118pub struct OperatorLink {
119    tx: tokio::sync::mpsc::Sender<Vec<u8>>,
120    pending: std::sync::Arc<std::sync::Mutex<std::collections::HashMap<u64, tokio::sync::oneshot::Sender<G4aMessage>>>>,
121    browser_pending: std::sync::Arc<std::sync::Mutex<std::collections::HashMap<u64, tokio::sync::oneshot::Sender<BrowserOpReply>>>>,
122    next: std::sync::Arc<std::sync::atomic::AtomicU64>,
123    events: Option<tokio::sync::mpsc::UnboundedSender<G4aMessage>>,
124    browser_events: Option<tokio::sync::mpsc::UnboundedSender<BrowserStreamEvent>>,
125}
126
127#[derive(Debug, thiserror::Error, PartialEq, Eq)]
128pub enum OperatorError {
129    #[error("node control stream closed")]
130    Closed,
131    #[error("node did not answer in time")]
132    Timeout,
133    #[error("request does not serialize")]
134    Encode,
135}
136
137impl OperatorLink {
138    /// Returns the link and the outbound queue to hand to `DrainConsumer::with_outbound`.
139    pub fn new(events: Option<tokio::sync::mpsc::UnboundedSender<G4aMessage>>) -> (Self, tokio::sync::mpsc::Receiver<Vec<u8>>) {
140        Self::new_with_browser(events, None)
141    }
142
143    pub fn new_with_browser(
144        events: Option<tokio::sync::mpsc::UnboundedSender<G4aMessage>>,
145        browser_events: Option<tokio::sync::mpsc::UnboundedSender<BrowserStreamEvent>>,
146    ) -> (Self, tokio::sync::mpsc::Receiver<Vec<u8>>) {
147        let (tx, rx) = tokio::sync::mpsc::channel(64);
148        (
149            Self {
150                tx,
151                pending: Default::default(),
152                browser_pending: Default::default(),
153                next: std::sync::Arc::new(1.into()),
154                events,
155                browser_events,
156            },
157            rx,
158        )
159    }
160
161    /// Send a browser sidecar op and wait for the node's BrowserOpReply.
162    pub async fn browser_op(
163        &self,
164        op: &BrowserClientOp,
165        timeout: std::time::Duration,
166    ) -> Result<BrowserOpReply, OperatorError> {
167        let op_id = self.next.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
168        let bytes = serde_json::to_vec(&serde_json::json!({ "op_id": op_id, "browser": op }))
169            .map_err(|_| OperatorError::Encode)?;
170        let (done, wait) = tokio::sync::oneshot::channel();
171        self.browser_pending.lock().expect("pending").insert(op_id, done);
172        if self.tx.send(bytes).await.is_err() {
173            self.browser_pending.lock().expect("pending").remove(&op_id);
174            return Err(OperatorError::Closed);
175        }
176        match tokio::time::timeout(timeout, wait).await {
177            Ok(Ok(msg)) => Ok(msg),
178            Ok(Err(_)) => Err(OperatorError::Closed),
179            Err(_) => {
180                self.browser_pending.lock().expect("pending").remove(&op_id);
181                Err(OperatorError::Timeout)
182            }
183        }
184    }
185
186    /// Enqueue one g4a request toward the node now (preserving send order) and
187    /// return a future that waits for its reply, so callers can pipeline.
188    pub async fn request_pipelined<T: serde::Serialize>(
189        &self,
190        frame: &T,
191        timeout: std::time::Duration,
192    ) -> Result<impl std::future::Future<Output = Result<G4aMessage, OperatorError>> + Send + 'static, OperatorError> {
193        let op_id = self.next.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
194        let g4a = serde_json::to_value(frame).map_err(|_| OperatorError::Encode)?;
195        let bytes = serde_json::to_vec(&serde_json::json!({ "op_id": op_id, "g4a": g4a })).map_err(|_| OperatorError::Encode)?;
196        let (done, wait) = tokio::sync::oneshot::channel();
197        self.pending.lock().expect("pending").insert(op_id, done);
198        if self.tx.send(bytes).await.is_err() {
199            self.pending.lock().expect("pending").remove(&op_id);
200            return Err(OperatorError::Closed);
201        }
202        let pending = self.pending.clone();
203        Ok(async move {
204            match tokio::time::timeout(timeout, wait).await {
205                Ok(Ok(msg)) => Ok(msg),
206                Ok(Err(_)) => Err(OperatorError::Closed),
207                Err(_) => {
208                    pending.lock().expect("pending").remove(&op_id);
209                    Err(OperatorError::Timeout)
210                }
211            }
212        })
213    }
214
215    /// Enqueue a browser op now (ordered) and return a future for its reply.
216    pub async fn browser_op_pipelined(
217        &self,
218        op: &BrowserClientOp,
219        timeout: std::time::Duration,
220    ) -> Result<impl std::future::Future<Output = Result<BrowserOpReply, OperatorError>> + Send + 'static, OperatorError> {
221        let op_id = self.next.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
222        let bytes = serde_json::to_vec(&serde_json::json!({ "op_id": op_id, "browser": op }))
223            .map_err(|_| OperatorError::Encode)?;
224        let (done, wait) = tokio::sync::oneshot::channel();
225        self.browser_pending.lock().expect("pending").insert(op_id, done);
226        if self.tx.send(bytes).await.is_err() {
227            self.browser_pending.lock().expect("pending").remove(&op_id);
228            return Err(OperatorError::Closed);
229        }
230        let pending = self.browser_pending.clone();
231        Ok(async move {
232            match tokio::time::timeout(timeout, wait).await {
233                Ok(Ok(msg)) => Ok(msg),
234                Ok(Err(_)) => Err(OperatorError::Closed),
235                Err(_) => {
236                    pending.lock().expect("pending").remove(&op_id);
237                    Err(OperatorError::Timeout)
238                }
239            }
240        })
241    }
242
243    /// Send one g4a request (e.g. a `RoutedNodeRequest`) and wait for its reply.
244    pub async fn request<T: serde::Serialize>(&self, frame: &T, timeout: std::time::Duration) -> Result<G4aMessage, OperatorError> {
245        let op_id = self.next.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
246        let g4a = serde_json::to_value(frame).map_err(|_| OperatorError::Encode)?;
247        let bytes = serde_json::to_vec(&serde_json::json!({ "op_id": op_id, "g4a": g4a })).map_err(|_| OperatorError::Encode)?;
248        let (done, wait) = tokio::sync::oneshot::channel();
249        self.pending.lock().expect("pending").insert(op_id, done);
250        if self.tx.send(bytes).await.is_err() {
251            self.pending.lock().expect("pending").remove(&op_id);
252            return Err(OperatorError::Closed);
253        }
254        match tokio::time::timeout(timeout, wait).await {
255            Ok(Ok(msg)) => Ok(msg),
256            Ok(Err(_)) => Err(OperatorError::Closed),
257            Err(_) => {
258                self.pending.lock().expect("pending").remove(&op_id);
259                Err(OperatorError::Timeout)
260            }
261        }
262    }
263
264    /// Drain handler for the node's g4a control channel: parses g4a on core only.
265    pub fn handler(self) -> crate::drain_consumer::CoreHandler {
266        std::sync::Arc::new(move |_node: &str, kind: ChannelKind, plain: &[u8]| {
267            if !is_g4a_channel(kind) {
268                return None;
269            }
270            let Ok(v) = serde_json::from_slice::<Value>(plain) else { return None };
271            // Browser sidecar (same Noise channel, not g4a).
272            if let Some(browser) = v.get("browser") {
273                if let Some(id) = v.get("op_id").and_then(Value::as_u64) {
274                    if let Ok(reply) = serde_json::from_value::<BrowserOpReply>(browser.clone()) {
275                        if let Some(done) = self.browser_pending.lock().expect("pending").remove(&id) {
276                            let _ = done.send(reply);
277                        }
278                    }
279                } else if let Ok(ev) = serde_json::from_value::<BrowserStreamEvent>(browser.clone()) {
280                    if let Some(tx) = &self.browser_events {
281                        let _ = tx.send(ev);
282                    }
283                }
284                return None;
285            }
286            let g4a = v.get("g4a").cloned().unwrap_or(Value::Null);
287            let Ok(bytes) = serde_json::to_vec(&g4a) else { return None };
288            let Ok(msg) = parse_g4a_payload(&bytes) else { return None };
289            match v.get("op_id").and_then(Value::as_u64) {
290                Some(id) => {
291                    if let Some(done) = self.pending.lock().expect("pending").remove(&id) {
292                        let _ = done.send(msg);
293                    }
294                }
295                None => {
296                    if let Some(ev) = &self.events {
297                        let _ = ev.send(msg);
298                    }
299                }
300            }
301            None
302        })
303    }
304}
305
306#[cfg(test)]
307mod tests {
308    use super::*;
309    use hatchery_core::Frame;
310
311    #[test]
312    fn parses_unknown_json_as_forward_compatible() {
313        let msg = parse_g4a_payload(br#"{"hello":1}"#).unwrap();
314        assert!(matches!(msg, G4aMessage::Json(_)));
315    }
316
317    #[test]
318    fn rejects_non_g4a_channel() {
319        let env = DrainEnvelope {
320            node_id: "n".into(),
321            frame: Frame::data(1, 1, br#"{"x":1}"#.to_vec()).unwrap(),
322        };
323        let err = handle_drain_envelope(&env, Some(ChannelKind::M4a)).unwrap_err();
324        assert_eq!(err, NodeControllerError::NotG4a);
325    }
326
327    #[test]
328    fn g4a_channel_parses_payload() {
329        let env = DrainEnvelope {
330            node_id: "n".into(),
331            frame: Frame::data(1, 1, br#"{"x":1}"#.to_vec()).unwrap(),
332        };
333        let msg = handle_drain_envelope(&env, Some(ChannelKind::G4aEvents))
334            .unwrap()
335            .unwrap();
336        assert!(matches!(msg, G4aMessage::Json(_)));
337    }
338}
hatchery_server/node_controller.rs

hatchery_server/
node_controller.rs
