clawft_kernel/workload_ctl/
session.rs1use std::net::SocketAddr;
20use std::sync::Arc;
21use std::time::Duration;
22
23use async_trait::async_trait;
24
25use crate::ipc::{KernelMessage, MessagePayload, MessageTarget};
26use crate::mesh::{MeshError, MeshStream};
27use crate::mesh_artifact::ArtifactExchange;
28use crate::mesh_artifact_transfer::{PeerLink, PeerSet, ServeStats};
29use crate::mesh_ipc::{MeshIpcEnvelope, MeshRequest};
30use crate::mesh_noise::EncryptedChannel;
31use crate::mesh_swarm_state::ServePeer;
32
33use super::host_service::WorkloadHostService;
34use super::msg::{SignedCtl, WORKLOAD_HOST_SERVICE};
35
36pub fn is_envelope(raw: &[u8]) -> bool {
38 raw.first() == Some(&b'{')
39}
40
41pub fn request_envelope(
43 source: &str,
44 dest: &str,
45 method: &str,
46 signed: &SignedCtl,
47 timeout: Duration,
48) -> MeshRequest {
49 let msg = KernelMessage::new(
50 0,
51 MessageTarget::ServiceMethod {
52 service: WORKLOAD_HOST_SERVICE.to_string(),
53 method: method.to_string(),
54 },
55 MessagePayload::Json(serde_json::to_value(signed).unwrap_or_default()),
56 );
57 MeshRequest::new(
58 MeshIpcEnvelope::new(source.into(), dest.into(), msg),
59 timeout,
60 )
61}
62
63pub fn ctl_payload(env: &MeshIpcEnvelope) -> Result<(String, SignedCtl), String> {
65 let method = match env.inner_target() {
66 MessageTarget::ServiceMethod { service, method } if service == WORKLOAD_HOST_SERVICE => {
67 method.clone()
68 }
69 other => {
70 return Err(format!(
71 "not addressed to {WORKLOAD_HOST_SERVICE}: {other:?}"
72 ));
73 }
74 };
75 let MessagePayload::Json(v) = &env.message.payload else {
76 return Err("workload.ctl payload must be JSON".into());
77 };
78 let signed: SignedCtl =
79 serde_json::from_value(v.clone()).map_err(|e| format!("bad signed payload: {e}"))?;
80 Ok((method, signed))
81}
82
83const NOISE_CHUNK: usize = 65_000;
86const MAX_NOISE_MESSAGE: usize = 17 * 1024 * 1024;
88
89pub struct NoiseStream {
94 ch: Box<dyn EncryptedChannel>,
95 partial: Vec<u8>,
96}
97
98impl NoiseStream {
99 pub fn new(ch: Box<dyn EncryptedChannel>) -> Self {
101 Self {
102 ch,
103 partial: Vec::new(),
104 }
105 }
106}
107
108#[async_trait]
109impl MeshStream for NoiseStream {
110 async fn send(&mut self, data: &[u8]) -> Result<(), MeshError> {
111 let mut chunks = data.chunks(NOISE_CHUNK).peekable();
112 if chunks.peek().is_none() {
113 return self.ch.send_encrypted(&[1]).await;
114 }
115 while let Some(c) = chunks.next() {
116 let mut m = Vec::with_capacity(c.len() + 1);
117 m.push(u8::from(chunks.peek().is_none()));
118 m.extend_from_slice(c);
119 self.ch.send_encrypted(&m).await?;
120 }
121 Ok(())
122 }
123 async fn recv(&mut self) -> Result<Vec<u8>, MeshError> {
124 loop {
125 let m = self.ch.recv_encrypted().await?;
126 let (flag, body) = m
127 .split_first()
128 .ok_or_else(|| MeshError::Transport("empty noise fragment".into()))?;
129 if self.partial.len() + body.len() > MAX_NOISE_MESSAGE {
130 self.partial.clear();
131 return Err(MeshError::MessageTooLarge {
132 size: MAX_NOISE_MESSAGE + 1,
133 max: MAX_NOISE_MESSAGE,
134 });
135 }
136 self.partial.extend_from_slice(body);
137 if *flag == 1 {
138 return Ok(std::mem::take(&mut self.partial));
139 }
140 }
141 }
142 async fn close(&mut self) -> Result<(), MeshError> {
143 self.ch.close().await
144 }
145 fn remote_addr(&self) -> Option<SocketAddr> {
146 None
147 }
148}
149
150pub async fn serve_connection(
153 stream: Box<dyn MeshStream>,
154 svc: Arc<WorkloadHostService>,
155) -> Result<(), MeshError> {
156 let mut peers = PeerSet::single(PeerLink::new("controller", stream));
157 loop {
158 let raw = match peers.links_mut()[0].stream.recv().await {
159 Ok(r) => r,
160 Err(MeshError::ConnectionClosed) => return Ok(()),
161 Err(e) => return Err(e),
162 };
163 if !is_envelope(&raw) {
164 let _ = peers.links_mut()[0].stream.close().await;
166 return Err(MeshError::Transport("unsolicited artifact frame".into()));
167 }
168 let env = MeshIpcEnvelope::from_bytes(&raw)
169 .map_err(|e| MeshError::Transport(format!("bad envelope: {e}")))?;
170 let reply_to = env.source_node.clone();
171 let (method, signed) = match ctl_payload(&env) {
172 Ok(x) => x,
173 Err(e) => return Err(MeshError::Transport(e)),
174 };
175 peers.links_mut()[0].peer_id = reply_to.clone();
176 let (resp, authenticated) = svc.handle_checked(&method, &signed, Some(&mut peers)).await;
177 let mut msg = KernelMessage::new(
178 0,
179 MessageTarget::ServiceMethod {
180 service: WORKLOAD_HOST_SERVICE.to_string(),
181 method: method.clone(),
182 },
183 MessagePayload::Json(serde_json::to_value(&resp).unwrap_or_default()),
184 );
185 msg.correlation_id = env.message.correlation_id.clone();
186 let out = MeshIpcEnvelope::new(svc.node_id().to_string(), reply_to, msg);
187 let bytes = out
188 .to_bytes()
189 .map_err(|e| MeshError::Transport(e.to_string()))?;
190 let link = &mut peers.links_mut()[0];
191 if link.is_dead() {
192 return Ok(()); }
194 link.stream.send(&bytes).await?;
195 if !authenticated {
196 let _ = link.stream.close().await;
198 return Ok(());
199 }
200 }
201}
202
203#[derive(Debug, thiserror::Error)]
205pub enum CallError {
206 #[error("transport: {0}")]
208 Transport(String),
209 #[error("timed out after {0:?}")]
211 Timeout(Duration),
212 #[error("protocol: {0}")]
214 Protocol(String),
215}
216
217impl From<MeshError> for CallError {
218 fn from(e: MeshError) -> Self {
219 CallError::Transport(e.to_string())
220 }
221}
222
223pub struct CtlConnection {
225 stream: Box<dyn MeshStream>,
226 local_node: String,
227 peer_verified: bool,
232}
233
234impl CtlConnection {
235 pub fn new(stream: Box<dyn MeshStream>, local_node: impl Into<String>) -> Self {
237 Self {
238 stream,
239 local_node: local_node.into(),
240 peer_verified: false,
241 }
242 }
243
244 pub fn with_peer_verified(mut self, verified: bool) -> Self {
247 self.peer_verified = verified;
248 self
249 }
250
251 pub async fn call(
255 &mut self,
256 dest: &str,
257 method: &str,
258 signed: &SignedCtl,
259 exchange: Option<&ArtifactExchange>,
260 timeout: Duration,
261 ) -> Result<SignedCtl, CallError> {
262 let req = request_envelope(&self.local_node, dest, method, signed, timeout);
263 let bytes = req
264 .request
265 .to_bytes()
266 .map_err(|e| CallError::Protocol(e.to_string()))?;
267 self.stream.send(&bytes).await?;
268 let deadline = tokio::time::Instant::now() + timeout;
269 let mut stats = ServeStats::default();
270 loop {
271 let raw = tokio::time::timeout_at(deadline, self.stream.recv())
272 .await
273 .map_err(|_| CallError::Timeout(timeout))??;
274 if !is_envelope(&raw) {
275 let ex = exchange.ok_or_else(|| {
276 CallError::Protocol("artifact frame but nothing to serve".into())
277 })?;
278 let who = ServePeer {
279 node_id: dest.to_owned(),
280 verified: self.peer_verified,
281 };
282 ex.serve_frame_as(self.stream.as_mut(), &who, &raw, &mut stats)
283 .await
284 .map_err(|e| CallError::Protocol(format!("serving payload: {e}")))?;
285 continue;
286 }
287 let env = MeshIpcEnvelope::from_bytes(&raw)
288 .map_err(|e| CallError::Protocol(e.to_string()))?;
289 if !req.matches_response(&env) {
290 return Err(CallError::Protocol("uncorrelated response".into()));
291 }
292 let (_, signed) = ctl_payload(&env).map_err(CallError::Protocol)?;
293 return Ok(signed);
294 }
295 }
296
297 pub async fn close(mut self) {
299 let _ = self.stream.close().await;
300 }
301}