clawft_kernel/cog_ingest/
owner.rs

1//! The store owner's side of forwarding: the `cog-store` service that
2//! verifies signed [`ForwardRequest`]s and writes them to the owning store,
3//! its policy, and how it is served (in-process or mesh TCP).
4
5use std::collections::{HashMap, HashSet};
6use std::sync::Arc;
7
8use async_trait::async_trait;
9use ed25519_dalek::SigningKey;
10
11use super::forward::{
12    FORWARD_DOMAIN, FORWARD_RESPONSE_DOMAIN, FORWARD_VERSION, ForwardOutcome, ForwardRefusal,
13    ForwardRequest, ForwardResponse, MEM_SCHEME, STORE_INGEST_METHOD, STORE_SERVICE, now_ms,
14    owner_target,
15};
16use super::store::{Provenance, StoreDirectory, StoreError};
17use super::types::MAX_VECTORS_PER_BATCH;
18use crate::ipc::{KernelMessage, MessagePayload, MessageTarget};
19use crate::mesh::{MeshError, MeshStream, TransportListener};
20use crate::mesh_ipc::MeshIpcEnvelope;
21use crate::node_registry::node_id_from_pubkey;
22use crate::workload_ctl::msg::{
23    MAX_CTL_BYTES, MAX_SKEW_MS, MAX_TTL_MS, NonceGuard, SignedCtl, open, sign,
24};
25use crate::workload_ctl::session::NoiseStream;
26use crate::workload_ctl::transport::CtlConnector;
27use crate::workload_pkg::codec::is_lower_hex;
28
29
30/// Which bridge keys may forward, and for which projects.
31pub trait ForwardPolicy: Send + Sync {
32    /// True if a batch for `project_id` signed by `key` may be taken.
33    fn allows(&self, key: &[u8; 32], project_id: Option<&str>) -> bool;
34}
35
36/// Fixed key table. A key is allowed for every project, or for a listed set
37/// (a restricted key cannot forward project-less batches).
38#[derive(Default)]
39pub struct KeyPolicy {
40    keys: HashMap<[u8; 32], Option<HashSet<String>>>,
41}
42
43impl KeyPolicy {
44    /// Empty policy (nothing allowed).
45    pub fn new() -> Self {
46        Self::default()
47    }
48
49    /// Allow `key` for any project and for project-less batches.
50    pub fn allow_any(mut self, key: [u8; 32]) -> Self {
51        self.keys.insert(key, None);
52        self
53    }
54
55    /// Allow `key` for the listed projects only.
56    pub fn allow_projects(mut self, key: [u8; 32], projects: &[&str]) -> Self {
57        self.keys
58            .insert(key, Some(projects.iter().map(|p| p.to_string()).collect()));
59        self
60    }
61}
62
63impl ForwardPolicy for KeyPolicy {
64    fn allows(&self, key: &[u8; 32], project_id: Option<&str>) -> bool {
65        match self.keys.get(key) {
66            None => false,
67            Some(None) => true,
68            Some(Some(set)) => project_id.is_some_and(|p| set.contains(p)),
69        }
70    }
71}
72
73/// A node's `cog-store` service: verifies forwarded batches and writes them
74/// to the owning store.
75pub struct StoreOwnerService {
76    node_id: String,
77    key: SigningKey,
78    policy: Arc<dyn ForwardPolicy>,
79    directory: Arc<dyn StoreDirectory>,
80    nonces: NonceGuard,
81}
82
83impl StoreOwnerService {
84    /// Service for the node whose identity is `key`.
85    pub fn new(
86        key: SigningKey,
87        policy: Arc<dyn ForwardPolicy>,
88        directory: Arc<dyn StoreDirectory>,
89    ) -> Self {
90        Self {
91            node_id: node_id_from_pubkey(&key.verifying_key().to_bytes()),
92            key,
93            policy,
94            directory,
95            nonces: NonceGuard::new(),
96        }
97    }
98
99    /// This node's id.
100    pub fn node_id(&self) -> &str {
101        &self.node_id
102    }
103
104    /// Handle one signed request at `now_ms`. The flag is false when the
105    /// sender could not be authenticated (the connection is then dropped).
106    pub fn handle_at(&self, signed: &SignedCtl, now_ms: u64) -> (SignedCtl, bool) {
107        let (nonce, outcome, authed) = match self.check(signed, now_ms) {
108            Ok(req) => {
109                let nonce = req.nonce.clone();
110                (nonce, self.apply(req), true)
111            }
112            Err((nonce, code, reason, authed)) => {
113                (nonce, ForwardOutcome::Refused { code, reason }, authed)
114            }
115        };
116        let resp = ForwardResponse {
117            version: FORWARD_VERSION,
118            responder: self.node_id.clone(),
119            request_nonce: nonce,
120            outcome,
121        };
122        let payload = serde_json::to_string(&resp).unwrap_or_else(|_| "{}".into());
123        (sign(FORWARD_RESPONSE_DOMAIN, payload, &self.key), authed)
124    }
125
126    #[allow(clippy::type_complexity)]
127    fn check(
128        &self,
129        s: &SignedCtl,
130        now: u64,
131    ) -> Result<ForwardRequest, (String, ForwardRefusal, String, bool)> {
132        use ForwardRefusal as R;
133        let fail = |nonce: &str, c: R, m: &str, authed: bool| {
134            (nonce.to_string(), c, m.to_string(), authed)
135        };
136        let pk = open(FORWARD_DOMAIN, s).map_err(|r| fail("", R::Signature, &r.reason, false))?;
137        let req: ForwardRequest = serde_json::from_str(&s.payload)
138            .map_err(|_| fail("", R::Signature, "malformed request", false))?;
139        let n = req.nonce.clone();
140        if req.version != FORWARD_VERSION {
141            return Err(fail(&n, R::Invalid, "unsupported version", false));
142        }
143        if req.requester != node_id_from_pubkey(&pk) {
144            return Err(fail(&n, R::Signature, "requester is not the signing key's node", false));
145        }
146        if req.target != self.node_id {
147            return Err(fail(&n, R::NotForMe, "addressed to another node", false));
148        }
149        if req.expires_at_ms <= now {
150            return Err(fail(&n, R::Expired, "request expired", false));
151        }
152        if req.issued_at_ms > now.saturating_add(MAX_SKEW_MS)
153            || req.expires_at_ms < req.issued_at_ms
154            || req.expires_at_ms - req.issued_at_ms > MAX_TTL_MS
155        {
156            return Err(fail(&n, R::Expired, "request lifetime out of bounds", false));
157        }
158        if !is_lower_hex(&req.nonce, 32) {
159            return Err(fail(&n, R::Invalid, "nonce must be 32 hex", false));
160        }
161        if !self.policy.allows(&pk, req.project_id.as_deref()) {
162            return Err(fail(&n, R::Unauthorized, "key may not forward for this project", false));
163        }
164        if req.vectors.is_empty() || req.vectors.len() > MAX_VECTORS_PER_BATCH {
165            return Err(fail(&n, R::Invalid, "vector count out of bounds", true));
166        }
167        if req.vectors.iter().any(|v| v.values.iter().any(|f| !f.is_finite())) {
168            return Err(fail(&n, R::Invalid, "non-finite value", true));
169        }
170        if !self.nonces.admit(&req.nonce, req.expires_at_ms, now) {
171            return Err(fail(&n, R::Replay, "nonce already used", true));
172        }
173        Ok(req)
174    }
175
176    fn apply(&self, req: ForwardRequest) -> ForwardOutcome {
177        let refuse = |code, reason: &str| ForwardOutcome::Refused {
178            code,
179            reason: reason.into(),
180        };
181        let Some(store) = self.directory.store_for(req.project_id.as_deref()) else {
182            return refuse(ForwardRefusal::NoStore, "no store for this project here");
183        };
184        let from = Provenance {
185            instance_id: req.instance_id,
186            source_node: req.requester,
187        };
188        match store.ingest(&from, &req.vectors, req.dedup) {
189            Ok(result) => ForwardOutcome::Ok { result },
190            Err(StoreError::Full(_)) => refuse(ForwardRefusal::StoreFull, "store full"),
191            Err(StoreError::Backend(_)) => refuse(ForwardRefusal::Backend, "store backend error"),
192        }
193    }
194}
195
196/// Serve one connection as the owner's `cog-store` until the peer closes it.
197pub async fn serve_connection(
198    mut stream: Box<dyn MeshStream>,
199    svc: Arc<StoreOwnerService>,
200) -> Result<(), MeshError> {
201    loop {
202        let raw = match stream.recv().await {
203            Ok(r) => r,
204            Err(MeshError::ConnectionClosed) => return Ok(()),
205            Err(e) => return Err(e),
206        };
207        if raw.len() > MAX_CTL_BYTES + 4096 {
208            let _ = stream.close().await;
209            return Err(MeshError::Transport("oversize ingest frame".into()));
210        }
211        let env = MeshIpcEnvelope::from_bytes(&raw)
212            .map_err(|e| MeshError::Transport(format!("bad envelope: {e}")))?;
213        let signed = match (&env.message.target, &env.message.payload) {
214            (MessageTarget::ServiceMethod { service, method }, MessagePayload::Json(v))
215                if service == STORE_SERVICE && method == STORE_INGEST_METHOD =>
216            {
217                serde_json::from_value::<SignedCtl>(v.clone())
218                    .map_err(|e| MeshError::Transport(format!("bad signed payload: {e}")))?
219            }
220            _ => {
221                let _ = stream.close().await;
222                return Err(MeshError::Transport("not a cog-store request".into()));
223            }
224        };
225        let (resp, authed) = svc.handle_at(&signed, now_ms());
226        let mut msg = KernelMessage::new(
227            0,
228            owner_target(),
229            MessagePayload::Json(serde_json::to_value(&resp).unwrap_or_default()),
230        );
231        msg.correlation_id = env.message.correlation_id.clone();
232        let out = MeshIpcEnvelope::new(svc.node_id().to_string(), env.source_node, msg);
233        let bytes = out
234            .to_bytes()
235            .map_err(|e| MeshError::Transport(e.to_string()))?;
236        stream.send(&bytes).await?;
237        if !authed {
238            let _ = stream.close().await;
239            return Ok(());
240        }
241    }
242}
243
244/// Most concurrent sessions an owner listener serves.
245pub const MAX_SESSIONS: usize = 32;
246
247/// Accept connections on `listener` and serve each as a `cog-store` session
248/// (Noise XX responder when `noise` is set). Runs until the listener fails.
249pub async fn serve_listener(
250    mut listener: Box<dyn TransportListener>,
251    svc: Arc<StoreOwnerService>,
252    noise: bool,
253) -> Result<(), MeshError> {
254    let slots = Arc::new(tokio::sync::Semaphore::new(MAX_SESSIONS));
255    loop {
256        let (mut stream, peer) = listener.accept().await?;
257        let Ok(slot) = slots.clone().try_acquire_owned() else {
258            let _ = stream.close().await;
259            continue;
260        };
261        let svc = svc.clone();
262        tokio::spawn(async move {
263            let _slot = slot;
264            let stream: Box<dyn MeshStream> = if noise {
265                let cfg = crate::mesh_noise::NoiseConfig {
266                    pattern: crate::mesh_noise::NoisePattern::XX,
267                    local_private_key: rand::random(),
268                    remote_static_key: None,
269                };
270                match crate::mesh_noise::NoiseChannel::respond(stream, &cfg).await {
271                    Ok(ch) => Box::new(NoiseStream::new(Box::new(ch))),
272                    Err(e) => {
273                        tracing::warn!(%peer, error = %e, "cog-store noise handshake failed");
274                        return;
275                    }
276                }
277            } else {
278                stream
279            };
280            if let Err(e) = serve_connection(stream, svc).await {
281                tracing::debug!(%peer, error = %e, "cog-store session ended");
282            }
283        });
284    }
285}
286
287/// Reaches owners: `mem://<name>` to registered in-process services, any
288/// other address over mesh TCP (Noise XX when `noise` is set).
289#[derive(Default)]
290pub struct OwnerConnector {
291    noise: bool,
292    local: std::sync::RwLock<HashMap<String, Arc<StoreOwnerService>>>,
293}
294
295impl OwnerConnector {
296    /// Connector; `noise` wraps TCP streams in Noise XX.
297    pub fn new(noise: bool) -> Self {
298        Self {
299            noise,
300            local: Default::default(),
301        }
302    }
303
304    /// Register an in-process owner at `mem://<name>`; returns the address.
305    pub fn register_local(&self, name: &str, svc: Arc<StoreOwnerService>) -> String {
306        let addr = format!("{MEM_SCHEME}{name}");
307        if let Ok(mut m) = self.local.write() {
308            m.insert(addr.clone(), svc);
309        }
310        addr
311    }
312}
313
314#[async_trait]
315impl CtlConnector for OwnerConnector {
316    async fn connect(&self, addr: &str) -> Result<Box<dyn MeshStream>, MeshError> {
317        use crate::mesh::MeshTransport;
318        if addr.starts_with(MEM_SCHEME) {
319            let svc = self
320                .local
321                .read()
322                .ok()
323                .and_then(|m| m.get(addr).cloned())
324                .ok_or_else(|| MeshError::PeerNotConnected(addr.to_string()))?;
325            let (client, server) = crate::mesh_test_support::connected_pair().await?;
326            tokio::spawn(async move {
327                let _ = serve_connection(Box::new(server), svc).await;
328            });
329            return Ok(Box::new(client));
330        }
331        let stream = crate::mesh_tcp::TcpTransport.connect(addr).await?;
332        if !self.noise {
333            return Ok(stream);
334        }
335        let cfg = crate::mesh_noise::NoiseConfig {
336            pattern: crate::mesh_noise::NoisePattern::XX,
337            local_private_key: rand::random(),
338            remote_static_key: None,
339        };
340        let ch = crate::mesh_noise::NoiseChannel::initiate(stream, &cfg).await?;
341        Ok(Box::new(NoiseStream::new(Box::new(ch))))
342    }
343}
344
clawft_kernel/cog_ingest/owner.rs

clawft_kernel/cog_ingest/
owner.rs
