1use std::collections::HashSet;
17use std::future::Future;
18use std::sync::Arc;
19use std::time::Duration;
20
21use tokio::time::Instant;
22
23use crate::cbor::Value;
24use crate::frame::StreamMode;
25use crate::record::{self, RecordType, Trust, Verified};
26use crate::seal::KEY_ID_SIZE;
27use crate::station_link::{
28 self, Confidentiality, ConfidentialityError, ConfidentialityReason, Link, LinkError, Report,
29 Seal, Stream, DEFAULT_CALL_TIMEOUT,
30};
31use crate::transport::Target;
32
33use super::{Pool, PoolError, PoolInner};
34
35const MIN_CANDIDATE_SHARE: Duration = Duration::from_secs(1);
37
38#[derive(Debug, Clone, PartialEq)]
43pub struct Call {
44 pub realm: [u8; 32],
45 pub procedure: String,
46 pub provider: [u8; 32],
47 pub payload: Value,
48 pub timeout: Duration,
49 pub token: Option<Vec<u8>>,
50 pub proofs: Vec<Vec<u8>>,
51 pub confidential: Confidentiality,
52}
53
54impl Default for Call {
55 fn default() -> Self {
56 Call {
57 realm: [0; 32],
58 procedure: String::new(),
59 provider: [0; 32],
60 payload: Value::Map(Vec::new()),
61 timeout: Duration::ZERO,
62 token: None,
63 proofs: Vec::new(),
64 confidential: Confidentiality::Preferred,
65 }
66 }
67}
68
69#[derive(Debug, Clone, PartialEq)]
74pub struct StreamCall {
75 pub realm: [u8; 32],
76 pub procedure: String,
77 pub provider: [u8; 32],
78 pub mode: StreamMode,
79 pub payload: Value,
80 pub deadline: Duration,
81 pub token: Option<Vec<u8>>,
82 pub proofs: Vec<Vec<u8>>,
83 pub confidential: Confidentiality,
84}
85
86impl Default for StreamCall {
87 fn default() -> Self {
88 StreamCall {
89 realm: [0; 32],
90 procedure: String::new(),
91 provider: [0; 32],
92 mode: StreamMode::ServerStream,
93 payload: Value::Map(Vec::new()),
94 deadline: Duration::ZERO,
95 token: None,
96 proofs: Vec::new(),
97 confidential: Confidentiality::Preferred,
98 }
99 }
100}
101
102#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
104pub struct Provider {
105 pub node: [u8; 32],
106 pub station: [u8; 32],
107}
108
109#[derive(Debug, Clone, PartialEq, Eq)]
112pub(super) struct Candidate {
113 provider: Provider,
114 expires_at: u64,
115 created_at: u64,
116 kem_key: Option<Vec<u8>>,
117}
118
119impl Candidate {
120 fn seal(&self) -> Seal {
123 match &self.kem_key {
124 Some(key) => Seal::To(key.clone()),
125 None => Seal::Clear,
126 }
127 }
128
129 fn kem_key_id(&self) -> Option<[u8; KEY_ID_SIZE]> {
130 self.kem_key.as_deref().map(crate::seal::key_id)
131 }
132}
133
134#[derive(Debug, Clone, PartialEq, Eq, Hash)]
135pub(super) struct ResolvedKey {
136 realm: [u8; 32],
137 procedure: String,
138 provider: [u8; 32],
139}
140
141impl Pool {
142 pub async fn call(&self, c: Call) -> Result<Value, PoolError> {
147 self.call_report(c).await.map(|(result, _)| result)
148 }
149
150 pub async fn call_report(&self, c: Call) -> Result<(Value, Report), PoolError> {
157 let inner = &self.inner;
158 let realm_key = inner.realm_key_for(&c.realm, &c.procedure)?;
159 let timeout = if c.timeout.is_zero() {
160 DEFAULT_CALL_TIMEOUT
161 } else {
162 c.timeout
163 };
164 let deadline = Instant::now() + timeout;
165 let key = ResolvedKey {
166 realm: c.realm,
167 procedure: c.procedure.clone(),
168 provider: c.provider,
169 };
170 let candidates = bounded(deadline, inner.candidates(&key, realm_key)).await?;
171 let candidates = callable(candidates, c.confidential)?;
172 let mut tried = Vec::new();
173 let count = candidates.len();
174 for (i, cand) in candidates.into_iter().enumerate() {
175 match inner
176 .reach(&cand, candidate_share(deadline, count - i))
177 .await
178 {
179 Ok(link) => {
180 let outcome =
181 bounded(deadline, inner.call_at(&link, &cand, &c, deadline)).await;
182 return inner.settled(key, cand, outcome);
183 }
184 Err(PoolError::Closed) => return Err(PoolError::Closed),
185 Err(e) => {
186 inner.forget(&key);
187 tried.push((cand.provider, e));
188 if Instant::now() >= deadline {
189 break;
190 }
191 }
192 }
193 }
194 Err(PoolError::NoProvider(tried))
195 }
196
197 pub async fn providers(
201 &self,
202 realm: &[u8; 32],
203 procedure: &str,
204 ) -> Result<Vec<Provider>, PoolError> {
205 let realm_key = self.inner.realm_key_for(realm, procedure)?;
206 let key = ResolvedKey {
207 realm: *realm,
208 procedure: procedure.to_string(),
209 provider: [0; 32],
210 };
211 let found = self.inner.resolve(&key, realm_key).await?;
212 Ok(found.into_iter().map(|c| c.provider).collect())
213 }
214
215 pub async fn open_stream(&self, c: StreamCall) -> Result<Stream, PoolError> {
221 let inner = &self.inner;
222 let realm_key = inner.realm_key_for(&c.realm, &c.procedure)?;
223 let deadline = Instant::now() + DEFAULT_CALL_TIMEOUT;
224 let key = ResolvedKey {
225 realm: c.realm,
226 procedure: c.procedure.clone(),
227 provider: c.provider,
228 };
229 let candidates = bounded(deadline, inner.candidates(&key, realm_key)).await?;
230 let candidates = callable(candidates, c.confidential)?;
231 let mut tried = Vec::new();
232 let count = candidates.len();
233 for (i, cand) in candidates.into_iter().enumerate() {
234 match inner
235 .reach(&cand, candidate_share(deadline, count - i))
236 .await
237 {
238 Ok(link) => {
239 let outcome = bounded(deadline, inner.open_at(&link, &cand, &c)).await;
240 return inner.settled(key, cand, outcome);
241 }
242 Err(PoolError::Closed) => return Err(PoolError::Closed),
243 Err(e) => {
244 inner.forget(&key);
245 tried.push((cand.provider, e));
246 if Instant::now() >= deadline {
247 break;
248 }
249 }
250 }
251 }
252 Err(PoolError::NoProvider(tried))
253 }
254
255 pub async fn station_target(&self, station: &[u8; 32]) -> Result<Target, PoolError> {
258 self.inner.station_target(station).await
259 }
260
261 pub async fn link_to(&self, station: &[u8; 32]) -> Result<Link, PoolError> {
266 let deadline = Instant::now() + DEFAULT_CALL_TIMEOUT;
267 self.inner.link_to(station, deadline).await
268 }
269
270 pub async fn find_record(&self, key: &[u8; 32]) -> Result<Verified, PoolError> {
272 self.inner
273 .first_answer(|l| async move { l.find_record(key).await })
274 .await
275 }
276
277 pub async fn find_records(&self, key: &[u8; 32]) -> Result<(Vec<Verified>, usize), PoolError> {
280 self.inner
281 .first_answer(|l| async move { l.find_records(key).await })
282 .await
283 }
284
285 pub async fn find_records_by_type(
288 &self,
289 t: RecordType,
290 ) -> Result<(Vec<Verified>, usize), PoolError> {
291 self.inner
292 .first_answer(|l| async move { l.find_records_by_type(t).await })
293 .await
294 }
295
296 pub async fn put_record(&self, wire: &[u8]) -> Result<(), PoolError> {
298 self.inner
299 .first_answer(|l| async move { l.put_record(wire).await })
300 .await
301 }
302}
303
304impl PoolInner {
305 fn remember(&self, key: ResolvedKey, cand: Candidate) {
306 self.lock().remember.insert(key, cand);
307 }
308
309 fn forget(&self, key: &ResolvedKey) {
310 self.lock().remember.remove(key);
311 }
312
313 async fn candidates(
316 &self,
317 key: &ResolvedKey,
318 realm_key: Option<Vec<u8>>,
319 ) -> Result<Vec<Candidate>, PoolError> {
320 let remembered = self.lock().remember.get(key).cloned();
321 if let Some(cand) = remembered {
322 if cand.expires_at as i64 > now_ms() && self.linked_to(&cand.provider.station).is_some()
323 {
324 return Ok(vec![cand]);
325 }
326 }
327 self.resolve(key, realm_key).await
328 }
329
330 async fn resolve(
333 &self,
334 key: &ResolvedKey,
335 realm_key: Option<Vec<u8>>,
336 ) -> Result<Vec<Candidate>, PoolError> {
337 let slot = record::procedure_key(&key.realm, &key.procedure);
338 let (found, _) = self
339 .first_answer(|l| async move { l.find_records(&slot).await })
340 .await?;
341 let now = now_ms();
342 let trust = Trust {
343 profile: self.opts.identity.profile(),
344 realm_key,
345 };
346 let mut out: Vec<Candidate> = found
347 .iter()
348 .filter(|v| v.record().record_type == RecordType::PROCEDURE_ADVERTISEMENT)
349 .filter_map(|v| {
350 let ad = record::read_procedure_advertisement(v.record()).ok()?;
351 let wanted = ad.realm_id == key.realm
352 && ad.procedure == key.procedure
353 && (key.provider == [0; 32] || ad.advertiser_node == key.provider);
354 if !wanted || record::verify_authorization(v, &trust, now).is_err() {
355 return None;
356 }
357 Some(Candidate {
358 provider: Provider {
359 node: ad.advertiser_node,
360 station: ad.serving_station,
361 },
362 expires_at: v.record().expires_at,
363 created_at: v.record().created_at,
364 kem_key: ad.kem_key.map(|(key, _)| key),
365 })
366 })
367 .collect();
368 if out.is_empty() {
369 return Err(PoolError::NoProvider(Vec::new()));
370 }
371 out.sort_by_key(|c| std::cmp::Reverse(c.created_at));
372 Ok(out)
373 }
374
375 async fn reach(self: &Arc<Self>, cand: &Candidate, share: Instant) -> Result<Link, PoolError> {
378 bounded(share, self.link_to(&cand.provider.station, share)).await
379 }
380
381 fn settled<T>(
385 &self,
386 key: ResolvedKey,
387 cand: Candidate,
388 outcome: Result<T, PoolError>,
389 ) -> Result<T, PoolError> {
390 match &outcome {
391 Ok(_) | Err(PoolError::Link(LinkError::Provider { .. })) => self.remember(key, cand),
392 Err(_) => self.forget(&key),
393 }
394 outcome
395 }
396
397 async fn call_at(
402 &self,
403 link: &Link,
404 cand: &Candidate,
405 c: &Call,
406 deadline: Instant,
407 ) -> Result<(Value, Report), PoolError> {
408 let left = deadline.saturating_duration_since(Instant::now());
409 let result = link
410 .call(station_link::Call {
411 realm: c.realm,
412 procedure: c.procedure.clone(),
413 target: cand.provider.node,
414 payload: c.payload.clone(),
415 timeout: left.max(Duration::from_millis(1)),
416 token: c.token.clone(),
417 proofs: c.proofs.clone(),
418 seal: Some(cand.seal()),
419 })
420 .await?;
421 Ok((result, Report::of(cand.provider.node, cand.kem_key_id())))
422 }
423
424 async fn open_at(
426 &self,
427 link: &Link,
428 cand: &Candidate,
429 c: &StreamCall,
430 ) -> Result<Stream, PoolError> {
431 Ok(link
432 .open_stream(station_link::StreamCall {
433 realm: c.realm,
434 procedure: c.procedure.clone(),
435 target: cand.provider.node,
436 mode: c.mode,
437 payload: c.payload.clone(),
438 deadline: c.deadline,
439 token: c.token.clone(),
440 proofs: c.proofs.clone(),
441 seal: Some(cand.seal()),
442 })
443 .await?)
444 }
445
446 fn linked_to(&self, station: &[u8; 32]) -> Option<Link> {
448 self.links()
449 .into_iter()
450 .find(|l| l.station_node_id() == *station)
451 }
452
453 pub(super) async fn link_to(
454 self: &Arc<Self>,
455 station: &[u8; 32],
456 deadline: Instant,
457 ) -> Result<Link, PoolError> {
458 if let Some(link) = self.linked_to(station) {
459 return Ok(link);
460 }
461 let (existing, direct) = {
462 let state = self.lock();
463 if state.closed {
464 return Err(PoolError::Closed);
465 }
466 let existing = state
467 .members
468 .iter()
469 .find(|m| m.target.expected_node_id == *station)
470 .cloned();
471 (existing, state.members.iter().filter(|m| m.direct).count())
472 };
473 let fresh = existing.is_none();
474 let member = match existing {
475 Some(m) => m,
476 None => {
477 if direct >= self.opts.max_direct_links {
478 return Err(PoolError::DirectLinksFull);
479 }
480 let target = bounded(deadline, self.station_target(station)).await?;
481 self.start_member(target, true)
482 }
483 };
484 if let Some(link) = member.await_up(deadline, fresh).await {
485 return Ok(link);
486 }
487 if fresh {
488 self.drop_member(&member);
489 }
490 Err(PoolError::StationNotReached {
491 station: *station,
492 cause: member.last_error(),
493 })
494 }
495
496 async fn station_target(&self, station: &[u8; 32]) -> Result<Target, PoolError> {
497 let slot = record::station_endpoint_key(station);
498 let verified = self
499 .first_answer(|l| async move { l.find_record(&slot).await })
500 .await
501 .map_err(|e| match e {
502 PoolError::Link(link) => PoolError::NoStationEndpoint(Some(link)),
503 other => other,
504 })?;
505 let r = verified.record();
506 let signer = r.signed.as_ref().map(|s| s.key_id);
507 let endpoint = record::read_station_endpoint(r)
508 .map_err(|e| PoolError::NoStationEndpoint(Some(e.into())))?;
509 match (signer, endpoint.host_advertised.first()) {
510 (Some(signer), Some(host)) if signer == *station && endpoint.quic_port != 0 => {
511 Ok(Target {
512 host: host.clone(),
513 port: endpoint.quic_port,
514 profile: self.opts.identity.profile(),
515 expected_node_id: *station,
516 })
517 }
518 _ => Err(PoolError::NoStationEndpoint(None)),
519 }
520 }
521
522 pub(super) async fn first_answer<'a, T, F, Fut>(&self, ask: F) -> Result<T, PoolError>
526 where
527 F: Fn(Link) -> Fut,
528 Fut: Future<Output = Result<T, LinkError>> + 'a,
529 {
530 let links = self.links();
531 if links.is_empty() {
532 return Err(PoolError::NoLink(Vec::new()));
533 }
534 let mut errors = Vec::new();
535 for link in links {
536 match ask(link).await {
537 Err(e) if unreachable(&e) => errors.push(e),
538 answered => return answered.map_err(PoolError::Link),
539 }
540 }
541 Err(PoolError::NoLink(errors))
542 }
543}
544
545fn unreachable(e: &LinkError) -> bool {
548 matches!(
549 e,
550 LinkError::CallTimeout
551 | LinkError::Closed
552 | LinkError::LivenessLost
553 | LinkError::V5DowngradeRefused
554 | LinkError::Io(_)
555 | LinkError::Goodbye(_)
556 | LinkError::StatusExpired
557 | LinkError::BindingExpired
558 )
559}
560
561fn candidate_share(deadline: Instant, left: usize) -> Instant {
565 let now = Instant::now();
566 let remaining = deadline.saturating_duration_since(now);
567 (now + (remaining / left.max(1) as u32).max(MIN_CANDIDATE_SHARE)).min(deadline)
568}
569
570async fn bounded<T>(
572 deadline: Instant,
573 work: impl Future<Output = Result<T, PoolError>>,
574) -> Result<T, PoolError> {
575 tokio::time::timeout_at(deadline, work)
576 .await
577 .unwrap_or(Err(PoolError::Link(LinkError::CallTimeout)))
578}
579
580fn now_ms() -> i64 {
581 crate::uuid_v7::now_ms() as i64
582}
583
584fn callable(
594 candidates: Vec<Candidate>,
595 confidential: Confidentiality,
596) -> Result<Vec<Candidate>, PoolError> {
597 if confidential == Confidentiality::Off {
598 return Err(PoolError::InvalidOpts(
599 "confidential off is refused for a pool call or open: it is preferred or required"
600 .into(),
601 ));
602 }
603 let advertised: Vec<[u8; KEY_ID_SIZE]> = candidates
604 .iter()
605 .filter_map(Candidate::kem_key_id)
606 .collect();
607 let keyed: HashSet<[u8; 32]> = candidates
608 .iter()
609 .filter(|c| c.kem_key.is_some())
610 .map(|c| c.provider.node)
611 .collect();
612 let kept: Vec<Candidate> = candidates
613 .into_iter()
614 .filter(|c| match (confidential, &c.kem_key) {
615 (_, Some(_)) => true,
616 (Confidentiality::Preferred, None) => !keyed.contains(&c.provider.node),
617 (_, None) => false,
618 })
619 .collect();
620 if kept.is_empty() {
621 return Err(PoolError::Confidentiality(ConfidentialityError {
622 reason: ConfidentialityReason::NoKemKey,
623 advertised,
624 named: None,
625 }));
626 }
627 Ok(kept)
628}
629
630#[cfg(test)]
631mod tests {
632 use super::*;
633
634 fn cand(node: u8, kem_key: Option<u8>) -> Candidate {
635 Candidate {
636 provider: Provider {
637 node: [node; 32],
638 station: [9; 32],
639 },
640 expires_at: 0,
641 created_at: 0,
642 kem_key: kem_key.map(|k| vec![k; 1568]),
643 }
644 }
645
646 fn refused(r: Result<Vec<Candidate>, PoolError>) -> ConfidentialityError {
647 match r {
648 Err(PoolError::Confidentiality(e)) => e,
649 other => panic!("not a confidentiality refusal: {other:?}"),
650 }
651 }
652
653 #[test]
654 fn preferred_seals_to_a_keyed_provider_and_calls_a_keyless_one_in_the_clear() {
655 let kept = callable(
656 vec![cand(1, Some(7)), cand(2, None), cand(3, None)],
657 Confidentiality::Preferred,
658 )
659 .unwrap();
660 assert_eq!(kept, vec![cand(1, Some(7)), cand(2, None), cand(3, None)]);
661 assert_eq!(kept[0].seal(), Seal::To(vec![7; 1568]));
662 assert_eq!(kept[1].seal(), Seal::Clear);
663 }
664
665 #[test]
666 fn a_provider_that_names_a_key_is_not_called_through_its_older_keyless_ad() {
667 let kept = callable(
668 vec![cand(1, Some(7)), cand(1, None), cand(2, None)],
669 Confidentiality::Preferred,
670 )
671 .unwrap();
672 assert_eq!(kept, vec![cand(1, Some(7)), cand(2, None)]);
673 }
674
675 #[test]
676 fn required_calls_only_keyed_providers_and_refuses_when_there_are_none() {
677 let kept = callable(
678 vec![cand(1, None), cand(2, Some(7))],
679 Confidentiality::Required,
680 )
681 .unwrap();
682 assert_eq!(kept, vec![cand(2, Some(7))]);
683 let e = refused(callable(vec![cand(1, None)], Confidentiality::Required));
684 assert_eq!(e.reason, ConfidentialityReason::NoKemKey);
685 assert!(e.advertised.is_empty());
686 assert_eq!(e.to_string(), "confidentiality: no_kem_key");
687 }
688
689 #[test]
690 fn confidential_parses_and_off_is_refused_for_a_pool_call() {
691 assert_eq!(Confidentiality::default(), Confidentiality::Preferred);
692 for (text, parsed) in [
693 ("", Confidentiality::Preferred),
694 ("preferred", Confidentiality::Preferred),
695 ("required", Confidentiality::Required),
696 ("off", Confidentiality::Off),
697 ] {
698 assert_eq!(text.parse::<Confidentiality>(), Ok(parsed), "{text}");
699 }
700 for refused in ["Required", "none", "optional"] {
701 assert!(refused.parse::<Confidentiality>().is_err(), "{refused}");
702 }
703 assert!(matches!(
704 callable(vec![cand(1, Some(7))], Confidentiality::Off),
705 Err(PoolError::InvalidOpts(_))
706 ));
707 }
708}