Skip to main content

clawft_kernel/cog_ingest/
hooks.rs

1//! The ingest bridge as the workload host sees it: authorise the placing
2//! project, issue a token with the host contract, bind it to the placement,
3//! and revoke it when the instance stops or unloads (a package revocation
4//! that force-unloads an instance goes through the same unload; no such
5//! path exists in the host yet).
6//!
7//! A cog instance registered by [`IngestHooks::lease`] lives in the host's
8//! in-memory instance table. Anything that re-creates an instance record
9//! without going through the host's `place` (a controller adopting an
10//! in-flight placement, a host restart that re-adopts running instances)
11//! must call `lease` again, or the cog's token is unknown and its posts are
12//! refused.
13
14use std::collections::{HashMap, HashSet};
15use std::net::{IpAddr, SocketAddr};
16use std::sync::Arc;
17
18use super::bridge::{BridgeHandle, BridgeScope, IngestBridge};
19use super::registry::{InstanceBinding, TokenRegistry};
20use super::types::{INGEST_PATH, IngestError};
21use crate::workload_runtime::HostContract;
22
23/// Which node a project's key is bound to, from the identity records.
24pub trait ProjectDirectory: Send + Sync {
25    /// Node id of the key currently bound to `project_id`; `None` when the
26    /// project is unregistered or every key of it is revoked.
27    fn bound_node(&self, project_id: &str) -> Option<String>;
28}
29
30/// A node's ingest wiring, shared by its workload hosts.
31#[derive(Clone)]
32pub struct IngestHooks {
33    /// This node's id: its own requests may place for any project it routes.
34    own_node: String,
35    registry: Arc<TokenRegistry>,
36    /// `None`: the bridge could not start (see `reason`).
37    bridge: Option<Arc<IngestBridge>>,
38    /// URL of the shared loopback bridge, handed to native cogs.
39    url: String,
40    /// Address to bind instance-scoped listeners on for container relays
41    /// (the engine or VM gateway). `None`: containers get no upstream and
42    /// rely on `network=host`.
43    scoped_bind: Option<IpAddr>,
44    reason: Option<String>,
45    /// Per project: the controllers (node ids) besides this node that may
46    /// place cogs for it.
47    controllers: HashMap<String, HashSet<String>>,
48    directory: Option<Arc<dyn ProjectDirectory>>,
49}
50
51impl IngestHooks {
52    /// Hooks over `registry` and `bridge`; native cogs are told `shared`
53    /// (where the shared listener is bound).
54    pub fn new(
55        own_node: impl Into<String>,
56        registry: Arc<TokenRegistry>,
57        bridge: Arc<IngestBridge>,
58        shared: SocketAddr,
59        scoped_bind: Option<IpAddr>,
60    ) -> Self {
61        Self {
62            own_node: own_node.into(),
63            registry,
64            bridge: Some(bridge),
65            url: format!("http://{shared}{INGEST_PATH}"),
66            scoped_bind,
67            reason: None,
68            controllers: HashMap::new(),
69            directory: None,
70        }
71    }
72
73    /// Hooks for a node whose bridge could not start. Cogs are placed with
74    /// no token and no URL, and every place result, status and
75    /// advertisement says `ingest: disabled`.
76    pub fn disabled(own_node: impl Into<String>, reason: impl Into<String>) -> Self {
77        Self {
78            own_node: own_node.into(),
79            registry: Arc::new(TokenRegistry::new()),
80            bridge: None,
81            url: String::new(),
82            scoped_bind: None,
83            reason: Some(reason.into()),
84            controllers: HashMap::new(),
85            directory: None,
86        }
87    }
88
89    /// Controllers (node ids) besides this node allowed to place for each
90    /// project.
91    pub fn with_project_controllers(mut self, c: HashMap<String, HashSet<String>>) -> Self {
92        self.controllers = c;
93        self
94    }
95
96    /// Identity records: a project's bound key may place for it.
97    pub fn with_project_directory(mut self, d: Arc<dyn ProjectDirectory>) -> Self {
98        self.directory = Some(d);
99        self
100    }
101
102    /// True when the bridge is running.
103    pub fn is_enabled(&self) -> bool {
104        self.bridge.is_some()
105    }
106
107    /// `enabled` or `disabled`.
108    pub fn state(&self) -> &'static str {
109        if self.is_enabled() { "enabled" } else { "disabled" }
110    }
111
112    /// Why the bridge is disabled.
113    pub fn disabled_reason(&self) -> Option<&str> {
114        self.reason.as_deref()
115    }
116
117    /// The token registry.
118    pub fn registry(&self) -> &Arc<TokenRegistry> {
119        &self.registry
120    }
121
122    /// May the controller in `b.controller_node` place a cog for
123    /// `b.project_id`, and is there a store owner for the placement? The
124    /// instance id is not known yet at place time and is ignored.
125    ///
126    /// A project needs one of: the controller is this node; the controller
127    /// is listed for the project in the config; the controller is the node
128    /// of the project's bound key. Whatever the project, a route to a store
129    /// owner must exist, so a placement whose vectors could not be delivered
130    /// is refused instead of placed.
131    pub fn authorize(&self, b: &InstanceBinding) -> Result<(), IngestError> {
132        if let Some(p) = &b.project_id {
133            let me = b.controller_node == self.own_node;
134            let listed = self
135                .controllers
136                .get(p)
137                .is_some_and(|s| s.contains(&b.controller_node));
138            let bound = self
139                .directory
140                .as_ref()
141                .and_then(|d| d.bound_node(p))
142                .is_some_and(|n| n == b.controller_node);
143            if !(me || listed || bound) {
144                return Err(IngestError::Forbidden);
145            }
146        }
147        // With the bridge down nothing is delivered, so there is no route to
148        // require; the project check above still holds, so a degraded record
149        // never carries an unverified project id.
150        let Some(bridge) = &self.bridge else {
151            return Ok(());
152        };
153        if !bridge.has_route(b) {
154            return Err(IngestError::NotRouted(match &b.project_id {
155                Some(p) => format!("project {p} is not routed here"),
156                None => format!("no store route for controller {}", b.controller_node),
157            }));
158        }
159        Ok(())
160    }
161
162    /// Complete `contract` for an instance on `route`: native cogs get the
163    /// shared bridge URL; other routes get their own token-scoped listener
164    /// (when `scoped_bind` is set) as their `ingest_upstream`.
165    pub async fn prepare(
166        &self,
167        route: &str,
168        contract: HostContract,
169    ) -> Result<(HostContract, Option<BridgeHandle>), IngestError> {
170        let Some(bridge) = &self.bridge else {
171            return Ok((contract, None));
172        };
173        let Some(ip) = self.scoped_bind.filter(|_| route != "native") else {
174            return Ok((contract.with_ingest_url(self.url.clone()), None));
175        };
176        let h = *blake3::hash(contract.token.expose().as_bytes()).as_bytes();
177        let handle = bridge
178            .bind(SocketAddr::new(ip, 0), BridgeScope::Token(h))
179            .await?;
180        let addr = handle.addr();
181        let contract = contract
182            .with_ingest_upstream(addr)
183            .with_ingest_url(format!("http://{addr}{INGEST_PATH}"));
184        Ok((contract, Some(handle)))
185    }
186
187    /// Bind a loaded instance to its placement and accept its token.
188    pub fn lease(
189        &self,
190        binding: InstanceBinding,
191        contract: HostContract,
192        listener: Option<BridgeHandle>,
193    ) -> Result<IngestLease, IngestError> {
194        self.registry.register(binding.clone(), &contract)?;
195        Ok(IngestLease {
196            binding,
197            contract,
198            _listener: listener,
199        })
200    }
201
202    /// Accept the token again (after a stop and a later start).
203    pub fn activate(&self, l: &IngestLease) -> Result<(), IngestError> {
204        if self.registry.contains(&l.binding.instance_id) {
205            return Ok(());
206        }
207        self.registry.register(l.binding.clone(), &l.contract)
208    }
209
210    /// Stop accepting the instance's token.
211    pub fn deactivate(&self, l: &IngestLease) {
212        self.registry.revoke(&l.binding.instance_id);
213        if let Some(b) = &self.bridge {
214            b.forget(&l.binding.instance_id);
215        }
216    }
217}
218
219/// One instance's ingest registration. Dropping it closes its scoped
220/// listener; revoke the token with [`IngestHooks::deactivate`].
221pub struct IngestLease {
222    /// The placement binding.
223    pub binding: InstanceBinding,
224    contract: HostContract,
225    _listener: Option<BridgeHandle>,
226}