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| answer_fetch_while_pooled(pool.clone(), r, s));
253 let procedure = record::own_procedure(&self.self_id, CONTENT_PROCEDURE);
254 let served = Pool {
255 inner: self.clone(),
256 }
257 .serve(Offer::stream(
258 *realm,
259 &procedure,
260 StreamMode::ServerStream,
261 answer,
262 ))
263 .await?;
264 self.content.lock().insert(
265 *realm,
266 SharedRealm {
267 _served: served,
268 roots: HashMap::new(),
269 chunks: HashMap::new(),
270 announcements: HashMap::new(),
271 },
272 );
273 Ok(())
274 }
275
276 async fn answer_fetch(&self, realm: &[u8; 32], s: &Stream) -> Result<(), LinkError> {
279 let args = &s.request().payload;
280 let mcid: Option<Mcid> = match wire_bytes(args, "mcid") {
281 Some(b) if b.len() == 50 && b[0] == 2 => b.as_slice().try_into().ok(),
282 _ => None,
283 };
284 let want = wire_text(args, "want");
285 let (Some(mcid), Some(want @ ("root" | "block"))) = (mcid, want.as_deref()) else {
286 return s
287 .abort(
288 "malformed",
289 "a fetch names one content id and wants root or block",
290 )
291 .await;
292 };
293 let body = self
294 .content
295 .lock()
296 .get(realm)
297 .and_then(|shared| shared.body(want, &mcid));
298 match body {
299 None => {
300 s.abort("not_shared", "this node does not share that content")
301 .await
302 }
303 Some(body) => {
304 s.send_value(body).await?;
305 s.close().await
306 }
307 }
308 }
309
310 fn forget_shared(&self, realm: &[u8; 32], mcid: &Mcid) -> Option<Announced> {
313 let mut realms = self.content.lock();
314 let shared = realms.get_mut(realm)?;
315 if let Some(Root::Chunked(m)) = shared.roots.remove(mcid) {
316 forget_chunks(&mut shared.chunks, &m);
317 }
318 shared.announcements.remove(mcid)
319 }
320
321 async fn announce_content(
324 self: &Arc<Self>,
325 realm: &[u8; 32],
326 mcid: &Mcid,
327 root: &Root,
328 ) -> Result<Record, PoolError> {
329 let station = self
330 .links()
331 .first()
332 .map(Link::station_node_id)
333 .ok_or(PoolError::NoLink(Vec::new()))?;
334 let mut opts = ContentAnnouncementOptions {
335 realm_id: *realm,
336 serving_station: station,
337 procedure: record::own_procedure(&self.self_id, CONTENT_PROCEDURE),
338 ttl_ms: ANNOUNCEMENT_TTL.as_millis() as u64,
339 ..ContentAnnouncementOptions::default()
340 };
341 match root {
342 Root::Chunked(m) => {
343 opts.name = String::from_utf8_lossy(&m.name).into_owned();
344 opts.size = Some(m.size);
345 opts.chunk_count = Some(m.chunk_count);
346 }
347 Root::Block(b) => opts.size = Some(b.len() as u64),
348 }
349 let unsigned =
350 new_content_announcement(&self.self_id, mcid, &opts).map_err(LinkError::from)?;
351 let signed = record::sign(&unsigned, &self.opts.identity).map_err(LinkError::from)?;
352 let wire = record::encode(&signed).map_err(LinkError::from)?;
353 Pool {
354 inner: self.clone(),
355 }
356 .put_record(&wire)
357 .await?;
358 Ok(signed)
359 }
360
361 async fn content_sharers(
364 &self,
365 realm: &[u8; 32],
366 mcid: &Mcid,
367 ) -> Result<Vec<Sharing>, PoolError> {
368 let key = record::content_key(mcid).map_err(LinkError::from)?;
369 let found = match self
370 .first_answer(|l| async move { l.find_records(&key).await })
371 .await
372 {
373 Ok((found, _)) => found,
374 Err(PoolError::Link(LinkError::RecordNotFound)) => Vec::new(),
375 Err(e) => return Err(e),
376 };
377 let mut out: Vec<Sharing> = found
378 .iter()
379 .filter(|v| v.record().record_type == RecordType::CONTENT_ANNOUNCEMENT)
380 .filter_map(|v| record::read_content_announcement(v.record()).ok())
381 .filter(|a| {
382 a.mcid.as_slice() == mcid.as_slice()
383 && a.realm_id == *realm
384 && a.serving_station != [0; 32]
385 && content_procedure_bound(&a.procedure, &a.announcer_node)
386 })
387 .map(|a| Sharing {
388 node: a.announcer_node,
389 station: a.serving_station,
390 procedure: a.procedure,
391 })
392 .collect();
393 shuffle(&mut out);
394 Ok(out)
395 }
396
397 async fn fetch_from(
399 self: &Arc<Self>,
400 realm: &[u8; 32],
401 s: &Sharing,
402 mcid: &Mcid,
403 opts: &ContentOptions,
404 ) -> Result<Vec<u8>, PoolError> {
405 let link = self
406 .link_to(&s.station, Instant::now() + DEFAULT_CALL_TIMEOUT)
407 .await?;
408 let (kind, body) = fetch_one(&link, realm, s, mcid, "root", opts.chunk_timeout).await?;
409 if kind == "block" {
410 return verified_block(&body, mcid)
411 .ok_or_else(|| PoolError::ContentMismatch("the block".into()));
412 }
413 let m = manifest::from_wire(body.get("manifest").unwrap_or(&Value::Null))
414 .map_err(|e| PoolError::ContentReply(format!("the manifest: {e}")))?;
415 manifest::verify_mcid(&m, mcid)
416 .map_err(|e| PoolError::ContentMismatch(format!("the manifest: {e}")))?;
417 if m.size > opts.max_bytes || m.chunk_count > opts.max_chunks {
418 return Err(PoolError::ContentTooLarge(format!(
419 "{} bytes in {} chunks",
420 m.size, m.chunk_count
421 )));
422 }
423 manifest::check_whole(&m)
424 .map_err(|e| PoolError::ContentMismatch(format!("the manifest: {e}")))?;
425 fetch_chunks(link, *realm, s.clone(), Arc::new(m), opts).await
426 }
427}
428
429impl SharedRealm {
430 fn body(&self, want: &str, mcid: &Mcid) -> Option<Value> {
433 let block = |b: &[u8]| {
434 Value::Map(vec![
435 (Value::text("kind"), Value::text("block")),
436 (Value::text("mcid"), Value::Bytes(mcid.to_vec())),
437 (Value::text("bytes"), Value::Bytes(b.to_vec())),
438 ])
439 };
440 if want == "block" {
441 return self.chunks.get(mcid).map(|b| block(b));
442 }
443 match self.roots.get(mcid)? {
444 Root::Block(b) => Some(block(b)),
445 Root::Chunked(m) => Some(Value::Map(vec![
446 (Value::text("kind"), Value::text("manifest")),
447 (Value::text("mcid"), Value::Bytes(mcid.to_vec())),
448 (Value::text("manifest"), manifest::to_wire(m)),
449 ])),
450 }
451 }
452}
453
454async fn answer_fetch_while_pooled(
457 pool: Weak<PoolInner>,
458 realm: [u8; 32],
459 s: Stream,
460) -> Result<(), String> {
461 match pool.upgrade() {
462 Some(inner) => inner.answer_fetch(&realm, &s).await,
463 None => {
464 s.abort("not_shared", "the node no longer shares content")
465 .await
466 }
467 }
468 .map_err(|e| e.to_string())
469}
470
471fn forget_chunks(chunks: &mut HashMap<Mcid, Vec<u8>>, m: &Manifest) {
473 for i in 0..m.chunks.len() {
474 if let Some(c) = chunk_mcid(m, i) {
475 chunks.remove(&c);
476 }
477 }
478}
479
480fn verified_block(body: &Value, mcid: &Mcid) -> Option<Vec<u8>> {
482 let bytes = wire_bytes(body, "bytes").unwrap_or_default();
483 (block_mcid(&bytes) == *mcid).then_some(bytes)
484}
485
486async fn renew_announcement(
490 pool: Weak<PoolInner>,
491 realm: [u8; 32],
492 mcid: Mcid,
493 root: Root,
494 latest: Arc<Mutex<Record>>,
495 mut stopped: watch::Receiver<bool>,
496) {
497 loop {
498 tokio::select! {
499 _ = stopped.wait_for(|s| *s) => return,
500 _ = tokio::time::sleep(ANNOUNCEMENT_TTL / 2) => {}
501 }
502 let Some(inner) = pool.upgrade() else { return };
503 if inner.lock().closed {
504 return;
505 }
506 let renewed = tokio::time::timeout(
507 DEFAULT_CALL_TIMEOUT,
508 inner.announce_content(&realm, &mcid, &root),
509 )
510 .await;
511 if let Ok(Ok(record)) = renewed {
512 *latest.lock().unwrap_or_else(|p| p.into_inner()) = record;
513 }
514 }
515}
516
517async fn fetch_chunks(
521 link: Link,
522 realm: [u8; 32],
523 s: Sharing,
524 m: Arc<Manifest>,
525 opts: &ContentOptions,
526) -> Result<Vec<u8>, PoolError> {
527 let count = m.chunks.len();
528 let parts: Arc<Mutex<Vec<Option<Vec<u8>>>>> = Arc::new(Mutex::new(vec![None; count]));
529 let next = Arc::new(AtomicUsize::new(0));
530 let (failed_tx, failed) = watch::channel::<Option<PoolError>>(None);
531 let failed_tx = Arc::new(failed_tx);
532 let mut workers = tokio::task::JoinSet::new();
533 for _ in 0..opts.parallel.max(1).min(count) {
534 workers.spawn(fetch_chunks_in_turn(ChunkWorker {
535 link: link.clone(),
536 realm,
537 s: s.clone(),
538 m: m.clone(),
539 parts: parts.clone(),
540 next: next.clone(),
541 failed_tx: failed_tx.clone(),
542 timeout: opts.chunk_timeout,
543 }));
544 }
545 while workers.join_next().await.is_some() {}
546 if let Some(e) = failed.borrow().clone() {
547 return Err(e);
548 }
549 let parts = std::mem::take(&mut *parts.lock().unwrap_or_else(|p| p.into_inner()));
550 let mut whole = Vec::with_capacity(m.size as usize);
551 for part in parts {
552 whole.extend(part.ok_or_else(|| PoolError::ContentReply("a chunk never arrived".into()))?);
553 }
554 manifest::verify(&m, &whole)
555 .map_err(|e| PoolError::ContentMismatch(format!("the whole: {e}")))?;
556 Ok(whole)
557}
558
559struct ChunkWorker {
564 link: Link,
565 realm: [u8; 32],
566 s: Sharing,
567 m: Arc<Manifest>,
568 parts: Arc<Mutex<Vec<Option<Vec<u8>>>>>,
569 next: Arc<AtomicUsize>,
570 failed_tx: Arc<watch::Sender<Option<PoolError>>>,
571 timeout: Duration,
572}
573
574async fn fetch_chunks_in_turn(w: ChunkWorker) {
578 loop {
579 if w.failed_tx.borrow().is_some() {
580 return;
581 }
582 let i = w.next.fetch_add(1, Ordering::SeqCst);
583 let Some(want) = chunk_mcid(&w.m, i) else {
584 return;
585 };
586 let fetched = fetch_one(&w.link, &w.realm, &w.s, &want, "block", w.timeout).await;
587 let outcome = fetched.and_then(|(_, body)| {
588 verified_block(&body, &want)
589 .ok_or_else(|| PoolError::ContentMismatch(format!("chunk {i}")))
590 });
591 let bytes = match outcome {
592 Ok(bytes) => bytes,
593 Err(e) => return keep_first_failure(&w.failed_tx, e),
594 };
595 w.parts.lock().unwrap_or_else(|p| p.into_inner())[i] = Some(bytes);
596 }
597}
598
599fn keep_first_failure(failed_tx: &watch::Sender<Option<PoolError>>, e: PoolError) {
601 failed_tx.send_if_modified(|f| {
602 let first = f.is_none();
603 if first {
604 *f = Some(e);
605 }
606 first
607 });
608}
609
610async fn fetch_one(
613 link: &Link,
614 realm: &[u8; 32],
615 s: &Sharing,
616 mcid: &Mcid,
617 want: &str,
618 timeout: Duration,
619) -> Result<(String, Value), PoolError> {
620 let asked = async {
621 let stream = link
622 .open_stream(station_link::StreamCall {
623 realm: *realm,
624 procedure: s.procedure.clone(),
625 target: s.node,
626 mode: StreamMode::ServerStream,
627 payload: Value::Map(vec![
628 (Value::text("mcid"), Value::Bytes(mcid.to_vec())),
629 (Value::text("want"), Value::text(want)),
630 ]),
631 deadline: timeout,
632 seal: Some(station_link::Seal::Clear),
636 ..station_link::StreamCall::default()
637 })
638 .await?;
639 let event = stream.recv().await;
640 let _ = stream.close().await;
641 read_body(event, mcid, want)
642 };
643 tokio::time::timeout(timeout, asked)
644 .await
645 .unwrap_or(Err(PoolError::Link(LinkError::CallTimeout)))
646}
647
648fn read_body(
649 event: Result<StreamEvent, LinkError>,
650 mcid: &Mcid,
651 want: &str,
652) -> Result<(String, Value), PoolError> {
653 let body = match event {
654 Err(LinkError::Stream { code, .. }) if code == "not_shared" => {
655 return Err(PoolError::NotShared)
656 }
657 Err(LinkError::EndOfStream) => {
658 return Err(PoolError::ContentReply(
659 "the stream ended with no body".into(),
660 ))
661 }
662 Err(e) => return Err(e.into()),
663 Ok(StreamEvent::Data { body, .. }) => body,
664 Ok(_) => return Err(PoolError::ContentReply("a frame that is not DATA".into())),
665 };
666 let kind = wire_text(&body, "kind").unwrap_or_default();
667 if wire_bytes(&body, "mcid").as_deref() != Some(mcid.as_slice()) {
668 return Err(PoolError::ContentReply(
669 "a body for another content id".into(),
670 ));
671 }
672 let bytes = wire_bytes(&body, "bytes");
673 match (kind.as_str(), bytes) {
674 ("block", Some(b)) if b.len() as u64 > MAX_BLOCK_BYTES => Err(PoolError::ContentTooLarge(
675 format!("a block of {} bytes", b.len()),
676 )),
677 ("block", Some(_)) => Ok((kind, body)),
678 ("manifest", _) if want == "root" => Ok((kind, body)),
679 _ => Err(PoolError::ContentReply(format!("kind {kind:?}"))),
680 }
681}
682
683fn wire_text(v: &Value, name: &str) -> Option<String> {
685 match v.get(name) {
686 Some(Value::Text(t)) => Some(t.clone()),
687 Some(Value::Bytes(b)) => String::from_utf8(b.clone()).ok(),
688 _ => None,
689 }
690}
691
692fn wire_bytes(v: &Value, name: &str) -> Option<Vec<u8>> {
694 match v.get(name) {
695 Some(Value::Bytes(b)) => Some(b.clone()),
696 _ => None,
697 }
698}