Skip to main content

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_c2_protocol::{
8    ChannelKind, DrainEnvelope, RoutedNodeEvent, RoutedNodeRequest, RoutedNodeResponse,
9    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    next: std::sync::Arc<std::sync::atomic::AtomicU64>,
122    events: Option<tokio::sync::mpsc::UnboundedSender<G4aMessage>>,
123}
124
125#[derive(Debug, thiserror::Error, PartialEq, Eq)]
126pub enum OperatorError {
127    #[error("node control stream closed")]
128    Closed,
129    #[error("node did not answer in time")]
130    Timeout,
131    #[error("request does not serialize")]
132    Encode,
133}
134
135impl OperatorLink {
136    /// Returns the link and the outbound queue to hand to `DrainConsumer::with_outbound`.
137    pub fn new(events: Option<tokio::sync::mpsc::UnboundedSender<G4aMessage>>) -> (Self, tokio::sync::mpsc::Receiver<Vec<u8>>) {
138        let (tx, rx) = tokio::sync::mpsc::channel(64);
139        (Self { tx, pending: Default::default(), next: std::sync::Arc::new(1.into()), events }, rx)
140    }
141
142    /// Send one g4a request (e.g. a `RoutedNodeRequest`) and wait for its reply.
143    pub async fn request<T: serde::Serialize>(&self, frame: &T, timeout: std::time::Duration) -> Result<G4aMessage, OperatorError> {
144        let op_id = self.next.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
145        let g4a = serde_json::to_value(frame).map_err(|_| OperatorError::Encode)?;
146        let bytes = serde_json::to_vec(&serde_json::json!({ "op_id": op_id, "g4a": g4a })).map_err(|_| OperatorError::Encode)?;
147        let (done, wait) = tokio::sync::oneshot::channel();
148        self.pending.lock().expect("pending").insert(op_id, done);
149        if self.tx.send(bytes).await.is_err() {
150            self.pending.lock().expect("pending").remove(&op_id);
151            return Err(OperatorError::Closed);
152        }
153        match tokio::time::timeout(timeout, wait).await {
154            Ok(Ok(msg)) => Ok(msg),
155            Ok(Err(_)) => Err(OperatorError::Closed),
156            Err(_) => {
157                self.pending.lock().expect("pending").remove(&op_id);
158                Err(OperatorError::Timeout)
159            }
160        }
161    }
162
163    /// Drain handler for the node's g4a control channel: parses g4a on core only.
164    pub fn handler(self) -> crate::drain_consumer::CoreHandler {
165        std::sync::Arc::new(move |_node: &str, kind: ChannelKind, plain: &[u8]| {
166            if !is_g4a_channel(kind) {
167                return None;
168            }
169            let Ok(v) = serde_json::from_slice::<Value>(plain) else { return None };
170            let g4a = v.get("g4a").cloned().unwrap_or(Value::Null);
171            let Ok(bytes) = serde_json::to_vec(&g4a) else { return None };
172            let Ok(msg) = parse_g4a_payload(&bytes) else { return None };
173            match v.get("op_id").and_then(Value::as_u64) {
174                Some(id) => {
175                    if let Some(done) = self.pending.lock().expect("pending").remove(&id) {
176                        let _ = done.send(msg);
177                    }
178                }
179                None => {
180                    if let Some(ev) = &self.events {
181                        let _ = ev.send(msg);
182                    }
183                }
184            }
185            None
186        })
187    }
188}
189
190#[cfg(test)]
191mod tests {
192    use super::*;
193    use hatchery_c2_protocol::Frame;
194
195    #[test]
196    fn parses_unknown_json_as_forward_compatible() {
197        let msg = parse_g4a_payload(br#"{"hello":1}"#).unwrap();
198        assert!(matches!(msg, G4aMessage::Json(_)));
199    }
200
201    #[test]
202    fn rejects_non_g4a_channel() {
203        let env = DrainEnvelope {
204            node_id: "n".into(),
205            frame: Frame::data(1, 1, br#"{"x":1}"#.to_vec()).unwrap(),
206        };
207        let err = handle_drain_envelope(&env, Some(ChannelKind::M4a)).unwrap_err();
208        assert_eq!(err, NodeControllerError::NotG4a);
209    }
210
211    #[test]
212    fn g4a_channel_parses_payload() {
213        let env = DrainEnvelope {
214            node_id: "n".into(),
215            frame: Frame::data(1, 1, br#"{"x":1}"#.to_vec()).unwrap(),
216        };
217        let msg = handle_drain_envelope(&env, Some(ChannelKind::G4aEvents))
218            .unwrap()
219            .unwrap();
220        assert!(matches!(msg, G4aMessage::Json(_)));
221    }
222}