clawft_kernel/
mesh_artifact_tunnel.rs

1//! Artifact streams over machine-mesh deliveries (ADR-106 section 5.4).
2//!
3//! The daemon holds the [`ArtifactExchange`] but not the mesh connections:
4//! the machine mesh service owns those and hands the daemon `deliver` frames
5//! and takes its `send`s. This module carries the piece protocol across that
6//! seam. A [`TunnelStream`] is a [`MeshStream`] whose bytes travel as
7//! `mesh.artifact.tunnel` messages addressed to a peer node:
8//!
9//! - **Dial** ([`TunnelDialer`], the production [`PeerDialer`]): opens a
10//!   session id and speaks the fetcher side.
11//! - **Serve** (inbound): the first `req` frame of an unknown session starts
12//!   [`ArtifactExchange::serve_as`] with a [`ServePeer`] built from the
13//!   service-stamped origin. Only a verified `node` peer is served; the
14//!   redistribution policy still decides what.
15//!
16//! Replies are accepted only from the peer that was dialled and only while
17//! that peer is a verified node, so an unadmitted sender cannot inject
18//! frames into a fetch. Frames larger than [`CHUNK`] are split and
19//! reassembled (a mesh-local line is capped at 1 MiB).
20
21use std::net::SocketAddr;
22use std::sync::atomic::{AtomicU64, Ordering};
23use std::sync::{Arc, Mutex};
24
25use async_trait::async_trait;
26use dashmap::DashMap;
27use rand::RngCore;
28use serde::{Deserialize, Serialize};
29use tokio::sync::mpsc;
30
31use crate::error::KernelResult;
32use crate::ipc::{KernelMessage, MessagePayload, MessageTarget};
33use crate::mesh::{MeshError, MeshStream};
34use crate::mesh_admit::PeerClass;
35use crate::mesh_artifact::ArtifactExchange;
36use crate::mesh_artifact_types::MAX_ARTIFACT_FRAME;
37use crate::mesh_delivery::PeerCtx;
38use crate::mesh_swarm_fetch::PeerDialer;
39use crate::mesh_swarm_state::ServePeer;
40use crate::workload_pkg::codec::hex_encode;
41
42/// Topic of tunnel frames.
43pub const TOPIC_TUNNEL: &str = "mesh.artifact.tunnel";
44/// Raw bytes per tunnel message.
45pub const CHUNK: usize = 192 * 1024;
46/// Most serve sessions from one peer.
47pub const MAX_SERVE_PER_PEER: usize = 8;
48/// Most serve sessions in all.
49pub const MAX_SERVE_TOTAL: usize = 64;
50/// Most dialled sessions open at once.
51pub const MAX_DIAL_TOTAL: usize = 64;
52/// Frames queued per session before the sender is dropped.
53const SESSION_QUEUE: usize = 128;
54
55/// Sends one kernel message to a peer node through the machine mesh. The
56/// daemon implements it with its service link; tests with an in-memory net.
57#[async_trait]
58pub trait PeerSender: Send + Sync + 'static {
59    /// Deliver `msg` to `node_id`.
60    async fn send_to_node(&self, node_id: &str, msg: KernelMessage) -> KernelResult<()>;
61}
62
63/// Which side of the session a frame was sent by.
64#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
65#[serde(rename_all = "snake_case")]
66enum Dir {
67    /// The fetcher (the dialler) to the holder.
68    Req,
69    /// The holder back to the fetcher.
70    Rsp,
71}
72
73#[derive(Debug, Serialize, Deserialize)]
74#[serde(deny_unknown_fields)]
75struct TunnelFrame {
76    v: u32,
77    sid: String,
78    dir: Dir,
79    /// Chunk index within one stream frame.
80    #[serde(default)]
81    part: u32,
82    /// Last chunk of the frame.
83    #[serde(default)]
84    last: bool,
85    /// The stream was closed by the sender.
86    #[serde(default)]
87    close: bool,
88    /// Chunk bytes, hex.
89    #[serde(default)]
90    data: String,
91}
92
93/// Counters, all monotonic.
94#[derive(Debug, Default)]
95pub struct TunnelCounters {
96    /// Frames refused because the peer is not a verified node.
97    pub refused_unverified: AtomicU64,
98    /// Frames refused for a limit or a malformed shape.
99    pub refused_other: AtomicU64,
100    /// Serve sessions started.
101    pub served_sessions: AtomicU64,
102    /// Sessions dropped because their queue was full (a slow consumer).
103    pub dropped_slow: AtomicU64,
104}
105
106struct Session {
107    tx: mpsc::Sender<Vec<u8>>,
108    partial: Mutex<Vec<u8>>,
109    next_part: Mutex<u32>,
110}
111
112/// Both ends of the artifact tunnel for one node.
113pub struct ArtifactTunnel {
114    exchange: Arc<ArtifactExchange>,
115    sender: Arc<dyn PeerSender>,
116    /// Sessions this node serves, by (peer, sid).
117    serving: DashMap<(String, String), Arc<Session>>,
118    /// Sessions this node dialled, by (peer, sid).
119    dialled: DashMap<(String, String), Arc<Session>>,
120    /// Counters.
121    pub counters: TunnelCounters,
122}
123
124/// Whether `from` may be served licensed content or sent licence floods:
125/// admission verified it (the service-stamped `AdmittedPeer`) and its class
126/// is `node`. A verified leaf is not: leaves never get licensed bytes or
127/// grant floods. Every serve, checkout and flood decision goes through this.
128pub fn licensed_peer(from: &PeerCtx) -> bool {
129    from.node_verified && from.class == PeerClass::Node
130}
131
132/// Lower- or upper-case hex to bytes; `None` on odd length or a bad digit.
133fn hex_decode(s: &str) -> Option<Vec<u8>> {
134    if !s.len().is_multiple_of(2) {
135        return None;
136    }
137    s.as_bytes()
138        .chunks(2)
139        .map(|p| {
140            let hi = (p[0] as char).to_digit(16)?;
141            let lo = (p[1] as char).to_digit(16)?;
142            Some((hi * 16 + lo) as u8)
143        })
144        .collect()
145}
146
147fn new_sid() -> String {
148    let mut b = [0u8; 8];
149    rand::rngs::OsRng.fill_bytes(&mut b);
150    hex_encode(&b)
151}
152
153impl ArtifactTunnel {
154    /// A tunnel serving `exchange` and sending through `sender`.
155    pub fn new(exchange: Arc<ArtifactExchange>, sender: Arc<dyn PeerSender>) -> Arc<Self> {
156        Arc::new(Self {
157            exchange,
158            sender,
159            serving: DashMap::new(),
160            dialled: DashMap::new(),
161            counters: TunnelCounters::default(),
162        })
163    }
164
165    /// The production [`PeerDialer`] over this tunnel.
166    pub fn dialer(self: &Arc<Self>) -> Arc<dyn PeerDialer> {
167        Arc::new(TunnelDialer { tunnel: self.clone() })
168    }
169
170    /// Serve sessions open now.
171    pub fn serving_sessions(&self) -> usize {
172        self.serving.len()
173    }
174
175    /// Handle one inbound `mesh.artifact.tunnel` delivery. `from` is built
176    /// from the service-stamped origin; anything but a verified node is
177    /// dropped and counted.
178    pub async fn on_message(self: &Arc<Self>, from: &PeerCtx, msg: KernelMessage) {
179        if !licensed_peer(from) {
180            self.counters.refused_unverified.fetch_add(1, Ordering::Relaxed);
181            return;
182        }
183        let MessagePayload::Json(v) = msg.payload else {
184            self.counters.refused_other.fetch_add(1, Ordering::Relaxed);
185            return;
186        };
187        let Ok(f) = serde_json::from_value::<TunnelFrame>(v) else {
188            self.counters.refused_other.fetch_add(1, Ordering::Relaxed);
189            return;
190        };
191        if f.v != 1 || f.sid.is_empty() || f.sid.len() > 32 || f.data.len() > CHUNK * 2 {
192            self.counters.refused_other.fetch_add(1, Ordering::Relaxed);
193            return;
194        }
195        let key = (from.peer_id.clone(), f.sid.clone());
196        if f.close {
197            // Dropping the registry's handle closes the session's receiver,
198            // which ends the serve loop or fails the fetch.
199            match f.dir {
200                Dir::Req => self.serving.remove(&key),
201                Dir::Rsp => self.dialled.remove(&key),
202            };
203            return;
204        }
205        match f.dir {
206            Dir::Req => self.inbound_req(key, f).await,
207            Dir::Rsp => {
208                if let Some(s) = self.dialled.get(&key).map(|s| s.clone()) {
209                    self.feed(&key, Dir::Rsp, &s, &f);
210                } else {
211                    self.counters.refused_other.fetch_add(1, Ordering::Relaxed);
212                }
213            }
214        }
215    }
216
217    async fn inbound_req(self: &Arc<Self>, key: (String, String), f: TunnelFrame) {
218        let existing = self.serving.get(&key).map(|s| s.clone());
219        let session = match existing {
220            Some(s) => s,
221            None => {
222                let per_peer = self.serving.iter().filter(|e| e.key().0 == key.0).count();
223                if per_peer >= MAX_SERVE_PER_PEER || self.serving.len() >= MAX_SERVE_TOTAL {
224                    self.counters.refused_other.fetch_add(1, Ordering::Relaxed);
225                    return;
226                }
227                let (tx, rx) = mpsc::channel(SESSION_QUEUE);
228                let s = Arc::new(Session {
229                    tx,
230                    partial: Mutex::default(),
231                    next_part: Mutex::default(),
232                });
233                self.serving.insert(key.clone(), s.clone());
234                self.counters.served_sessions.fetch_add(1, Ordering::Relaxed);
235                let me = self.clone();
236                let (peer, sid) = key.clone();
237                tokio::spawn(async move {
238                    let mut stream = TunnelStream::new(me.clone(), peer.clone(), sid.clone(), Dir::Rsp, rx);
239                    let who = ServePeer::verified(peer.clone());
240                    let _ = me.exchange.serve_as(&mut stream, &who).await;
241                    me.serving.remove(&(peer, sid));
242                });
243                s
244            }
245        };
246        self.feed(&key, Dir::Req, &session, &f);
247    }
248
249    /// Reassemble `f` into whole stream frames and queue them for the session.
250    /// Never waits: this runs on the daemon's single delivery worker, so a
251    /// session whose queue is full is dropped (counted) instead of stalling
252    /// every other delivery behind it.
253    fn feed(&self, key: &(String, String), dir: Dir, s: &Session, f: &TunnelFrame) {
254        let Some(bytes) = hex_decode(&f.data) else { return };
255        let whole = {
256            let mut part = s.next_part.lock().unwrap_or_else(|p| p.into_inner());
257            let mut buf = s.partial.lock().unwrap_or_else(|p| p.into_inner());
258            if f.part != *part || buf.len() + bytes.len() > MAX_ARTIFACT_FRAME + 64 {
259                buf.clear();
260                *part = 0;
261                return;
262            }
263            buf.extend_from_slice(&bytes);
264            if f.last {
265                *part = 0;
266                Some(std::mem::take(&mut *buf))
267            } else {
268                *part += 1;
269                None
270            }
271        };
272        if let Some(w) = whole
273            && s.tx.try_send(w).is_err()
274        {
275            self.counters.dropped_slow.fetch_add(1, Ordering::Relaxed);
276            match dir {
277                Dir::Req => self.serving.remove(key),
278                Dir::Rsp => self.dialled.remove(key),
279            };
280        }
281    }
282}
283
284struct TunnelDialer {
285    tunnel: Arc<ArtifactTunnel>,
286}
287
288#[async_trait]
289impl PeerDialer for TunnelDialer {
290    async fn dial(&self, peer_id: &str) -> Result<Box<dyn MeshStream>, MeshError> {
291        let t = &self.tunnel;
292        if t.dialled.len() >= MAX_DIAL_TOTAL {
293            return Err(MeshError::Transport("too many artifact sessions open".into()));
294        }
295        let sid = new_sid();
296        let (tx, rx) = mpsc::channel(SESSION_QUEUE);
297        let s = Arc::new(Session {
298            tx,
299            partial: Mutex::default(),
300            next_part: Mutex::default(),
301        });
302        t.dialled.insert((peer_id.to_owned(), sid.clone()), s);
303        Ok(Box::new(TunnelStream::new(t.clone(), peer_id.to_owned(), sid, Dir::Req, rx)))
304    }
305}
306
307/// One artifact session as a [`MeshStream`].
308pub struct TunnelStream {
309    tunnel: Arc<ArtifactTunnel>,
310    peer: String,
311    sid: String,
312    /// The direction this end sends in.
313    out: Dir,
314    rx: mpsc::Receiver<Vec<u8>>,
315    closed: bool,
316}
317
318impl TunnelStream {
319    fn new(
320        tunnel: Arc<ArtifactTunnel>,
321        peer: String,
322        sid: String,
323        out: Dir,
324        rx: mpsc::Receiver<Vec<u8>>,
325    ) -> Self {
326        Self { tunnel, peer, sid, out, rx, closed: false }
327    }
328
329    fn registry_remove(&self) {
330        let key = (self.peer.clone(), self.sid.clone());
331        match self.out {
332            Dir::Req => self.tunnel.dialled.remove(&key),
333            Dir::Rsp => self.tunnel.serving.remove(&key),
334        };
335    }
336
337    async fn emit(&self, f: TunnelFrame) -> Result<(), MeshError> {
338        let value = serde_json::to_value(&f).map_err(|e| MeshError::Io(e.to_string()))?;
339        let msg = KernelMessage::new(
340            0,
341            MessageTarget::Topic(TOPIC_TUNNEL.into()),
342            MessagePayload::Json(value),
343        );
344        self.tunnel
345            .sender
346            .send_to_node(&self.peer, msg)
347            .await
348            .map_err(|e| MeshError::Transport(e.to_string()))
349    }
350}
351
352#[async_trait]
353impl MeshStream for TunnelStream {
354    async fn send(&mut self, data: &[u8]) -> Result<(), MeshError> {
355        if self.closed {
356            return Err(MeshError::ConnectionClosed);
357        }
358        let n = data.chunks(CHUNK).count().max(1);
359        for (i, chunk) in data.chunks(CHUNK).chain(data.is_empty().then_some(&[][..])).enumerate() {
360            self.emit(TunnelFrame {
361                v: 1,
362                sid: self.sid.clone(),
363                dir: self.out,
364                part: i as u32,
365                last: i + 1 == n,
366                close: false,
367                data: hex_encode(chunk),
368            })
369            .await?;
370        }
371        Ok(())
372    }
373
374    async fn recv(&mut self) -> Result<Vec<u8>, MeshError> {
375        self.rx.recv().await.ok_or(MeshError::ConnectionClosed)
376    }
377
378    async fn close(&mut self) -> Result<(), MeshError> {
379        if std::mem::replace(&mut self.closed, true) {
380            return Ok(());
381        }
382        self.registry_remove();
383        let _ = self
384            .emit(TunnelFrame {
385                v: 1,
386                sid: self.sid.clone(),
387                dir: self.out,
388                part: 0,
389                last: true,
390                close: true,
391                data: String::new(),
392            })
393            .await;
394        Ok(())
395    }
396
397    fn remote_addr(&self) -> Option<SocketAddr> {
398        None
399    }
400}
401
402impl Drop for TunnelStream {
403    fn drop(&mut self) {
404        if !self.closed {
405            self.registry_remove();
406        }
407    }
408}
409
410#[cfg(test)]
411mod tests {
412    use std::time::Duration;
413
414    use super::*;
415    use crate::artifact_store::ArtifactStore;
416    use crate::mesh_artifact::ExchangeConfig;
417
418    struct Null;
419
420    #[async_trait]
421    impl PeerSender for Null {
422        async fn send_to_node(&self, _: &str, _: KernelMessage) -> KernelResult<()> {
423            Ok(())
424        }
425    }
426
427    fn node_ctx(id: &str) -> PeerCtx {
428        PeerCtx {
429            peer_id: id.into(),
430            node_verified: true,
431            class: PeerClass::Node,
432            remote_static: None,
433            src_scope: None,
434        }
435    }
436
437    fn rsp(sid: &str, part: u32) -> KernelMessage {
438        KernelMessage::new(
439            0,
440            MessageTarget::Topic(TOPIC_TUNNEL.into()),
441            MessagePayload::Json(serde_json::json!({
442                "v": 1, "sid": sid, "dir": "rsp", "part": part, "last": true, "data": "00"
443            })),
444        )
445    }
446
447    #[tokio::test]
448    async fn a_session_that_stops_reading_is_dropped_and_never_blocks_the_delivery_worker() {
449        let ex = Arc::new(
450            ArtifactExchange::new("x", Arc::new(ArtifactStore::new_memory()), ExchangeConfig::default())
451                .unwrap(),
452        );
453        let t = ArtifactTunnel::new(ex, Arc::new(Null));
454        // Dial a peer and never read: find the session id the tunnel made.
455        let _stream = t.dialer().dial("peer-a").await.unwrap();
456        let sid = t.dialled.iter().next().map(|e| e.key().1.clone()).unwrap();
457        // Far more whole frames than the queue holds, delivered inline: every
458        // call must return at once (a plain await on a full queue would hang
459        // this test, and the daemon's only delivery worker with it).
460        let all = tokio::time::timeout(Duration::from_secs(5), async {
461            for _ in 0..(SESSION_QUEUE + 10) {
462                t.on_message(&node_ctx("peer-a"), rsp(&sid, 0)).await;
463            }
464        })
465        .await;
466        assert!(all.is_ok(), "delivery blocked on a slow session");
467        assert_eq!(t.counters.dropped_slow.load(Ordering::Relaxed), 1);
468        assert!(t.dialled.is_empty(), "the slow session was dropped");
469    }
470}
clawft_kernel/mesh_artifact_tunnel.rs

clawft_kernel/
mesh_artifact_tunnel.rs
