clawft_kernel/cog_ingest/
forward.rs

1//! Forwarding validated batches to the store owner (the owner's side is in
2//! [`super::owner`]).
3//!
4//! The owner is this node ([`LocalForwarder`]) or another node
5//! ([`MeshForwarder`] to a [`StoreOwnerService`]). Remote forwarding is a
6//! [`ForwardRequest`] signed by the bridge node's Ed25519 key, carried as a
7//! [`MeshIpcEnvelope`] addressed to `ServiceMethod { cog-store, store.ingest }`
8//! and answered by a response signed by the owner and bound to the request
9//! nonce. The owner checks, in order: size, signature, structure, key
10//! binding, addressee, time window, forwarder policy (key and project),
11//! nonce freshness. The nonce is recorded last, so a forged request cannot
12//! burn a legitimate one.
13
14use std::sync::Arc;
15use std::time::{Duration, SystemTime, UNIX_EPOCH};
16
17use async_trait::async_trait;
18use ed25519_dalek::SigningKey;
19use serde::{Deserialize, Serialize};
20
21use super::registry::InstanceBinding;
22use super::store::{IngestOutcome, Provenance, StoreDirectory, StoreError};
23use super::types::{IngestBatch, IngestError, IngestVector};
24use crate::ipc::{KernelMessage, MessagePayload, MessageTarget};
25use crate::mesh::MeshStream;
26use crate::mesh_ipc::{MeshIpcEnvelope, MeshRequest};
27use crate::node_registry::node_id_from_pubkey;
28use crate::workload_ctl::msg::{SignedCtl, fresh_nonce, open, sign};
29use crate::workload_ctl::transport::CtlConnector;
30
31/// Service the store owner answers as.
32pub const STORE_SERVICE: &str = "cog-store";
33/// Method of the service.
34pub const STORE_INGEST_METHOD: &str = "store.ingest";
35/// Wire version.
36pub const FORWARD_VERSION: u32 = 1;
37/// Domain tag signed before a request payload.
38pub const FORWARD_DOMAIN: &[u8] = b"weftos.cog_ingest.forward.v1\0";
39/// Domain tag signed before a response payload.
40pub const FORWARD_RESPONSE_DOMAIN: &[u8] = b"weftos.cog_ingest.forward_response.v1\0";
41/// Address scheme for in-process owners.
42pub const MEM_SCHEME: &str = "mem://";
43
44pub(super) fn now_ms() -> u64 {
45    SystemTime::now()
46        .duration_since(UNIX_EPOCH)
47        .map(|d| d.as_millis() as u64)
48        .unwrap_or(0)
49}
50
51/// Takes a batch to the node that owns the store.
52#[async_trait]
53pub trait Forwarder: Send + Sync {
54    /// Forward `batch` for the instance `b`.
55    async fn forward(
56        &self,
57        b: &InstanceBinding,
58        batch: &IngestBatch,
59    ) -> Result<IngestOutcome, IngestError>;
60}
61
62fn store_err(e: StoreError) -> IngestError {
63    IngestError::Unavailable(e.to_string())
64}
65
66/// The store owner is this node.
67pub struct LocalForwarder {
68    node_id: String,
69    directory: Arc<dyn StoreDirectory>,
70}
71
72impl LocalForwarder {
73    /// Forwarder into `directory` on node `node_id`.
74    pub fn new(node_id: impl Into<String>, directory: Arc<dyn StoreDirectory>) -> Self {
75        Self {
76            node_id: node_id.into(),
77            directory,
78        }
79    }
80}
81
82#[async_trait]
83impl Forwarder for LocalForwarder {
84    async fn forward(
85        &self,
86        b: &InstanceBinding,
87        batch: &IngestBatch,
88    ) -> Result<IngestOutcome, IngestError> {
89        let store = self
90            .directory
91            .store_for(b.project_id.as_deref())
92            .ok_or_else(|| IngestError::Unavailable("no store for this placement".into()))?;
93        let from = Provenance {
94            instance_id: b.instance_id.clone(),
95            source_node: self.node_id.clone(),
96        };
97        store
98            .ingest(&from, &batch.vectors, batch.dedup)
99            .map_err(store_err)
100    }
101}
102
103// ── Wire messages ────────────────────────────────────────────────────────
104
105/// A forwarded batch.
106#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
107#[serde(deny_unknown_fields)]
108pub struct ForwardRequest {
109    /// [`FORWARD_VERSION`].
110    pub version: u32,
111    /// Forwarding node id (derived from the signing key).
112    pub requester: String,
113    /// Store-owner node id.
114    pub target: String,
115    /// 32 hex characters, fresh per request.
116    pub nonce: String,
117    /// Issue time, ms since the Unix epoch.
118    pub issued_at_ms: u64,
119    /// Expiry, ms since the Unix epoch.
120    pub expires_at_ms: u64,
121    /// Project whose store takes the batch (`None`: the controller's).
122    #[serde(default, skip_serializing_if = "Option::is_none")]
123    pub project_id: Option<String>,
124    /// Instance that posted it (provenance).
125    pub instance_id: String,
126    /// Placing controller (provenance).
127    pub controller_node: String,
128    /// Dedup flag, honoured by the store.
129    pub dedup: bool,
130    /// The vectors.
131    pub vectors: Vec<IngestVector>,
132}
133
134/// Why an owner refused. Stable names.
135#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
136#[serde(rename_all = "snake_case")]
137pub enum ForwardRefusal {
138    /// Signature, key or structure check failed.
139    Signature,
140    /// Outside its time window.
141    Expired,
142    /// Nonce already seen.
143    Replay,
144    /// Addressed to another node.
145    NotForMe,
146    /// Key (or key and project) not authorised to forward here.
147    Unauthorized,
148    /// Malformed batch.
149    Invalid,
150    /// No store for the project here.
151    NoStore,
152    /// Store at capacity.
153    StoreFull,
154    /// Store backend failure.
155    Backend,
156}
157
158/// What the owner did.
159#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
160#[serde(tag = "outcome", rename_all = "snake_case")]
161pub enum ForwardOutcome {
162    /// Written.
163    Ok {
164        /// Counts.
165        result: IngestOutcome,
166    },
167    /// Refused.
168    Refused {
169        /// Code.
170        code: ForwardRefusal,
171        /// Reason.
172        reason: String,
173    },
174}
175
176/// Signed answer.
177#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
178#[serde(deny_unknown_fields)]
179pub struct ForwardResponse {
180    /// [`FORWARD_VERSION`].
181    pub version: u32,
182    /// Responding node id.
183    pub responder: String,
184    /// Nonce of the request answered.
185    pub request_nonce: String,
186    /// Result.
187    pub outcome: ForwardOutcome,
188}
189
190impl ForwardRequest {
191    /// Build and sign a request for `batch` from `key`'s node.
192    pub fn signed(
193        key: &SigningKey,
194        target: &str,
195        b: &InstanceBinding,
196        batch: &IngestBatch,
197        ttl_ms: u64,
198    ) -> Result<(Self, SignedCtl), IngestError> {
199        let now = now_ms();
200        let req = Self {
201            version: FORWARD_VERSION,
202            requester: node_id_from_pubkey(&key.verifying_key().to_bytes()),
203            target: target.to_string(),
204            nonce: fresh_nonce(),
205            issued_at_ms: now,
206            expires_at_ms: now.saturating_add(ttl_ms),
207            project_id: b.project_id.clone(),
208            instance_id: b.instance_id.clone(),
209            controller_node: b.controller_node.clone(),
210            dedup: batch.dedup,
211            vectors: batch.vectors.clone(),
212        };
213        let payload = serde_json::to_string(&req)
214            .map_err(|e| IngestError::Unavailable(format!("encode: {e}")))?;
215        let signed = sign(FORWARD_DOMAIN, payload, key);
216        Ok((req, signed))
217    }
218}
219
220pub(super) fn owner_target() -> MessageTarget {
221    MessageTarget::ServiceMethod {
222        service: STORE_SERVICE.to_string(),
223        method: STORE_INGEST_METHOD.to_string(),
224    }
225}
226
227// ── Bridge side ──────────────────────────────────────────────────────────
228
229/// Forwards to a remote owner over a kept connection (reconnects once on a
230/// transport failure).
231pub struct MeshForwarder {
232    key: SigningKey,
233    local_node: String,
234    owner_node: String,
235    owner_key: [u8; 32],
236    addr: String,
237    connector: Arc<dyn CtlConnector>,
238    timeout: Duration,
239    conn: tokio::sync::Mutex<Option<Box<dyn MeshStream>>>,
240}
241
242impl MeshForwarder {
243    /// Forwarder signing as `key`'s node to the owner `owner_node` (whose
244    /// key must be `owner_key`) at `addr`.
245    pub fn new(
246        key: SigningKey,
247        owner_node: &str,
248        owner_key: [u8; 32],
249        addr: &str,
250        connector: Arc<dyn CtlConnector>,
251    ) -> Self {
252        Self {
253            local_node: node_id_from_pubkey(&key.verifying_key().to_bytes()),
254            key,
255            owner_node: owner_node.to_string(),
256            owner_key,
257            addr: addr.to_string(),
258            connector,
259            timeout: Duration::from_secs(5),
260            conn: tokio::sync::Mutex::new(None),
261        }
262    }
263
264    async fn round_trip(
265        &self,
266        guard: &mut Option<Box<dyn MeshStream>>,
267        req: &MeshRequest,
268        bytes: &[u8],
269    ) -> Result<MeshIpcEnvelope, String> {
270        if guard.is_none() {
271            *guard = Some(
272                self.connector
273                    .connect(&self.addr)
274                    .await
275                    .map_err(|e| e.to_string())?,
276            );
277        }
278        let s = guard.as_mut().ok_or("no connection")?;
279        s.send(bytes).await.map_err(|e| e.to_string())?;
280        let raw = tokio::time::timeout(self.timeout, s.recv())
281            .await
282            .map_err(|_| "timed out".to_string())?
283            .map_err(|e| e.to_string())?;
284        let env = MeshIpcEnvelope::from_bytes(&raw).map_err(|e| e.to_string())?;
285        if !req.matches_response(&env) {
286            return Err("uncorrelated response".into());
287        }
288        Ok(env)
289    }
290}
291
292#[async_trait]
293impl Forwarder for MeshForwarder {
294    async fn forward(
295        &self,
296        b: &InstanceBinding,
297        batch: &IngestBatch,
298    ) -> Result<IngestOutcome, IngestError> {
299        let unavailable = |m: String| IngestError::Unavailable(m);
300        let (sent, signed) = ForwardRequest::signed(&self.key, &self.owner_node, b, batch, 30_000)?;
301        let msg = KernelMessage::new(
302            0,
303            owner_target(),
304            MessagePayload::Json(serde_json::to_value(&signed).unwrap_or_default()),
305        );
306        let req = MeshRequest::new(
307            MeshIpcEnvelope::new(self.local_node.clone(), self.owner_node.clone(), msg),
308            self.timeout,
309        );
310        let bytes = req.request.to_bytes().map_err(|e| unavailable(e.to_string()))?;
311        let mut guard = self.conn.lock().await;
312        let env = match self.round_trip(&mut guard, &req, &bytes).await {
313            Ok(e) => e,
314            Err(first) => {
315                // A kept connection may have gone away. The same signed
316                // request is resent once; if the owner had already taken it
317                // the answer is a replay refusal, never a double write.
318                *guard = None;
319                self.round_trip(&mut guard, &req, &bytes)
320                    .await
321                    .map_err(|e| unavailable(format!("{first}; retry: {e}")))?
322            }
323        };
324        drop(guard);
325        let MessagePayload::Json(v) = env.message.payload else {
326            return Err(unavailable("response payload is not JSON".into()));
327        };
328        let signed: SignedCtl = serde_json::from_value(v).map_err(|e| unavailable(e.to_string()))?;
329        let pk = open(FORWARD_RESPONSE_DOMAIN, &signed).map_err(|r| unavailable(r.reason))?;
330        if pk != self.owner_key {
331            return Err(unavailable("response signed by another key".into()));
332        }
333        let resp: ForwardResponse =
334            serde_json::from_str(&signed.payload).map_err(|e| unavailable(e.to_string()))?;
335        if resp.responder != self.owner_node || resp.request_nonce != sent.nonce {
336            return Err(unavailable("response answers another request".into()));
337        }
338        match resp.outcome {
339            ForwardOutcome::Ok { result } => Ok(result),
340            ForwardOutcome::Refused { code, reason } => {
341                Err(unavailable(format!("owner refused ({code:?}): {reason}")))
342            }
343        }
344    }
345}
clawft_kernel/cog_ingest/forward.rs

clawft_kernel/cog_ingest/
forward.rs
