use std::collections::HashSet;
use std::fmt;
use std::future::Future;
use std::str::FromStr;
use std::sync::Arc;
use std::time::Duration;
use tokio::time::Instant;
use crate::cbor::Value;
use crate::frame::StreamMode;
use crate::record::{self, RecordType, Trust, Verified};
use crate::seal::KEY_ID_SIZE;
use crate::station_link::{self, Link, LinkError, Stream, DEFAULT_CALL_TIMEOUT};
use crate::transport::Target;
use super::{Pool, PoolError, PoolInner};
const MIN_CANDIDATE_SHARE: Duration = Duration::from_secs(1);
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum Confidentiality {
#[default]
Preferred,
Required,
}
impl FromStr for Confidentiality {
type Err = PoolError;
fn from_str(s: &str) -> Result<Self, PoolError> {
match s {
"" | "preferred" => Ok(Confidentiality::Preferred),
"required" => Ok(Confidentiality::Required),
"off" => Err(PoolError::InvalidOpts(
"confidential off is refused for a call or an open: it is preferred or required"
.into(),
)),
other => Err(PoolError::InvalidOpts(format!(
"confidential is preferred or required, not {other:?}"
))),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ConfidentialityReason {
NoKemKey,
}
impl ConfidentialityReason {
pub fn name(self) -> &'static str {
match self {
ConfidentialityReason::NoKemKey => "no_kem_key",
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ConfidentialityError {
pub reason: ConfidentialityReason,
pub advertised: Vec<[u8; KEY_ID_SIZE]>,
}
impl fmt::Display for ConfidentialityError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "confidentiality: {}", self.reason.name())?;
for id in &self.advertised {
f.write_str(" ")?;
for b in id {
write!(f, "{b:02x}")?;
}
}
Ok(())
}
}
#[derive(Debug, Clone, PartialEq)]
pub struct Call {
pub realm: [u8; 32],
pub procedure: String,
pub provider: [u8; 32],
pub payload: Value,
pub timeout: Duration,
pub token: Option<Vec<u8>>,
pub proofs: Vec<Vec<u8>>,
pub confidential: Confidentiality,
}
impl Default for Call {
fn default() -> Self {
Call {
realm: [0; 32],
procedure: String::new(),
provider: [0; 32],
payload: Value::Map(Vec::new()),
timeout: Duration::ZERO,
token: None,
proofs: Vec::new(),
confidential: Confidentiality::Preferred,
}
}
}
#[derive(Debug, Clone, PartialEq)]
pub struct StreamCall {
pub realm: [u8; 32],
pub procedure: String,
pub provider: [u8; 32],
pub mode: StreamMode,
pub payload: Value,
pub deadline: Duration,
pub token: Option<Vec<u8>>,
pub proofs: Vec<Vec<u8>>,
pub confidential: Confidentiality,
}
impl Default for StreamCall {
fn default() -> Self {
StreamCall {
realm: [0; 32],
procedure: String::new(),
provider: [0; 32],
mode: StreamMode::ServerStream,
payload: Value::Map(Vec::new()),
deadline: Duration::ZERO,
token: None,
proofs: Vec::new(),
confidential: Confidentiality::Preferred,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct Provider {
pub node: [u8; 32],
pub station: [u8; 32],
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) struct Candidate {
provider: Provider,
expires_at: u64,
created_at: u64,
kem_key_id: Option<[u8; KEY_ID_SIZE]>,
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub(super) struct ResolvedKey {
realm: [u8; 32],
procedure: String,
provider: [u8; 32],
}
impl Pool {
pub async fn call(&self, c: Call) -> Result<Value, PoolError> {
let inner = &self.inner;
let realm_key = inner.realm_key_for(&c.realm, &c.procedure)?;
let timeout = if c.timeout.is_zero() {
DEFAULT_CALL_TIMEOUT
} else {
c.timeout
};
let deadline = Instant::now() + timeout;
let key = ResolvedKey {
realm: c.realm,
procedure: c.procedure.clone(),
provider: c.provider,
};
let candidates = bounded(deadline, inner.candidates(&key, realm_key)).await?;
let candidates = callable(candidates, c.confidential)?;
let mut tried = Vec::new();
let count = candidates.len();
for (i, cand) in candidates.into_iter().enumerate() {
match inner
.reach(&cand, candidate_share(deadline, count - i))
.await
{
Ok(link) => {
let outcome =
bounded(deadline, inner.call_at(&link, &cand, &c, deadline)).await;
return inner.settled(key, cand, outcome);
}
Err(PoolError::Closed) => return Err(PoolError::Closed),
Err(e) => {
inner.forget(&key);
tried.push((cand.provider, e));
if Instant::now() >= deadline {
break;
}
}
}
}
Err(PoolError::NoProvider(tried))
}
pub async fn providers(
&self,
realm: &[u8; 32],
procedure: &str,
) -> Result<Vec<Provider>, PoolError> {
let realm_key = self.inner.realm_key_for(realm, procedure)?;
let key = ResolvedKey {
realm: *realm,
procedure: procedure.to_string(),
provider: [0; 32],
};
let found = self.inner.resolve(&key, realm_key).await?;
Ok(found.into_iter().map(|c| c.provider).collect())
}
pub async fn open_stream(&self, c: StreamCall) -> Result<Stream, PoolError> {
let inner = &self.inner;
let realm_key = inner.realm_key_for(&c.realm, &c.procedure)?;
let deadline = Instant::now() + DEFAULT_CALL_TIMEOUT;
let key = ResolvedKey {
realm: c.realm,
procedure: c.procedure.clone(),
provider: c.provider,
};
let candidates = bounded(deadline, inner.candidates(&key, realm_key)).await?;
let candidates = callable(candidates, c.confidential)?;
let mut tried = Vec::new();
let count = candidates.len();
for (i, cand) in candidates.into_iter().enumerate() {
match inner
.reach(&cand, candidate_share(deadline, count - i))
.await
{
Ok(link) => {
let outcome = bounded(deadline, inner.open_at(&link, &cand, &c)).await;
return inner.settled(key, cand, outcome);
}
Err(PoolError::Closed) => return Err(PoolError::Closed),
Err(e) => {
inner.forget(&key);
tried.push((cand.provider, e));
if Instant::now() >= deadline {
break;
}
}
}
}
Err(PoolError::NoProvider(tried))
}
pub async fn station_target(&self, station: &[u8; 32]) -> Result<Target, PoolError> {
self.inner.station_target(station).await
}
pub async fn link_to(&self, station: &[u8; 32]) -> Result<Link, PoolError> {
let deadline = Instant::now() + DEFAULT_CALL_TIMEOUT;
self.inner.link_to(station, deadline).await
}
pub async fn find_record(&self, key: &[u8; 32]) -> Result<Verified, PoolError> {
self.inner
.first_answer(|l| async move { l.find_record(key).await })
.await
}
pub async fn find_records(&self, key: &[u8; 32]) -> Result<(Vec<Verified>, usize), PoolError> {
self.inner
.first_answer(|l| async move { l.find_records(key).await })
.await
}
pub async fn find_records_by_type(
&self,
t: RecordType,
) -> Result<(Vec<Verified>, usize), PoolError> {
self.inner
.first_answer(|l| async move { l.find_records_by_type(t).await })
.await
}
pub async fn put_record(&self, wire: &[u8]) -> Result<(), PoolError> {
self.inner
.first_answer(|l| async move { l.put_record(wire).await })
.await
}
}
impl PoolInner {
fn remember(&self, key: ResolvedKey, cand: Candidate) {
self.lock().remember.insert(key, cand);
}
fn forget(&self, key: &ResolvedKey) {
self.lock().remember.remove(key);
}
async fn candidates(
&self,
key: &ResolvedKey,
realm_key: Option<Vec<u8>>,
) -> Result<Vec<Candidate>, PoolError> {
let remembered = self.lock().remember.get(key).copied();
if let Some(cand) = remembered {
if cand.expires_at as i64 > now_ms() && self.linked_to(&cand.provider.station).is_some()
{
return Ok(vec![cand]);
}
}
self.resolve(key, realm_key).await
}
async fn resolve(
&self,
key: &ResolvedKey,
realm_key: Option<Vec<u8>>,
) -> Result<Vec<Candidate>, PoolError> {
let slot = record::procedure_key(&key.realm, &key.procedure);
let (found, _) = self
.first_answer(|l| async move { l.find_records(&slot).await })
.await?;
let now = now_ms();
let trust = Trust {
profile: self.opts.identity.profile(),
realm_key,
};
let mut out: Vec<Candidate> = found
.iter()
.filter(|v| v.record().record_type == RecordType::PROCEDURE_ADVERTISEMENT)
.filter_map(|v| {
let ad = record::read_procedure_advertisement(v.record()).ok()?;
let wanted = ad.realm_id == key.realm
&& ad.procedure == key.procedure
&& (key.provider == [0; 32] || ad.advertiser_node == key.provider);
if !wanted || record::verify_authorization(v, &trust, now).is_err() {
return None;
}
Some(Candidate {
provider: Provider {
node: ad.advertiser_node,
station: ad.serving_station,
},
expires_at: v.record().expires_at,
created_at: v.record().created_at,
kem_key_id: ad.kem_key.map(|(_, id)| id),
})
})
.collect();
if out.is_empty() {
return Err(PoolError::NoProvider(Vec::new()));
}
out.sort_by_key(|c| std::cmp::Reverse(c.created_at));
Ok(out)
}
async fn reach(self: &Arc<Self>, cand: &Candidate, share: Instant) -> Result<Link, PoolError> {
bounded(share, self.link_to(&cand.provider.station, share)).await
}
fn settled<T>(
&self,
key: ResolvedKey,
cand: Candidate,
outcome: Result<T, PoolError>,
) -> Result<T, PoolError> {
match &outcome {
Ok(_) | Err(PoolError::Link(LinkError::Provider { .. })) => self.remember(key, cand),
Err(_) => self.forget(&key),
}
outcome
}
async fn call_at(
&self,
link: &Link,
cand: &Candidate,
c: &Call,
deadline: Instant,
) -> Result<Value, PoolError> {
let left = deadline.saturating_duration_since(Instant::now());
Ok(link
.call(station_link::Call {
realm: c.realm,
procedure: c.procedure.clone(),
target: cand.provider.node,
payload: c.payload.clone(),
timeout: left.max(Duration::from_millis(1)),
token: c.token.clone(),
proofs: c.proofs.clone(),
})
.await?)
}
async fn open_at(
&self,
link: &Link,
cand: &Candidate,
c: &StreamCall,
) -> Result<Stream, PoolError> {
Ok(link
.open_stream(station_link::StreamCall {
realm: c.realm,
procedure: c.procedure.clone(),
target: cand.provider.node,
mode: c.mode,
payload: c.payload.clone(),
deadline: c.deadline,
token: c.token.clone(),
proofs: c.proofs.clone(),
})
.await?)
}
fn linked_to(&self, station: &[u8; 32]) -> Option<Link> {
self.links()
.into_iter()
.find(|l| l.station_node_id() == *station)
}
pub(super) async fn link_to(
self: &Arc<Self>,
station: &[u8; 32],
deadline: Instant,
) -> Result<Link, PoolError> {
if let Some(link) = self.linked_to(station) {
return Ok(link);
}
let (existing, direct) = {
let state = self.lock();
if state.closed {
return Err(PoolError::Closed);
}
let existing = state
.members
.iter()
.find(|m| m.target.expected_node_id == *station)
.cloned();
(existing, state.members.iter().filter(|m| m.direct).count())
};
let fresh = existing.is_none();
let member = match existing {
Some(m) => m,
None => {
if direct >= self.opts.max_direct_links {
return Err(PoolError::DirectLinksFull);
}
let target = bounded(deadline, self.station_target(station)).await?;
self.start_member(target, true)
}
};
if let Some(link) = member.await_up(deadline, fresh).await {
return Ok(link);
}
if fresh {
self.drop_member(&member);
}
Err(PoolError::StationNotReached {
station: *station,
cause: member.last_error(),
})
}
async fn station_target(&self, station: &[u8; 32]) -> Result<Target, PoolError> {
let slot = record::station_endpoint_key(station);
let verified = self
.first_answer(|l| async move { l.find_record(&slot).await })
.await
.map_err(|e| match e {
PoolError::Link(link) => PoolError::NoStationEndpoint(Some(link)),
other => other,
})?;
let r = verified.record();
let signer = r.signed.as_ref().map(|s| s.key_id);
let endpoint = record::read_station_endpoint(r)
.map_err(|e| PoolError::NoStationEndpoint(Some(e.into())))?;
match (signer, endpoint.host_advertised.first()) {
(Some(signer), Some(host)) if signer == *station && endpoint.quic_port != 0 => {
Ok(Target {
host: host.clone(),
port: endpoint.quic_port,
profile: self.opts.identity.profile(),
expected_node_id: *station,
})
}
_ => Err(PoolError::NoStationEndpoint(None)),
}
}
pub(super) async fn first_answer<'a, T, F, Fut>(&self, ask: F) -> Result<T, PoolError>
where
F: Fn(Link) -> Fut,
Fut: Future<Output = Result<T, LinkError>> + 'a,
{
let links = self.links();
if links.is_empty() {
return Err(PoolError::NoLink(Vec::new()));
}
let mut errors = Vec::new();
for link in links {
match ask(link).await {
Err(e) if unreachable(&e) => errors.push(e),
answered => return answered.map_err(PoolError::Link),
}
}
Err(PoolError::NoLink(errors))
}
}
fn unreachable(e: &LinkError) -> bool {
matches!(
e,
LinkError::CallTimeout
| LinkError::Closed
| LinkError::LivenessLost
| LinkError::V5DowngradeRefused
| LinkError::Io(_)
| LinkError::Goodbye(_)
| LinkError::StatusExpired
| LinkError::BindingExpired
)
}
fn candidate_share(deadline: Instant, left: usize) -> Instant {
let now = Instant::now();
let remaining = deadline.saturating_duration_since(now);
(now + (remaining / left.max(1) as u32).max(MIN_CANDIDATE_SHARE)).min(deadline)
}
async fn bounded<T>(
deadline: Instant,
work: impl Future<Output = Result<T, PoolError>>,
) -> Result<T, PoolError> {
tokio::time::timeout_at(deadline, work)
.await
.unwrap_or(Err(PoolError::Link(LinkError::CallTimeout)))
}
fn now_ms() -> i64 {
crate::uuid_v7::now_ms() as i64
}
fn callable(
candidates: Vec<Candidate>,
confidential: Confidentiality,
) -> Result<Vec<Candidate>, PoolError> {
let advertised: Vec<[u8; KEY_ID_SIZE]> =
candidates.iter().filter_map(|c| c.kem_key_id).collect();
let keyed: HashSet<[u8; 32]> = candidates
.iter()
.filter(|c| c.kem_key_id.is_some())
.map(|c| c.provider.node)
.collect();
let clear: Vec<Candidate> = match confidential {
Confidentiality::Preferred => candidates
.into_iter()
.filter(|c| !keyed.contains(&c.provider.node))
.collect(),
Confidentiality::Required => Vec::new(),
};
if clear.is_empty() {
return Err(PoolError::Confidentiality(ConfidentialityError {
reason: ConfidentialityReason::NoKemKey,
advertised,
}));
}
Ok(clear)
}
#[cfg(test)]
mod tests {
use super::*;
fn cand(node: u8, kem_key_id: Option<[u8; KEY_ID_SIZE]>) -> Candidate {
Candidate {
provider: Provider {
node: [node; 32],
station: [9; 32],
},
expires_at: 0,
created_at: 0,
kem_key_id,
}
}
fn refused(r: Result<Vec<Candidate>, PoolError>) -> ConfidentialityError {
match r {
Err(PoolError::Confidentiality(e)) => e,
other => panic!("not a confidentiality refusal: {other:?}"),
}
}
#[test]
fn preferred_calls_only_the_providers_that_name_no_key() {
let kept = callable(
vec![cand(1, Some([7; 8])), cand(2, None), cand(3, None)],
Confidentiality::Preferred,
)
.unwrap();
assert_eq!(kept, vec![cand(2, None), cand(3, None)]);
}
#[test]
fn a_provider_that_names_a_key_is_never_called_in_the_clear() {
let e = refused(callable(
vec![cand(1, Some([7; 8])), cand(2, Some([8; 8]))],
Confidentiality::Preferred,
));
assert_eq!(e.reason, ConfidentialityReason::NoKemKey);
assert_eq!(e.advertised, vec![[7; 8], [8; 8]]);
assert_eq!(
e.to_string(),
"confidentiality: no_kem_key 0707070707070707 0808080808080808"
);
}
#[test]
fn a_provider_that_names_a_key_is_not_called_through_its_older_keyless_ad() {
let e = refused(callable(
vec![cand(1, Some([7; 8])), cand(1, None)],
Confidentiality::Preferred,
));
assert_eq!(e.advertised, vec![[7; 8]]);
let kept = callable(
vec![cand(1, Some([7; 8])), cand(1, None), cand(2, None)],
Confidentiality::Preferred,
)
.unwrap();
assert_eq!(kept, vec![cand(2, None)]);
}
#[test]
fn required_calls_no_provider_until_this_sdk_seals() {
let e = refused(callable(
vec![cand(1, None), cand(2, Some([7; 8]))],
Confidentiality::Required,
));
assert_eq!(e.reason, ConfidentialityReason::NoKemKey);
assert_eq!(e.advertised, vec![[7; 8]]);
let e = refused(callable(vec![cand(1, None)], Confidentiality::Required));
assert!(e.advertised.is_empty());
}
#[test]
fn confidential_is_preferred_or_required_and_off_is_refused() {
assert_eq!(Confidentiality::default(), Confidentiality::Preferred);
assert_eq!(
"".parse::<Confidentiality>().unwrap(),
Confidentiality::Preferred
);
assert_eq!(
"preferred".parse::<Confidentiality>().unwrap(),
Confidentiality::Preferred
);
assert_eq!(
"required".parse::<Confidentiality>().unwrap(),
Confidentiality::Required
);
for refused in ["off", "Required", "none", "optional"] {
assert!(
matches!(
refused.parse::<Confidentiality>(),
Err(PoolError::InvalidOpts(_))
),
"{refused}"
);
}
}
}