clawft_kernel/cog_ingest/
hooks.rs1use 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
23pub trait ProjectDirectory: Send + Sync {
25 fn bound_node(&self, project_id: &str) -> Option<String>;
28}
29
30#[derive(Clone)]
32pub struct IngestHooks {
33 own_node: String,
35 registry: Arc<TokenRegistry>,
36 bridge: Option<Arc<IngestBridge>>,
38 url: String,
40 scoped_bind: Option<IpAddr>,
44 reason: Option<String>,
45 controllers: HashMap<String, HashSet<String>>,
48 directory: Option<Arc<dyn ProjectDirectory>>,
49}
50
51impl IngestHooks {
52 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 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 pub fn with_project_controllers(mut self, c: HashMap<String, HashSet<String>>) -> Self {
92 self.controllers = c;
93 self
94 }
95
96 pub fn with_project_directory(mut self, d: Arc<dyn ProjectDirectory>) -> Self {
98 self.directory = Some(d);
99 self
100 }
101
102 pub fn is_enabled(&self) -> bool {
104 self.bridge.is_some()
105 }
106
107 pub fn state(&self) -> &'static str {
109 if self.is_enabled() { "enabled" } else { "disabled" }
110 }
111
112 pub fn disabled_reason(&self) -> Option<&str> {
114 self.reason.as_deref()
115 }
116
117 pub fn registry(&self) -> &Arc<TokenRegistry> {
119 &self.registry
120 }
121
122 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 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 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 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 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 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
219pub struct IngestLease {
222 pub binding: InstanceBinding,
224 contract: HostContract,
225 _listener: Option<BridgeHandle>,
226}