clawft_kernel/infer_proxy/
mesh_forward.rs

1//! Forwarding over the mesh: the consumer side ([`forward_remote`]) and the
2//! serving side ([`serve_infer`]).
3//!
4//! Security shape:
5//! - The consumer dials only a node id the table resolved, and only through
6//!   a [`MeshDialer`] that vouches for the peer being admitted and verified.
7//!   The client's request never names a host.
8//! - The server serves only a verified peer, only a role this node exposed
9//!   to the mesh, and only a *local* instance: it never forwards onward.
10//! - Both sides bound frame counts by bytes and time, and the server
11//!   revalidates the request as it would a client's.
12
13use std::collections::HashMap;
14use std::sync::{Arc, Mutex};
15
16use async_trait::async_trait;
17use tokio::sync::{OwnedSemaphorePermit, Semaphore};
18
19use super::mesh_policy::{mesh_path_allowed, pin_body};
20use super::table::{MeshLocal, PlacementTable};
21use super::types::qualifies;
22use super::types::{MeshDialer, ProxyAudit, ProxyError, ProxyLimits, ProxyRequest, ResponseSink};
23use super::upstream::Upstream;
24use super::wire::{self, Resp};
25use crate::mesh::MeshStream;
26use crate::mesh_admit::Grant;
27use crate::mesh_framing::{FrameType, MeshFrame, read_frame as rf, write_frame as wf};
28
29fn mesh(e: impl std::fmt::Display) -> ProxyError {
30    ProxyError::Mesh(e.to_string())
31}
32
33async fn write_frame(s: &mut dyn MeshStream, f: &MeshFrame) -> Result<(), ProxyError> {
34    wf(s, f).await.map_err(mesh)
35}
36
37async fn read_frame(s: &mut dyn MeshStream) -> Result<MeshFrame, ProxyError> {
38    rf(s).await.map_err(mesh)
39}
40
41async fn recv_timeout(
42    s: &mut dyn MeshStream,
43    limits: &ProxyLimits,
44) -> Result<MeshFrame, ProxyError> {
45    tokio::time::timeout(limits.stall_timeout, read_frame(s))
46        .await
47        .map_err(|_| ProxyError::Timeout("mesh peer stalled".into()))?
48}
49
50/// Forward `req` to the instance of its role on `node_id` and relay the
51/// response into `sink`.
52pub async fn forward_remote(
53    dialer: &dyn MeshDialer,
54    node_id: &str,
55    req: &ProxyRequest,
56    sink: &mut dyn ResponseSink,
57    limits: &ProxyLimits,
58) -> Result<(), ProxyError> {
59    if !dialer.is_admitted(node_id) {
60        return Err(ProxyError::Refused(format!(
61            "peer {node_id} is not an admitted peer"
62        )));
63    }
64    if req.body.len() > limits.max_mesh_request_body {
65        return Err(ProxyError::TooLarge("request body over the mesh limit".into()));
66    }
67    let mut stream = dialer.dial(node_id).await?;
68    let result = async {
69        let payload = wire::encode_request(req)?;
70        write_frame(
71            stream.as_mut(),
72            &MeshFrame {
73                frame_type: FrameType::InferRequest,
74                payload,
75            },
76        )
77        .await?;
78        let mut total: u64 = 0;
79        let mut got_head = false;
80        loop {
81            let f = recv_timeout(stream.as_mut(), limits).await?;
82            if f.frame_type != FrameType::InferResponse {
83                return Err(mesh("unexpected frame from peer"));
84            }
85            match wire::decode_resp(&f.payload)? {
86                Resp::Head {
87                    status,
88                    content_type,
89                } => {
90                    if got_head {
91                        return Err(mesh("second response head"));
92                    }
93                    got_head = true;
94                    sink.head(status, content_type.as_deref()).await?;
95                }
96                Resp::Chunk(d) => {
97                    if !got_head {
98                        return Err(mesh("body before head"));
99                    }
100                    total += d.len() as u64;
101                    if total > limits.max_response_body {
102                        return Err(ProxyError::TooLarge("response body".into()));
103                    }
104                    sink.chunk(&d).await?;
105                }
106                Resp::End if got_head => return Ok(()),
107                Resp::End => return Err(mesh("end before head")),
108                Resp::Error(m) => return Err(ProxyError::Upstream(format!("peer: {m}"))),
109            }
110        }
111    }
112    .await;
113    let _ = stream.close().await;
114    result
115}
116
117/// The peer on the other end of a stream, as established by admission.
118#[derive(Debug, Clone)]
119pub struct InferPeer {
120    /// Node id.
121    pub node_id: String,
122    /// The peer's admission grant, from the connection (`None`: the id
123    /// was not verified against it). Serving requires [`qualifies`].
124    pub grant: Option<Grant>,
125}
126
127/// What a served exchange did, for the caller's accounting.
128#[derive(Debug, Clone, PartialEq, Eq)]
129pub enum Served {
130    /// Request answered.
131    Ok,
132    /// Request refused or failed; the reason was sent to the peer.
133    Failed(String),
134}
135
136struct FrameSink<'a> {
137    stream: &'a mut dyn MeshStream,
138}
139
140impl FrameSink<'_> {
141    async fn send(&mut self, r: Resp) -> Result<(), ProxyError> {
142        write_frame(
143            self.stream,
144            &MeshFrame {
145                frame_type: FrameType::InferResponse,
146                payload: wire::encode_resp(&r),
147            },
148        )
149        .await
150    }
151}
152
153#[async_trait]
154impl ResponseSink for FrameSink<'_> {
155    async fn head(&mut self, status: u16, content_type: Option<&str>) -> Result<(), ProxyError> {
156        self.send(Resp::Head {
157            status,
158            content_type: content_type.map(str::to_string),
159        })
160        .await
161    }
162
163    async fn chunk(&mut self, data: &[u8]) -> Result<(), ProxyError> {
164        for piece in data.chunks(wire::CHUNK_BYTES) {
165            self.send(Resp::Chunk(piece.to_vec())).await?;
166        }
167        Ok(())
168    }
169}
170
171/// Concurrency limits on the serving side: a peer cannot occupy the model
172/// server with parallel requests, and the node as a whole stays bounded.
173pub struct ServeGate {
174    total: Arc<Semaphore>,
175    per_peer_max: usize,
176    per_peer: Mutex<HashMap<String, Arc<Semaphore>>>,
177}
178
179impl ServeGate {
180    /// At most `per_peer` concurrent requests from one peer and `total` in
181    /// all.
182    pub fn new(per_peer: usize, total: usize) -> Self {
183        Self {
184            total: Arc::new(Semaphore::new(total)),
185            per_peer_max: per_peer,
186            per_peer: Mutex::new(HashMap::new()),
187        }
188    }
189
190    fn acquire(&self, peer: &str) -> Option<(OwnedSemaphorePermit, OwnedSemaphorePermit)> {
191        let peer_sem = {
192            let mut m = self.per_peer.lock().unwrap();
193            // Entries with no request in flight are dropped so the table
194            // cannot grow with the peers seen.
195            m.retain(|_, s| s.available_permits() < self.per_peer_max);
196            m.entry(peer.to_string())
197                .or_insert_with(|| Arc::new(Semaphore::new(self.per_peer_max)))
198                .clone()
199        };
200        let a = peer_sem.try_acquire_owned().ok()?;
201        let b = self.total.clone().try_acquire_owned().ok()?;
202        Some((a, b))
203    }
204}
205
206#[cfg(test)]
207impl ServeGate {
208    /// Take a slot and leak it (the test never releases it).
209    pub(super) fn acquire_for_test(&self, peer: &str) -> bool {
210        match self.acquire(peer) {
211            Some((a, b)) => {
212                std::mem::forget(a);
213                std::mem::forget(b);
214                true
215            }
216            None => false,
217        }
218    }
219}
220
221#[cfg(test)]
222impl ServeGate {
223    /// Take and immediately release a slot: whether one was free.
224    pub(super) fn acquire_for_test_and_release(&self, peer: &str) -> bool {
225        self.acquire(peer).is_some()
226    }
227}
228
229impl Default for ServeGate {
230    fn default() -> Self {
231        Self::new(2, 8)
232    }
233}
234
235/// What a peer is told. Local detail (ports, paths, server messages) goes to
236/// the audit trail, never over the wire.
237fn peer_reason(e: &ProxyError) -> &'static str {
238    match e {
239        ProxyError::Refused(_) => "refused",
240        ProxyError::NoInstance(_) => "role not served here",
241        ProxyError::BadRequest(_)
242        | ProxyError::Forbidden(_)
243        | ProxyError::MethodNotAllowed
244        | ProxyError::TooLarge(_) => "request not accepted",
245        ProxyError::Timeout(_) => "timeout",
246        _ => "upstream error",
247    }
248}
249
250/// What the serving side consults about a peer. The standing checks that
251/// need no frame (`peer_listed`) run before a slot is taken or a byte is read.
252pub trait ServePolicy: Send + Sync {
253    /// The loopback instance this node exposes to `peer` for `role`, or
254    /// `None` (not exposed, or `peer` is not on the role's serve allowlist).
255    fn instance_for(&self, role: &str, peer: &str) -> Option<MeshLocal>;
256    /// Whether `peer` is on any role's serve allowlist.
257    fn peer_listed(&self, peer: &str) -> bool;
258}
259
260impl ServePolicy for PlacementTable {
261    fn instance_for(&self, role: &str, peer: &str) -> Option<MeshLocal> {
262        self.local_for_peer(role, peer)
263    }
264    fn peer_listed(&self, peer: &str) -> bool {
265        self.peer_listed_any(peer)
266    }
267}
268
269/// A bare closure is a policy that lists every peer (unit tests).
270impl<F: Fn(&str, &str) -> Option<MeshLocal> + Send + Sync> ServePolicy for F {
271    fn instance_for(&self, role: &str, peer: &str) -> Option<MeshLocal> {
272        self(role, peer)
273    }
274    fn peer_listed(&self, _: &str) -> bool {
275        true
276    }
277}
278
279/// Serve one forwarded request from `peer` on `stream`. The peer's grant
280/// must qualify and the peer must be on a serve allowlist before a slot is
281/// taken or a frame read; then the first frame must arrive within the head
282/// timeout. The request is held to the mesh path allowlist, resolved to the
283/// instance the policy exposes to this peer for the role, and its body is
284/// pinned to that instance's model.
285pub async fn serve_infer(
286    stream: &mut dyn MeshStream,
287    peer: &InferPeer,
288    policy: &dyn ServePolicy,
289    upstream: &Upstream,
290    gate: &ServeGate,
291    audit: Option<&dyn ProxyAudit>,
292) -> Result<Served, ProxyError> {
293    let limits = upstream.limits().clone();
294    let refuse = |e: &ProxyError| {
295        if let Some(a) = audit {
296            a.record(
297                "infer.mesh.failed",
298                serde_json::json!({"peer": peer.node_id, "standing": peer.grant.as_ref().map(qualifies), "why": e.to_string()}),
299            );
300        }
301        peer_reason(e).to_string()
302    };
303    let mut sink = FrameSink { stream };
304
305    // Standing first: a peer that does not qualify holds no slot and costs
306    // no read.
307    let standing = if !peer.grant.as_ref().is_some_and(qualifies) {
308        Some(ProxyError::Refused("peer is not an enforced-admission full node".into()))
309    } else if !policy.peer_listed(&peer.node_id) {
310        Some(ProxyError::Refused("peer is on no serve allowlist".into()))
311    } else {
312        None
313    };
314    if let Some(e) = standing {
315        let why = refuse(&e);
316        let _ = sink.send(Resp::Error(why.clone())).await;
317        return Ok(Served::Failed(why));
318    }
319    // Then the slots, before reading anything: a peer over its limit must
320    // not make this node buffer a frame per stream.
321    let Some(_permits) = gate.acquire(&peer.node_id) else {
322        let e = ProxyError::Refused("too many concurrent requests".into());
323        let why = refuse(&e);
324        let _ = sink.send(Resp::Error(why.clone())).await;
325        return Ok(Served::Failed(why));
326    };
327    // The first frame gets the short head timeout, not the stall timeout.
328    let first = tokio::time::timeout(limits.head_timeout, read_frame(sink.stream)).await;
329    let frame = match first {
330        Ok(Ok(f)) => f,
331        Ok(Err(e)) => return Ok(Served::Failed(refuse(&e))),
332        Err(_) => {
333            return Ok(Served::Failed(refuse(&ProxyError::Timeout("first frame".into()))));
334        }
335    };
336    let outcome: Result<(), ProxyError> = async {
337        if frame.frame_type != FrameType::InferRequest {
338            return Err(ProxyError::BadRequest("expected an infer request".into()));
339        }
340        let mut req = wire::decode_request(
341            &frame.payload,
342            &ProxyLimits {
343                max_request_body: limits.max_mesh_request_body,
344                ..limits.clone()
345            },
346        )?;
347        if !mesh_path_allowed(&req.path) {
348            return Err(ProxyError::Forbidden("path is not served to peers".into()));
349        }
350        let local = policy.instance_for(&req.role, &peer.node_id).ok_or_else(|| {
351            ProxyError::NoInstance(format!("{} (not served to this peer here)", req.role))
352        })?;
353        req.body = pin_body(&req.path, &req.body, &local)?;
354        let fwd = upstream.forward(&local.base, &req, false, &mut sink);
355        tokio::time::timeout(limits.request_timeout, fwd)
356            .await
357            .map_err(|_| ProxyError::Timeout("request".into()))??;
358        sink.send(Resp::End).await
359    }
360    .await;
361    match outcome {
362        Ok(()) => Ok(Served::Ok),
363        Err(e) => {
364            let why = refuse(&e);
365            let _ = sink.send(Resp::Error(why.clone())).await;
366            Ok(Served::Failed(why))
367        }
368    }
369}
clawft_kernel/infer_proxy/mesh_forward.rs

clawft_kernel/infer_proxy/
mesh_forward.rs
