1use std::collections::HashMap;
12use std::sync::atomic::{AtomicUsize, Ordering};
13use std::sync::{Arc, Mutex, MutexGuard, Weak};
14use std::time::Duration;
15
16use tokio::sync::watch;
17use tokio::time::Instant;
18
19use crate::cbor::Value;
20use crate::frame::StreamMode;
21use crate::manifest::{self, block_mcid, chunk_mcid, Manifest, Mcid, DEFAULT_CHUNK_SIZE};
22use crate::record::{
23 self, new_content_announcement, ContentAnnouncementOptions, Reason, Record, RecordType,
24 TombstoneOptions,
25};
26use crate::station_link::{
27 self, stream_handler, Link, LinkError, Stream, StreamEvent, DEFAULT_CALL_TIMEOUT,
28};
29
30use super::{shuffle, Offer, Pool, PoolError, PoolInner, Served};
31
32pub const CONTENT_PROCEDURE: &str = "content_v1";
34
35const ANNOUNCEMENT_TTL: Duration = Duration::from_secs(60 * 60);
37
38const MAX_BLOCK_BYTES: u64 = DEFAULT_CHUNK_SIZE;
40
41#[derive(Debug, Clone, Copy, PartialEq, Eq)]
44pub struct ContentOptions {
45 pub max_bytes: u64,
46 pub max_chunks: u64,
47 pub parallel: usize,
48 pub chunk_timeout: Duration,
49}
50
51impl Default for ContentOptions {
52 fn default() -> Self {
53 ContentOptions {
54 max_bytes: 256 << 20,
55 max_chunks: 16_384,
56 parallel: 4,
57 chunk_timeout: Duration::from_secs(15),
58 }
59 }
60}
61
62#[derive(Default)]
66pub(super) struct Sharer {
67 realms: Mutex<HashMap<[u8; 32], SharedRealm>>,
68 serving: tokio::sync::Mutex<()>,
70}
71
72struct SharedRealm {
73 _served: Served,
74 roots: HashMap<Mcid, Root>,
75 chunks: HashMap<Mcid, Vec<u8>>,
76 announcements: HashMap<Mcid, Announced>,
77}
78
79#[derive(Clone)]
80enum Root {
81 Block(Vec<u8>),
82 Chunked(Arc<Manifest>),
83}
84
85struct Announced {
86 latest: Arc<Mutex<Record>>,
87 stop: watch::Sender<bool>,
88}
89
90impl Sharer {
91 fn lock(&self) -> MutexGuard<'_, HashMap<[u8; 32], SharedRealm>> {
92 self.realms.lock().unwrap_or_else(|p| p.into_inner())
93 }
94}
95
96pub fn content_procedure_bound(procedure: &str, node: &[u8; 32]) -> bool {
101 let node_hex: String = node.iter().map(|b| format!("{b:02x}")).collect();
102 if procedure == record::own_procedure(node, CONTENT_PROCEDURE) {
103 return true;
104 }
105 match procedure.split_once('/') {
106 Some((org, name)) => {
107 !org.is_empty()
108 && org != "_"
109 && !org.starts_with(record::OWN_NAMESPACE_PREFIX)
110 && name == format!("{CONTENT_PROCEDURE}_{node_hex}")
111 }
112 None => false,
113 }
114}
115
116impl Pool {
117 pub async fn share_content(
123 &self,
124 realm: &[u8; 32],
125 data: &[u8],
126 name: &str,
127 ) -> Result<Mcid, PoolError> {
128 let inner = &self.inner;
129 let (mcid, root, chunks) = if data.len() as u64 <= MAX_BLOCK_BYTES {
130 (block_mcid(data), Root::Block(data.to_vec()), HashMap::new())
131 } else {
132 let (m, parts) = manifest::create(data, name, DEFAULT_CHUNK_SIZE)
133 .map_err(|e| PoolError::ContentMismatch(e.to_string()))?;
134 let chunks: HashMap<Mcid, Vec<u8>> = parts
135 .into_iter()
136 .enumerate()
137 .filter_map(|(i, part)| chunk_mcid(&m, i).map(|c| (c, part)))
138 .collect();
139 (m.mcid, Root::Chunked(Arc::new(m)), chunks)
140 };
141 inner.shared_realm(realm).await?;
142 let already = {
143 let mut realms = inner.content.lock();
144 let shared = realms.get_mut(realm).ok_or(PoolError::Closed)?;
145 shared.roots.insert(mcid, root.clone());
146 shared.chunks.extend(chunks);
147 shared.announcements.contains_key(&mcid)
148 };
149 if already {
150 return Ok(mcid);
151 }
152 let latest = match inner.announce_content(realm, &mcid, &root).await {
153 Ok(latest) => latest,
154 Err(e) => {
155 inner.forget_shared(realm, &mcid);
156 return Err(e);
157 }
158 };
159 let latest = Arc::new(Mutex::new(latest));
160 let (stop, stopped) = watch::channel(false);
161 if let Some(shared) = inner.content.lock().get_mut(realm) {
162 shared.announcements.insert(
163 mcid,
164 Announced {
165 latest: latest.clone(),
166 stop,
167 },
168 );
169 }
170 tokio::spawn(renew_announcement(
171 Arc::downgrade(inner),
172 *realm,
173 mcid,
174 root,
175 latest,
176 stopped,
177 ));
178 Ok(mcid)
179 }
180
181 pub async fn unshare_content(&self, realm: &[u8; 32], mcid: &Mcid) -> Result<(), PoolError> {
185 let Some(announced) = self.inner.forget_shared(realm, mcid) else {
186 return Ok(());
187 };
188 announced.stop.send_replace(true);
189 let latest = announced
190 .latest
191 .lock()
192 .unwrap_or_else(|p| p.into_inner())
193 .clone();
194 let tombstone =
195 record::new_tombstone(&latest, Reason::Shutdown, &TombstoneOptions::default())
196 .map_err(LinkError::from)?;
197 let signed =
198 record::sign(&tombstone, &self.inner.opts.identity).map_err(LinkError::from)?;
199 let wire = record::encode(&signed).map_err(LinkError::from)?;
200 self.put_record(&wire).await
201 }
202
203 pub async fn get_content(
215 &self,
216 realm: &[u8; 32],
217 mcid: &Mcid,
218 opts: ContentOptions,
219 ) -> Result<Vec<u8>, PoolError> {
220 let sharers = self.inner.content_sharers(realm, mcid).await?;
221 if sharers.is_empty() {
222 return Err(PoolError::NotShared);
223 }
224 let mut tried = Vec::new();
225 for sharer in sharers {
226 match self.inner.fetch_from(realm, &sharer, mcid, &opts).await {
227 Ok(data) => return Ok(data),
228 Err(e) => tried.push((sharer.node, e)),
229 }
230 }
231 Err(PoolError::ContentUnavailable(tried))
232 }
233}
234
235#[derive(Debug, Clone)]
237struct Sharing {
238 node: [u8; 32],
239 station: [u8; 32],
240 procedure: String,
241}
242
243impl PoolInner {
244 async fn shared_realm(self: &Arc<Self>, realm: &[u8; 32]) -> Result<(), PoolError> {
246 let _serving = self.content.serving.lock().await;
247 if self.content.lock().contains_key(realm) {
248 return Ok(());
249 }
250 let pool = Arc::downgrade(self);
251 let r = *realm;
252 let answer = stream_handler(move |s| {
253 let pool = pool.clone();
254 async move {
255 match pool.upgrade() {
256 Some(inner) => inner.answer_fetch(&r, &s).await,
257 None => {
258 s.abort("not_shared", "the node no longer shares content")
259 .await
260 }
261 }
262 .map_err(|e| e.to_string())
263 }
264 });
265 let procedure = record::own_procedure(&self.self_id, CONTENT_PROCEDURE);
266 let served = Pool {
267 inner: self.clone(),
268 }
269 .serve(Offer::stream(
270 *realm,
271 &procedure,
272 StreamMode::ServerStream,
273 answer,
274 ))
275 .await?;
276 self.content.lock().insert(
277 *realm,
278 SharedRealm {
279 _served: served,
280 roots: HashMap::new(),
281 chunks: HashMap::new(),
282 announcements: HashMap::new(),
283 },
284 );
285 Ok(())
286 }
287
288 async fn answer_fetch(&self, realm: &[u8; 32], s: &Stream) -> Result<(), LinkError> {
291 let args = &s.request().payload;
292 let mcid: Option<Mcid> = match wire_bytes(args, "mcid") {
293 Some(b) if b.len() == 50 && b[0] == 2 => b.as_slice().try_into().ok(),
294 _ => None,
295 };
296 let want = wire_text(args, "want");
297 let (Some(mcid), Some(want @ ("root" | "block"))) = (mcid, want.as_deref()) else {
298 return s
299 .abort(
300 "malformed",
301 "a fetch names one content id and wants root or block",
302 )
303 .await;
304 };
305 let body = self
306 .content
307 .lock()
308 .get(realm)
309 .and_then(|shared| shared.body(want, &mcid));
310 match body {
311 None => {
312 s.abort("not_shared", "this node does not share that content")
313 .await
314 }
315 Some(body) => {
316 s.send_value(body).await?;
317 s.close().await
318 }
319 }
320 }
321
322 fn forget_shared(&self, realm: &[u8; 32], mcid: &Mcid) -> Option<Announced> {
325 let mut realms = self.content.lock();
326 let shared = realms.get_mut(realm)?;
327 if let Some(Root::Chunked(m)) = shared.roots.remove(mcid) {
328 for i in 0..m.chunks.len() {
329 if let Some(c) = chunk_mcid(&m, i) {
330 shared.chunks.remove(&c);
331 }
332 }
333 }
334 shared.announcements.remove(mcid)
335 }
336
337 async fn announce_content(
340 self: &Arc<Self>,
341 realm: &[u8; 32],
342 mcid: &Mcid,
343 root: &Root,
344 ) -> Result<Record, PoolError> {
345 let station = self
346 .links()
347 .first()
348 .map(Link::station_node_id)
349 .ok_or(PoolError::NoLink(Vec::new()))?;
350 let mut opts = ContentAnnouncementOptions {
351 realm_id: *realm,
352 serving_station: station,
353 procedure: record::own_procedure(&self.self_id, CONTENT_PROCEDURE),
354 ttl_ms: ANNOUNCEMENT_TTL.as_millis() as u64,
355 ..ContentAnnouncementOptions::default()
356 };
357 match root {
358 Root::Chunked(m) => {
359 opts.name = String::from_utf8_lossy(&m.name).into_owned();
360 opts.size = Some(m.size);
361 opts.chunk_count = Some(m.chunk_count);
362 }
363 Root::Block(b) => opts.size = Some(b.len() as u64),
364 }
365 let unsigned =
366 new_content_announcement(&self.self_id, mcid, &opts).map_err(LinkError::from)?;
367 let signed = record::sign(&unsigned, &self.opts.identity).map_err(LinkError::from)?;
368 let wire = record::encode(&signed).map_err(LinkError::from)?;
369 Pool {
370 inner: self.clone(),
371 }
372 .put_record(&wire)
373 .await?;
374 Ok(signed)
375 }
376
377 async fn content_sharers(
380 &self,
381 realm: &[u8; 32],
382 mcid: &Mcid,
383 ) -> Result<Vec<Sharing>, PoolError> {
384 let key = record::content_key(mcid).map_err(LinkError::from)?;
385 let found = match self
386 .first_answer(|l| async move { l.find_records(&key).await })
387 .await
388 {
389 Ok((found, _)) => found,
390 Err(PoolError::Link(LinkError::RecordNotFound)) => Vec::new(),
391 Err(e) => return Err(e),
392 };
393 let mut out: Vec<Sharing> = found
394 .iter()
395 .filter(|v| v.record().record_type == RecordType::CONTENT_ANNOUNCEMENT)
396 .filter_map(|v| record::read_content_announcement(v.record()).ok())
397 .filter(|a| {
398 a.mcid.as_slice() == mcid.as_slice()
399 && a.realm_id == *realm
400 && a.serving_station != [0; 32]
401 && content_procedure_bound(&a.procedure, &a.announcer_node)
402 })
403 .map(|a| Sharing {
404 node: a.announcer_node,
405 station: a.serving_station,
406 procedure: a.procedure,
407 })
408 .collect();
409 shuffle(&mut out);
410 Ok(out)
411 }
412
413 async fn fetch_from(
415 self: &Arc<Self>,
416 realm: &[u8; 32],
417 s: &Sharing,
418 mcid: &Mcid,
419 opts: &ContentOptions,
420 ) -> Result<Vec<u8>, PoolError> {
421 let link = self
422 .link_to(&s.station, Instant::now() + DEFAULT_CALL_TIMEOUT)
423 .await?;
424 let (kind, body) = fetch_one(&link, realm, s, mcid, "root", opts.chunk_timeout).await?;
425 if kind == "block" {
426 let bytes = wire_bytes(&body, "bytes").unwrap_or_default();
427 if block_mcid(&bytes) != *mcid {
428 return Err(PoolError::ContentMismatch("the block".into()));
429 }
430 return Ok(bytes);
431 }
432 let m = manifest::from_wire(body.get("manifest").unwrap_or(&Value::Null))
433 .map_err(|e| PoolError::ContentReply(format!("the manifest: {e}")))?;
434 manifest::verify_mcid(&m, mcid)
435 .map_err(|e| PoolError::ContentMismatch(format!("the manifest: {e}")))?;
436 if m.size > opts.max_bytes || m.chunk_count > opts.max_chunks {
437 return Err(PoolError::ContentTooLarge(format!(
438 "{} bytes in {} chunks",
439 m.size, m.chunk_count
440 )));
441 }
442 manifest::check_whole(&m)
443 .map_err(|e| PoolError::ContentMismatch(format!("the manifest: {e}")))?;
444 fetch_chunks(link, *realm, s.clone(), Arc::new(m), opts).await
445 }
446}
447
448impl SharedRealm {
449 fn body(&self, want: &str, mcid: &Mcid) -> Option<Value> {
452 let block = |b: &[u8]| {
453 Value::Map(vec![
454 (Value::text("kind"), Value::text("block")),
455 (Value::text("mcid"), Value::Bytes(mcid.to_vec())),
456 (Value::text("bytes"), Value::Bytes(b.to_vec())),
457 ])
458 };
459 if want == "block" {
460 return self.chunks.get(mcid).map(|b| block(b));
461 }
462 match self.roots.get(mcid)? {
463 Root::Block(b) => Some(block(b)),
464 Root::Chunked(m) => Some(Value::Map(vec![
465 (Value::text("kind"), Value::text("manifest")),
466 (Value::text("mcid"), Value::Bytes(mcid.to_vec())),
467 (Value::text("manifest"), manifest::to_wire(m)),
468 ])),
469 }
470 }
471}
472
473async fn renew_announcement(
477 pool: Weak<PoolInner>,
478 realm: [u8; 32],
479 mcid: Mcid,
480 root: Root,
481 latest: Arc<Mutex<Record>>,
482 mut stopped: watch::Receiver<bool>,
483) {
484 loop {
485 tokio::select! {
486 _ = stopped.wait_for(|s| *s) => return,
487 _ = tokio::time::sleep(ANNOUNCEMENT_TTL / 2) => {}
488 }
489 let Some(inner) = pool.upgrade() else { return };
490 if inner.lock().closed {
491 return;
492 }
493 let renewed = tokio::time::timeout(
494 DEFAULT_CALL_TIMEOUT,
495 inner.announce_content(&realm, &mcid, &root),
496 )
497 .await;
498 if let Ok(Ok(record)) = renewed {
499 *latest.lock().unwrap_or_else(|p| p.into_inner()) = record;
500 }
501 }
502}
503
504async fn fetch_chunks(
508 link: Link,
509 realm: [u8; 32],
510 s: Sharing,
511 m: Arc<Manifest>,
512 opts: &ContentOptions,
513) -> Result<Vec<u8>, PoolError> {
514 let count = m.chunks.len();
515 let parts: Arc<Mutex<Vec<Option<Vec<u8>>>>> = Arc::new(Mutex::new(vec![None; count]));
516 let next = Arc::new(AtomicUsize::new(0));
517 let (failed_tx, failed) = watch::channel::<Option<PoolError>>(None);
518 let failed_tx = Arc::new(failed_tx);
519 let mut workers = tokio::task::JoinSet::new();
520 for _ in 0..opts.parallel.max(1).min(count) {
521 let (link, s, m, parts, next, failed_tx) = (
522 link.clone(),
523 s.clone(),
524 m.clone(),
525 parts.clone(),
526 next.clone(),
527 failed_tx.clone(),
528 );
529 let timeout = opts.chunk_timeout;
530 workers.spawn(async move {
531 loop {
532 if failed_tx.borrow().is_some() {
533 return;
534 }
535 let i = next.fetch_add(1, Ordering::SeqCst);
536 let Some(want) = chunk_mcid(&m, i) else {
537 return;
538 };
539 let fetched = fetch_one(&link, &realm, &s, &want, "block", timeout).await;
540 let outcome = fetched.and_then(|(_, body)| {
541 let bytes = wire_bytes(&body, "bytes").unwrap_or_default();
542 if block_mcid(&bytes) == want {
543 Ok(bytes)
544 } else {
545 Err(PoolError::ContentMismatch(format!("chunk {i}")))
546 }
547 });
548 match outcome {
549 Ok(bytes) => parts.lock().unwrap_or_else(|p| p.into_inner())[i] = Some(bytes),
550 Err(e) => {
551 failed_tx.send_if_modified(|f| {
552 let first = f.is_none();
553 if first {
554 *f = Some(e);
555 }
556 first
557 });
558 return;
559 }
560 }
561 }
562 });
563 }
564 while workers.join_next().await.is_some() {}
565 if let Some(e) = failed.borrow().clone() {
566 return Err(e);
567 }
568 let parts = std::mem::take(&mut *parts.lock().unwrap_or_else(|p| p.into_inner()));
569 let mut whole = Vec::with_capacity(m.size as usize);
570 for part in parts {
571 whole.extend(part.ok_or_else(|| PoolError::ContentReply("a chunk never arrived".into()))?);
572 }
573 manifest::verify(&m, &whole)
574 .map_err(|e| PoolError::ContentMismatch(format!("the whole: {e}")))?;
575 Ok(whole)
576}
577
578async fn fetch_one(
581 link: &Link,
582 realm: &[u8; 32],
583 s: &Sharing,
584 mcid: &Mcid,
585 want: &str,
586 timeout: Duration,
587) -> Result<(String, Value), PoolError> {
588 let asked = async {
589 let stream = link
590 .open_stream(station_link::StreamCall {
591 realm: *realm,
592 procedure: s.procedure.clone(),
593 target: s.node,
594 mode: StreamMode::ServerStream,
595 payload: Value::Map(vec![
596 (Value::text("mcid"), Value::Bytes(mcid.to_vec())),
597 (Value::text("want"), Value::text(want)),
598 ]),
599 deadline: timeout,
600 ..station_link::StreamCall::default()
601 })
602 .await?;
603 let event = stream.recv().await;
604 let _ = stream.close().await;
605 read_body(event, mcid, want)
606 };
607 tokio::time::timeout(timeout, asked)
608 .await
609 .unwrap_or(Err(PoolError::Link(LinkError::CallTimeout)))
610}
611
612fn read_body(
613 event: Result<StreamEvent, LinkError>,
614 mcid: &Mcid,
615 want: &str,
616) -> Result<(String, Value), PoolError> {
617 let body = match event {
618 Err(LinkError::Stream { code, .. }) if code == "not_shared" => {
619 return Err(PoolError::NotShared)
620 }
621 Err(LinkError::EndOfStream) => {
622 return Err(PoolError::ContentReply(
623 "the stream ended with no body".into(),
624 ))
625 }
626 Err(e) => return Err(e.into()),
627 Ok(StreamEvent::Data { body, .. }) => body,
628 Ok(_) => return Err(PoolError::ContentReply("a frame that is not DATA".into())),
629 };
630 let kind = wire_text(&body, "kind").unwrap_or_default();
631 if wire_bytes(&body, "mcid").as_deref() != Some(mcid.as_slice()) {
632 return Err(PoolError::ContentReply(
633 "a body for another content id".into(),
634 ));
635 }
636 let bytes = wire_bytes(&body, "bytes");
637 match (kind.as_str(), bytes) {
638 ("block", Some(b)) if b.len() as u64 > MAX_BLOCK_BYTES => Err(PoolError::ContentTooLarge(
639 format!("a block of {} bytes", b.len()),
640 )),
641 ("block", Some(_)) => Ok((kind, body)),
642 ("manifest", _) if want == "root" => Ok((kind, body)),
643 _ => Err(PoolError::ContentReply(format!("kind {kind:?}"))),
644 }
645}
646
647fn wire_text(v: &Value, name: &str) -> Option<String> {
649 match v.get(name) {
650 Some(Value::Text(t)) => Some(t.clone()),
651 Some(Value::Bytes(b)) => String::from_utf8(b.clone()).ok(),
652 _ => None,
653 }
654}
655
656fn wire_bytes(v: &Value, name: &str) -> Option<Vec<u8>> {
658 match v.get(name) {
659 Some(Value::Bytes(b)) => Some(b.clone()),
660 _ => None,
661 }
662}