clawft_kernel/cog_ingest/
bridge.rs

1//! The node-local HTTP bridge: `POST /api/v1/store/ingest`.
2//!
3//! A deliberately small HTTP/1.1 server: one request per connection,
4//! `Content-Length` bodies only, strict limits, no keep-alive (the Seed
5//! ingest endpoint stalls keep-alive connections, so cogs already cope).
6//!
7//! Order of checks per request: request line, headers (size-bounded),
8//! token (401 unknown, 403 another instance's), per-instance request rate
9//! (429), body length (411/413), body read, shape validation (400/413),
10//! per-instance vector budget (429), then the store owner. Nothing is read
11//! from the body of an unauthenticated request.
12
13use std::net::SocketAddr;
14use std::sync::Arc;
15use std::sync::atomic::{AtomicU64, Ordering};
16use std::time::Duration;
17
18use tokio::io::{AsyncReadExt, AsyncWriteExt};
19use tokio::net::{TcpListener, TcpStream};
20use tokio::task::JoinHandle;
21
22use super::registry::{InstanceBinding, RateBudget, StoreRouter, TokenRegistry};
23use super::types::{INGEST_PATH, IngestError, MAX_BODY_BYTES, parse_batch};
24
25/// Largest request head (request line and headers).
26pub const MAX_HEAD_BYTES: usize = 8 * 1024;
27/// Most concurrent connections one bridge listener serves.
28pub const MAX_CONNECTIONS: usize = 128;
29
30/// Which instances a listener accepts.
31#[derive(Debug, Clone, PartialEq, Eq)]
32pub enum BridgeScope {
33    /// Only this instance's token (one listener per instance; the address
34    /// goes into [`HostContract::with_ingest_upstream`](crate::workload_runtime::HostContract)).
35    Instance(String),
36    /// Only the holder of one token, identified by its BLAKE3 hash. For a
37    /// listener bound before the instance id exists (the token is issued
38    /// with the host contract, before the adapter names the instance).
39    Token([u8; 32]),
40    /// Any registered instance (the node's shared `127.0.0.1:80`).
41    Any,
42}
43
44/// Bridge settings.
45#[derive(Debug, Clone)]
46pub struct BridgeConfig {
47    /// Time allowed to receive a request head and body.
48    pub read_timeout: Duration,
49    /// Time allowed from accept to an authenticated request head (the
50    /// pre-auth deadline; shorter than `read_timeout`).
51    pub preauth_timeout: Duration,
52    /// Most connections that have not authenticated yet, at once. A further
53    /// connection waits up to `preauth_timeout` for a slot, then is dropped,
54    /// so an anonymous flood cannot use up the connection cap meant for
55    /// authenticated cogs. Residual limit: a process on this host (the bridge
56    /// is loopback) that keeps this many connections open and renews them
57    /// every `preauth_timeout` still delays every cog by up to that long and
58    /// can starve it; the token cannot be checked before the head arrives.
59    pub max_unauthenticated: usize,
60    /// Allow binding a non-loopback address. Only for an instance-scoped
61    /// listener whose address a container relay needs (the VM gateway);
62    /// the token still applies.
63    pub allow_non_loopback: bool,
64}
65
66impl Default for BridgeConfig {
67    fn default() -> Self {
68        Self {
69            read_timeout: Duration::from_secs(5),
70            preauth_timeout: Duration::from_secs(2),
71            max_unauthenticated: 64,
72            allow_non_loopback: false,
73        }
74    }
75}
76
77/// Counters, for evidence and tests.
78#[derive(Debug, Default)]
79pub struct BridgeStats {
80    /// Batches forwarded and written.
81    pub accepted: AtomicU64,
82    /// Vectors written.
83    pub vectors: AtomicU64,
84    /// Requests refused before the owner (auth, shape, size, rate).
85    pub rejected: AtomicU64,
86    /// Requests the owner could not take.
87    pub failed: AtomicU64,
88}
89
90/// The bridge.
91pub struct IngestBridge {
92    registry: Arc<TokenRegistry>,
93    router: Arc<dyn StoreRouter>,
94    budget: RateBudget,
95    cfg: BridgeConfig,
96    stats: BridgeStats,
97}
98
99/// A running listener.
100pub struct BridgeHandle {
101    addr: SocketAddr,
102    task: JoinHandle<()>,
103}
104
105impl BridgeHandle {
106    /// Address bound.
107    pub fn addr(&self) -> SocketAddr {
108        self.addr
109    }
110}
111
112impl Drop for BridgeHandle {
113    fn drop(&mut self) {
114        self.task.abort();
115    }
116}
117
118impl IngestBridge {
119    /// Bridge over `registry` and `router`.
120    pub fn new(
121        registry: Arc<TokenRegistry>,
122        router: Arc<dyn StoreRouter>,
123        budget: RateBudget,
124        cfg: BridgeConfig,
125    ) -> Arc<Self> {
126        Arc::new(Self {
127            registry,
128            router,
129            budget,
130            cfg,
131            stats: BridgeStats::default(),
132        })
133    }
134
135    /// True if a store owner is known for `b`.
136    pub fn has_route(&self, b: &InstanceBinding) -> bool {
137        self.router.route(b).is_some()
138    }
139
140    /// Drop an instance's rate counters (at unload).
141    pub fn forget(&self, instance_id: &str) {
142        self.budget.forget(instance_id);
143    }
144
145    /// Counters.
146    pub fn stats(&self) -> &BridgeStats {
147        &self.stats
148    }
149
150    /// Bind `addr` and serve `scope`. A non-loopback address is refused
151    /// unless the config allows it and the scope is one instance.
152    pub async fn bind(
153        self: &Arc<Self>,
154        addr: SocketAddr,
155        scope: BridgeScope,
156    ) -> Result<BridgeHandle, IngestError> {
157        let one_instance = matches!(scope, BridgeScope::Instance(_) | BridgeScope::Token(_));
158        if !(addr.ip().is_loopback() || self.cfg.allow_non_loopback && one_instance) {
159            return Err(IngestError::Malformed(format!(
160                "refusing to bind {addr}: the bridge listens on loopback, or on one \
161                 instance-scoped address when allow_non_loopback is set"
162            )));
163        }
164        let listener = TcpListener::bind(addr)
165            .await
166            .map_err(|e| IngestError::Unavailable(format!("bind {addr}: {e}")))?;
167        let bound = listener
168            .local_addr()
169            .map_err(|e| IngestError::Unavailable(e.to_string()))?;
170        let me = self.clone();
171        let slots = Arc::new(tokio::sync::Semaphore::new(MAX_CONNECTIONS));
172        let unauth = Arc::new(tokio::sync::Semaphore::new(self.cfg.max_unauthenticated));
173        let task = tokio::spawn(async move {
174            loop {
175                let Ok((sock, _)) = listener.accept().await else {
176                    return;
177                };
178                let Ok(slot) = slots.clone().try_acquire_owned() else {
179                    continue; // over the connection cap: dropped
180                };
181                let (me, scope, unauth) = (me.clone(), scope.clone(), unauth.clone());
182                tokio::spawn(async move {
183                    let _slot = slot;
184                    // Connections that have not authenticated are capped
185                    // separately. A new one waits (first come, first served)
186                    // up to the pre-auth deadline for a slot, so holders that
187                    // never authenticate free theirs before a legitimate cog
188                    // gives up; beyond the wait it is dropped.
189                    let Ok(Ok(pre)) =
190                        tokio::time::timeout(me.cfg.preauth_timeout, unauth.acquire_owned()).await
191                    else {
192                        return;
193                    };
194                    me.serve(sock, &scope, pre).await;
195                });
196            }
197        });
198        Ok(BridgeHandle { addr: bound, task })
199    }
200
201    async fn serve(
202        &self,
203        mut sock: TcpStream,
204        scope: &BridgeScope,
205        pre: tokio::sync::OwnedSemaphorePermit,
206    ) {
207        let res = tokio::time::timeout(
208            self.cfg.read_timeout,
209            self.handle(&mut sock, scope, pre),
210        )
211        .await;
212        let (status, body, extra) = match res {
213            Ok(Ok(v)) => (200, v, None),
214            Ok(Err(e)) => {
215                match &e {
216                    IngestError::Unavailable(_) => self.stats.failed.fetch_add(1, Ordering::Relaxed),
217                    _ => self.stats.rejected.fetch_add(1, Ordering::Relaxed),
218                };
219                let extra = matches!(e, IngestError::RateLimited).then_some("Retry-After: 1\r\n");
220                // Owner-side detail stays in the log, not on the wire.
221                let msg = match &e {
222                    IngestError::Unavailable(d) => {
223                        tracing::warn!(detail = %d, "cog ingest forward failed");
224                        "store unavailable".to_string()
225                    }
226                    other => other.to_string(),
227                };
228                (e.status(), serde_json::json!({"ok": false, "error": msg}), extra)
229            }
230            Err(_) => {
231                self.stats.rejected.fetch_add(1, Ordering::Relaxed);
232                (408, serde_json::json!({"ok": false, "error": "timeout"}), None)
233            }
234        };
235        let body = body.to_string();
236        let head = format!(
237            "HTTP/1.1 {status} {}\r\nContent-Type: application/json\r\nContent-Length: {}\r\n{}Connection: close\r\n\r\n",
238            reason(status),
239            body.len(),
240            extra.unwrap_or("")
241        );
242        let _ = sock.write_all(head.as_bytes()).await;
243        let _ = sock.write_all(body.as_bytes()).await;
244        let _ = sock.shutdown().await;
245    }
246
247    async fn handle(
248        &self,
249        sock: &mut TcpStream,
250        scope: &BridgeScope,
251        pre: tokio::sync::OwnedSemaphorePermit,
252    ) -> Result<serde_json::Value, IngestError> {
253        // Until the token checks out the connection holds an
254        // unauthenticated slot and a short deadline.
255        let (req, mut body, binding) = tokio::time::timeout(self.cfg.preauth_timeout, async {
256            let (head, body) = read_head(sock).await?;
257            let req = parse_head(&head)?;
258            if req.path != INGEST_PATH {
259                return Err(IngestError::Malformed(format!("no such path `{}`", req.path)));
260            }
261            if req.method != "POST" {
262                return Err(IngestError::Malformed("POST only".into()));
263            }
264            let binding = self.authenticate(req.token.as_deref(), scope)?;
265            Ok((req, body, binding))
266        })
267        .await
268        .map_err(|_| IngestError::Malformed("request head not received in time".into()))??;
269        drop(pre);
270        if !self.budget.charge_request(&binding.instance_id) {
271            return Err(IngestError::RateLimited);
272        }
273        let len = req
274            .content_length
275            .ok_or_else(|| IngestError::Malformed("Content-Length required".into()))?;
276        if len > MAX_BODY_BYTES {
277            return Err(IngestError::TooLarge(format!(
278                "body over {MAX_BODY_BYTES} bytes"
279            )));
280        }
281        if body.len() > len {
282            body.truncate(len); // pipelined bytes are not a second request
283        }
284        while body.len() < len {
285            let mut chunk = vec![0u8; (len - body.len()).min(8192)];
286            let n = sock
287                .read(&mut chunk)
288                .await
289                .map_err(|e| IngestError::Malformed(format!("read: {e}")))?;
290            if n == 0 {
291                return Err(IngestError::Malformed("body shorter than Content-Length".into()));
292            }
293            body.extend_from_slice(&chunk[..n]);
294        }
295        let batch = parse_batch(&body)?;
296        if !self.budget.charge_vectors(&binding.instance_id, batch.vectors.len()) {
297            return Err(IngestError::RateLimited);
298        }
299        let fwd = self
300            .router
301            .route(&binding)
302            .ok_or_else(|| IngestError::Unavailable("no store owner for this placement".into()))?;
303        // The token is checked again immediately before the write, so a stop
304        // or unload that completed while this request was being read is not
305        // followed by a write.
306        if !self.registry.contains(&binding.instance_id) {
307            return Err(IngestError::Unauthorized);
308        }
309        let out = fwd.forward(&binding, &batch).await?;
310        self.stats.accepted.fetch_add(1, Ordering::Relaxed);
311        self.stats
312            .vectors
313            .fetch_add(out.accepted as u64, Ordering::Relaxed);
314        Ok(serde_json::json!({
315            "ok": true,
316            "accepted": out.accepted,
317            "deduped": out.deduped,
318        }))
319    }
320
321    fn authenticate(
322        &self,
323        token: Option<&str>,
324        scope: &BridgeScope,
325    ) -> Result<InstanceBinding, IngestError> {
326        let b = token
327            .and_then(|t| self.registry.lookup(t))
328            .ok_or(IngestError::Unauthorized)?;
329        match scope {
330            BridgeScope::Instance(id) if *id != b.instance_id => Err(IngestError::Forbidden),
331            BridgeScope::Token(h)
332                if token.map(|t| blake3::hash(t.as_bytes())) != Some(blake3::Hash::from(*h)) =>
333            {
334                Err(IngestError::Forbidden)
335            }
336            _ => Ok(b),
337        }
338    }
339}
340
341fn reason(status: u16) -> &'static str {
342    match status {
343        200 => "OK",
344        400 => "Bad Request",
345        401 => "Unauthorized",
346        403 => "Forbidden",
347        408 => "Request Timeout",
348        411 => "Length Required",
349        413 => "Payload Too Large",
350        429 => "Too Many Requests",
351        502 => "Bad Gateway",
352        _ => "Error",
353    }
354}
355
356struct Head {
357    method: String,
358    path: String,
359    content_length: Option<usize>,
360    token: Option<String>,
361}
362
363/// Read until the blank line; returns the head text and any body bytes
364/// already read.
365async fn read_head(sock: &mut TcpStream) -> Result<(String, Vec<u8>), IngestError> {
366    let mut buf = Vec::with_capacity(1024);
367    loop {
368        if let Some(i) = buf.windows(4).position(|w| w == b"\r\n\r\n") {
369            let rest = buf.split_off(i + 4);
370            buf.truncate(i);
371            let head = String::from_utf8(buf)
372                .map_err(|_| IngestError::Malformed("head is not UTF-8".into()))?;
373            return Ok((head, rest));
374        }
375        if buf.len() > MAX_HEAD_BYTES {
376            return Err(IngestError::TooLarge("request head".into()));
377        }
378        let mut chunk = [0u8; 1024];
379        let n = sock
380            .read(&mut chunk)
381            .await
382            .map_err(|e| IngestError::Malformed(format!("read: {e}")))?;
383        if n == 0 {
384            return Err(IngestError::Malformed("connection closed in head".into()));
385        }
386        buf.extend_from_slice(&chunk[..n]);
387    }
388}
389
390fn parse_head(head: &str) -> Result<Head, IngestError> {
391    let bad = |m: &str| IngestError::Malformed(m.to_string());
392    let mut lines = head.split("\r\n");
393    let line = lines.next().ok_or_else(|| bad("empty request"))?;
394    let mut parts = line.split(' ');
395    let (method, target, version) = (parts.next(), parts.next(), parts.next());
396    let (Some(method), Some(target), Some(version), None) = (method, target, version, parts.next())
397    else {
398        return Err(bad("bad request line"));
399    };
400    if version != "HTTP/1.1" && version != "HTTP/1.0" {
401        return Err(bad("unsupported HTTP version"));
402    }
403    let path = target.split('?').next().unwrap_or("").to_string();
404    let (mut content_length, mut token) = (None, None);
405    for l in lines {
406        let (k, v) = l.split_once(':').ok_or_else(|| bad("bad header line"))?;
407        let (k, v) = (k.trim().to_ascii_lowercase(), v.trim());
408        match k.as_str() {
409            "content-length" => {
410                if content_length.is_some() {
411                    return Err(bad("duplicate Content-Length"));
412                }
413                content_length = Some(v.parse::<usize>().map_err(|_| bad("bad Content-Length"))?);
414            }
415            "transfer-encoding" => return Err(bad("Transfer-Encoding is not supported")),
416            "authorization" => {
417                let t = v.strip_prefix("Bearer ").or_else(|| v.strip_prefix("bearer "));
418                token = token.or(t.map(|t| t.trim().to_string()));
419            }
420            "x-api-key" | "x-cog-token" => token = token.or(Some(v.to_string())),
421            _ => {}
422        }
423    }
424    Ok(Head {
425        method: method.to_string(),
426        path,
427        content_length,
428        token,
429    })
430}
clawft_kernel/cog_ingest/bridge.rs

clawft_kernel/cog_ingest/
bridge.rs
