Skip to main content

clawft_kernel/workload_ctl/
session.rs

1//! `workload.ctl` on the wire: [`MeshIpcEnvelope`]s carrying a
2//! [`KernelMessage`] addressed to `ServiceMethod { workload-host, <method> }`,
3//! correlated with [`MeshRequest`], plus the artifact piece protocol
4//! multiplexed on the same stream for fetch-before-load.
5//!
6//! One connection, controller to target:
7//!
8//! ```text
9//! controller                                  target (workload-host)
10//!   envelope(SignedCtl request)       ->
11//!                                     <-      artifact frames (target fetches
12//!   artifact frames (controller serves) ->     the package from the controller)
13//!                                     <-      envelope(SignedCtl response)
14//! ```
15//!
16//! Envelopes are JSON (first byte `{`); artifact frames start with a
17//! big-endian length whose first byte is 0 (frames are far below 16 MiB).
18
19use 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
36/// True if a raw frame is an IPC envelope (JSON) rather than an artifact frame.
37pub fn is_envelope(raw: &[u8]) -> bool {
38    raw.first() == Some(&b'{')
39}
40
41/// Wrap a signed request for `dest` as a correlated [`MeshRequest`].
42pub 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
63/// The signed control payload of an envelope addressed to `workload-host`.
64pub 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
83/// Largest plaintext per Noise message (a Noise message is at most 65535
84/// bytes including the 16-byte tag; one byte carries the fragment flag).
85const NOISE_CHUNK: usize = 65_000;
86/// Largest reassembled message (the IPC and artifact frame caps are lower).
87const MAX_NOISE_MESSAGE: usize = 17 * 1024 * 1024;
88
89/// A Noise channel as a [`MeshStream`]. Messages larger than one Noise
90/// message (artifact pieces) are split into fragments, each prefixed with
91/// a flag byte (1 = last), and reassembled on receipt. Partial fragments
92/// are kept on `self`, so `recv` stays cancel-safe.
93pub struct NoiseStream {
94    ch: Box<dyn EncryptedChannel>,
95    partial: Vec<u8>,
96}
97
98impl NoiseStream {
99    /// Wrap an established channel.
100    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
150/// Serve one connection as the target's `workload-host` until the peer
151/// closes it. Each request may fetch its payload over the same stream.
152pub 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            // Artifact frames are only valid while this side is fetching.
165            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(()); // a failed fetch closed the stream
193        }
194        link.stream.send(&bytes).await?;
195        if !authenticated {
196            // One signed refusal per unauthenticated connection, then drop it.
197            let _ = link.stream.close().await;
198            return Ok(());
199        }
200    }
201}
202
203/// Why a call did not produce a signed response.
204#[derive(Debug, thiserror::Error)]
205pub enum CallError {
206    /// Transport failure.
207    #[error("transport: {0}")]
208    Transport(String),
209    /// No response in time.
210    #[error("timed out after {0:?}")]
211    Timeout(Duration),
212    /// The peer spoke something else.
213    #[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
223/// Controller end of one connection.
224pub struct CtlConnection {
225    stream: Box<dyn MeshStream>,
226    local_node: String,
227    /// Whether admission verified the node at the other end of the stream
228    /// (the service-stamped `AdmittedPeer`, ADR-106 5.4). A claimed id is
229    /// never enough: the default is unverified, which the checkout
230    /// redistribution policy refuses to serve.
231    peer_verified: bool,
232}
233
234impl CtlConnection {
235    /// Wrap an open stream to a target.
236    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    /// State that admission verified the peer on this stream. Only a caller
245    /// that holds that evidence (a verified node connection) may set it.
246    pub fn with_peer_verified(mut self, verified: bool) -> Self {
247        self.peer_verified = verified;
248        self
249    }
250
251    /// Send one signed request to `dest` and wait for its correlated
252    /// response. Artifact frames that arrive meanwhile are served from
253    /// `exchange` (the target fetching the payload), if given.
254    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    /// Close the connection.
298    pub async fn close(mut self) {
299        let _ = self.stream.close().await;
300    }
301}