hatchery_server/
node_controller.rs1use hatchery_core::{
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#[derive(Debug)]
25pub enum G4aMessage {
26 ServerFrame(C2ServerFrame),
27 Event(RoutedNodeEvent),
28 Response(RoutedNodeResponse),
29 Request(RoutedNodeRequest),
30 Json(Value),
32}
33
34pub fn is_g4a_channel(kind: ChannelKind) -> bool {
36 matches!(
37 kind,
38 ChannelKind::G4aControl | ChannelKind::G4aPty | ChannelKind::G4aEvents
39 )
40}
41
42pub 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
61pub 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 if envelope.frame.payload.is_empty() {
73 return Ok(None);
74 }
75 Ok(Some(parse_g4a_payload(&envelope.frame.payload)?))
76}
77
78pub 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#[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 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 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 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_core::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}