1use std::collections::HashSet;
18use std::sync::Arc;
19
20use clawft_types::placement::{Capability, CapabilityId, NodeFacts, Provenance};
21
22use crate::chain::{
23 EVENT_KIND_ARTIFACT_EVICT, EVENT_KIND_ARTIFACT_PEER_BAN, EVENT_KIND_ARTIFACT_REVOKE,
24 EVENT_KIND_ARTIFACT_SEED,
25};
26use crate::mesh_artifact::ArtifactExchange;
27use crate::mesh_artifact_types::{ArtifactDescriptor, ArtifactId};
28use crate::mesh_swarm_picker::LinkStats;
29use crate::mesh_swarm_state::{Audience, GrantInfo};
30use crate::revocation::{RevocationList, RevokedSubject};
31use crate::workload_pkg::codec::hex_encode;
32
33pub const MAX_ADVERTISED_ARTIFACTS: usize = 128;
35
36#[derive(Debug, Clone, PartialEq, Eq)]
38pub struct Forgotten {
39 pub id: ArtifactId,
41 pub content_hash: [u8; 32],
43 pub bytes_freed: u64,
45}
46
47#[derive(Debug, Clone, PartialEq, Eq)]
49pub struct RevokedArtifact {
50 pub content_hash: [u8; 32],
52 pub package_id: String,
54 pub subject: RevokedSubject,
56 pub bytes_freed: u64,
58}
59
60impl ArtifactExchange {
61 pub fn set_revocations(&self, list: Arc<RevocationList>) -> bool {
65 self.swarm.revocations.set(list).is_ok()
66 }
67
68 pub(crate) fn is_revoked_subject(
69 &self,
70 package_id: &str,
71 signers: &[String],
72 content_hash: &[u8; 32],
73 ) -> bool {
74 let Some(list) = self.swarm.revocations.get() else {
75 return false;
76 };
77 let hex = hex_encode(content_hash);
78 list.first_revoked(Some(package_id), signers.iter(), std::iter::once(&hex))
79 .is_some()
80 }
81
82 pub fn is_revoked(&self, content_hash: &[u8; 32]) -> bool {
86 match self.grants.get(content_hash) {
87 Some(gs) if !gs.is_empty() => gs
88 .iter()
89 .all(|g| self.is_revoked_subject(&g.package_id, &g.signers, content_hash)),
90 _ => self.is_revoked_subject("", &[], content_hash),
91 }
92 }
93
94 pub(crate) fn servable_grant(
98 &self,
99 content_hash: &[u8; 32],
100 audience: &Audience<'_>,
101 ) -> Option<GrantInfo> {
102 let gs = self.grants.get(content_hash)?;
103 if !self.swarm.policy.allows(content_hash, &gs, audience) {
104 return None;
105 }
106 gs.iter()
107 .find(|g| !self.is_revoked_subject(&g.package_id, &g.signers, content_hash))
108 .cloned()
109 }
110
111 pub fn apply_revocations(&self) -> Vec<RevokedArtifact> {
128 let Some(list) = self.swarm.revocations.get().cloned() else {
129 return Vec::new();
130 };
131 let mut hits: Vec<([u8; 32], String, RevokedSubject)> = Vec::new();
133 let mut emptied: Vec<[u8; 32]> = Vec::new();
134 for entry in self.grants.iter() {
135 let hex = hex_encode(entry.key());
136 let mut left = 0;
137 for g in entry.value() {
138 match list.first_revoked(
139 Some(&g.package_id),
140 g.signers.iter(),
141 std::iter::once(&hex),
142 ) {
143 Some(subject) => hits.push((*entry.key(), g.package_id.clone(), subject)),
144 None => left += 1,
145 }
146 }
147 if left == 0 {
148 emptied.push(*entry.key());
149 }
150 }
151 for s in list.list_subjects(Some(crate::revocation::RevocationKind::ArtifactHash)) {
154 let Some(hash) = crate::workload_pkg::codec::hex_decode_exact::<32>(&s.id) else {
155 continue;
156 };
157 let held = self.by_content.contains_key(&hash)
158 || self.pending.iter().any(|p| p.content_hash == hash);
159 if held && !hits.iter().any(|h| h.0 == hash) {
160 hits.push((hash, String::new(), s));
161 emptied.push(hash);
162 }
163 }
164 let mut out = Vec::new();
165 for (content_hash, package_id, subject) in hits {
166 let evict = emptied.contains(&content_hash);
167 if let Some(mut gs) = self.grants.get_mut(&content_hash) {
168 gs.retain(|g| g.package_id != package_id);
169 }
170 self.grants.remove_if(&content_hash, |_, gs| gs.is_empty());
171 let mut ids: Vec<ArtifactId> = Vec::new();
172 let mut bytes_freed = 0;
173 if evict {
174 self.swarm.seeded.remove(&content_hash);
175 ids = self
176 .by_content
177 .get(&content_hash)
178 .map(|i| *i)
179 .into_iter()
180 .chain(
181 self.pending
182 .iter()
183 .filter(|p| p.content_hash == content_hash)
184 .map(|p| p.id()),
185 )
186 .collect();
187 for id in &ids {
188 if let Some(f) = self.forget(id, "revoked") {
189 bytes_freed += f.bytes_freed;
190 }
191 }
192 }
193 self.chain_event(
194 EVENT_KIND_ARTIFACT_REVOKE,
195 serde_json::json!({
196 "subject_kind": subject.kind.to_string(),
197 "subject_id": subject.id,
198 "reason": subject.reason,
199 "content_hash": hex_encode(&content_hash),
200 "package_id": package_id,
201 "artifact_ids": ids.iter().map(|i| i.to_string()).collect::<Vec<_>>(),
202 "evicted": evict,
203 "bytes_freed": bytes_freed,
204 "node": self.node_id(),
205 }),
206 );
207 out.push(RevokedArtifact {
208 content_hash,
209 package_id,
210 subject,
211 bytes_freed,
212 });
213 }
214 out
215 }
216
217 pub fn ban_peer(&self, peer: &str, reason: &str) -> bool {
222 if self.swarm.bans.insert(peer.to_string(), reason.to_string()).is_some() {
223 return false;
224 }
225 self.chain_event(
226 EVENT_KIND_ARTIFACT_PEER_BAN,
227 serde_json::json!({ "peer": peer, "reason": reason, "node": self.node_id() }),
228 );
229 true
230 }
231
232 pub fn unban_peer(&self, peer: &str) -> bool {
234 self.swarm.corrupt.remove(peer);
235 self.swarm.bans.remove(peer).is_some()
236 }
237
238 pub fn is_banned(&self, peer: &str) -> bool {
240 self.swarm.bans.contains_key(peer)
241 }
242
243 pub fn banned_peers(&self) -> Vec<(String, String)> {
245 let mut v: Vec<_> = self
246 .swarm
247 .bans
248 .iter()
249 .map(|e| (e.key().clone(), e.value().clone()))
250 .collect();
251 v.sort();
252 v
253 }
254
255 pub(crate) fn report_corrupt_piece(&self, id: &ArtifactId, index: u32, peer: &str) -> bool {
259 self.chain_piece_rejected(id, index, peer, "hash mismatch");
260 let n = {
261 let mut e = self.swarm.corrupt.entry(peer.to_string()).or_insert(0);
262 *e += 1;
263 *e
264 };
265 if n >= self.config().ban_after_corrupt {
266 self.ban_peer(peer, &format!("served a corrupt piece of {id}"));
267 }
268 self.is_banned(peer)
269 }
270
271 pub fn link_stats(&self) -> &LinkStats {
275 &self.swarm.links
276 }
277
278 pub(crate) fn on_granted(&self, content_hash: &[u8; 32]) {
279 self.maybe_seed(content_hash);
280 }
281
282 pub(crate) fn on_verified(&self, id: &ArtifactId, content_hash: &[u8; 32]) {
283 self.maybe_seed(content_hash);
284 if let Some(cache) = self.swarm.cache() {
285 cache.on_verified(id);
286 }
287 }
288
289 fn maybe_seed(&self, content_hash: &[u8; 32]) {
292 let Some(id) = self.by_content.get(content_hash).map(|i| *i) else {
293 return;
294 };
295 let Some(g) = self.servable_grant(content_hash, &Audience::Seed) else {
296 return;
297 };
298 if self.is_revoked(content_hash) || self.swarm.seeded.insert(*content_hash, ()).is_some() {
299 return;
300 }
301 let size = self.descriptor(&id).map_or(0, |d| d.total_size);
302 self.chain_event(
303 EVENT_KIND_ARTIFACT_SEED,
304 serde_json::json!({
305 "artifact_id": id.to_string(),
306 "content_hash": hex_encode(content_hash),
307 "package_id": g.package_id,
308 "size": size,
309 "node": self.node_id(),
310 }),
311 );
312 }
313
314 pub fn pending_bytes(&self) -> u64 {
316 self.pending.iter().map(|d| d.total_size).sum()
317 }
318
319 pub fn servable_artifacts(&self) -> Vec<ArtifactDescriptor> {
321 self.descriptors
322 .iter()
323 .map(|d| d.clone())
324 .filter(|d| self.is_servable(d))
325 .collect()
326 }
327
328 pub fn forget(&self, id: &ArtifactId, reason: &str) -> Option<Forgotten> {
337 if reason != "revoked" && self.swarm.cache().is_some_and(|c| c.is_pinned(id)) {
338 return None;
339 }
340 let verified = self.descriptors.remove(id).map(|(_, d)| d);
341 let pending = self.pending.remove(id).map(|(_, d)| d);
342 let d = verified.or(pending)?;
343 self.by_content.remove_if(&d.content_hash, |_, v| v == id);
344 self.swarm.seeded.remove(&d.content_hash);
345 self.peer_haves.retain(|k, _| k.0 != *id);
346 let mut keep: HashSet<[u8; 32]> = HashSet::new();
347 for v in self.descriptors.iter() {
348 keep.insert(v.content_hash);
349 keep.extend(v.pieces.iter().copied());
350 }
351 for p in self.pending.iter() {
352 keep.insert(p.content_hash);
353 keep.extend(p.pieces.iter().copied());
354 }
355 let mut freed = 0u64;
356 let candidates = d.pieces.iter().copied().chain(std::iter::once(d.content_hash));
357 let mut seen = HashSet::new();
358 for hash in candidates {
359 if keep.contains(&hash) || !seen.insert(hash) {
360 continue;
361 }
362 if !self.swarm.owned.contains_key(&hash) {
363 continue; }
365 let key = hex_encode(&hash);
366 if let Some((_, size, _, _, refs)) = self.store.metadata(&key) {
367 if refs > 1 {
368 self.swarm.owned.remove(&hash);
370 } else if self.store.remove(&key).is_ok() {
371 self.swarm.owned.remove(&hash);
372 freed += size;
373 }
374 }
375 }
376 if let Some(cache) = self.swarm.cache() {
377 cache.on_forgotten(id);
378 }
379 self.chain_event(
380 EVENT_KIND_ARTIFACT_EVICT,
381 serde_json::json!({
382 "artifact_id": id.to_string(),
383 "content_hash": d.content_hex(),
384 "size": d.total_size,
385 "bytes_freed": freed,
386 "reason": reason,
387 "node": self.node_id(),
388 }),
389 );
390 Some(Forgotten {
391 id: *id,
392 content_hash: d.content_hash,
393 bytes_freed: freed,
394 })
395 }
396
397 pub(crate) fn chain_event(&self, kind: &str, payload: serde_json::Value) {
398 if let Some(cm) = &self.chain {
399 cm.append("mesh_artifact", kind, Some(payload));
400 }
401 }
402
403 pub fn held_capabilities(&self, model_shards: &HashSet<[u8; 32]>) -> Vec<Capability> {
411 let mut held: Vec<ArtifactDescriptor> = self
414 .servable_artifacts()
415 .into_iter()
416 .filter(|d| {
417 self.grants
418 .get(&d.content_hash)
419 .is_some_and(|gs| gs.iter().all(GrantInfo::is_opt_in))
420 })
421 .collect();
422 held.sort_by_key(|d| d.content_hash);
423 held.truncate(MAX_ADVERTISED_ARTIFACTS);
424 let mut caps = Vec::new();
425 let mut shards = Vec::new();
426 for d in &held {
427 let hex = d.content_hex();
428 if let Ok(id) = CapabilityId::new(format!("store.artifact.{}", &hex[..16])) {
429 let size = i64::try_from(d.total_size).unwrap_or(i64::MAX);
430 caps.push(
431 Capability::new(id, Provenance::Probed)
432 .with_attr("content", hex.as_str())
433 .with_attr("size", size),
434 );
435 }
436 if model_shards.contains(&d.content_hash) {
437 shards.push(hex);
438 }
439 }
440 if !shards.is_empty()
441 && let Ok(id) = CapabilityId::new("model.present")
442 {
443 let list = crate::node_facts::probe::str_list(&shards);
444 caps.push(Capability::new(id, Provenance::Probed).with_attr("shards", list));
445 }
446 caps
447 }
448}
449
450pub fn set_held_capabilities(facts: &mut NodeFacts, held: Vec<Capability>) {
453 facts
454 .capabilities
455 .retain(|c| !c.id.is_under("store.artifact") && c.id.as_str() != "model.present");
456 facts.capabilities.extend(held);
457}