use std::collections::HashMap;
use std::future::Future;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use unb_runtime::{ClientSession, Pipe, SessionOutcome, Wire, WsError};
use web_time::Instant;
pub const DEFAULT_DIAL_TIMEOUT: Duration = Duration::from_secs(5);
pub const CACHE_TTL: Duration = Duration::from_secs(300);
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum TransportKind {
Unix,
WebTransport,
WebSocket,
}
impl TransportKind {
fn rank(self) -> u8 {
match self {
TransportKind::Unix => 0,
TransportKind::WebTransport => 1,
TransportKind::WebSocket => 2,
}
}
fn supported(self) -> bool {
match self {
TransportKind::Unix => cfg!(all(feature = "unix", unix)),
TransportKind::WebTransport => cfg!(feature = "webtransport"),
TransportKind::WebSocket => true,
}
}
}
#[derive(Debug, Clone)]
pub struct Endpoint {
pub kind: TransportKind,
pub address: String,
pub cert_hash: Option<[u8; 32]>,
}
#[derive(Debug, Clone, Default)]
pub struct EndpointSet(Vec<Endpoint>);
impl EndpointSet {
pub fn new() -> EndpointSet {
EndpointSet(Vec::new())
}
pub fn push(&mut self, endpoint: Endpoint) {
self.0.push(endpoint);
}
pub fn is_empty(&self) -> bool {
self.0.is_empty()
}
pub fn len(&self) -> usize {
self.0.len()
}
pub fn iter(&self) -> impl Iterator<Item = &Endpoint> {
self.0.iter()
}
pub fn cache_key(&self) -> String {
let mut addresses: Vec<&str> = self.0.iter().map(|e| e.address.as_str()).collect();
addresses.sort();
addresses.join("|")
}
}
impl From<Endpoint> for EndpointSet {
fn from(endpoint: Endpoint) -> EndpointSet {
EndpointSet(vec![endpoint])
}
}
impl<const N: usize> From<[Endpoint; N]> for EndpointSet {
fn from(endpoints: [Endpoint; N]) -> EndpointSet {
EndpointSet(endpoints.into())
}
}
impl From<Vec<Endpoint>> for EndpointSet {
fn from(endpoints: Vec<Endpoint>) -> EndpointSet {
EndpointSet(endpoints)
}
}
impl FromIterator<Endpoint> for EndpointSet {
fn from_iter<I: IntoIterator<Item = Endpoint>>(iter: I) -> EndpointSet {
EndpointSet(iter.into_iter().collect())
}
}
#[derive(Debug, Clone)]
pub struct DialConfig {
pub attempt_timeout: Duration,
pub supported: Option<Vec<TransportKind>>,
}
impl Default for DialConfig {
fn default() -> DialConfig {
DialConfig {
attempt_timeout: DEFAULT_DIAL_TIMEOUT,
supported: None,
}
}
}
impl DialConfig {
fn kind_supported(&self, kind: TransportKind) -> bool {
match &self.supported {
Some(kinds) => kinds.contains(&kind),
None => kind.supported(),
}
}
}
pub struct Peers {
cache: Mutex<HashMap<String, (TransportKind, Instant)>>,
config: DialConfig,
#[cfg(all(feature = "webtransport", not(target_arch = "wasm32")))]
webtransport: unb_transport::webtransport::ClientPool,
}
impl Peers {
pub fn new() -> Peers {
Peers::with_config(DialConfig::default())
}
pub fn with_config(config: DialConfig) -> Peers {
Peers {
cache: Mutex::new(HashMap::new()),
config,
#[cfg(all(feature = "webtransport", not(target_arch = "wasm32")))]
webtransport: unb_transport::webtransport::ClientPool::new(),
}
}
#[cfg(not(all(target_arch = "wasm32", target_os = "wasi")))]
pub async fn dial(&self, peer: &str, set: &EndpointSet) -> Result<Arc<ClientSession>, WsError> {
self.dial_with(peer, set, |endpoint| async move {
self.dial_candidate(&endpoint).await
})
.await
}
pub async fn dial_with<D, Fut>(
&self,
peer: &str,
set: &EndpointSet,
dialer: D,
) -> Result<Arc<ClientSession>, WsError>
where
D: Fn(Endpoint) -> Fut,
Fut: Future<Output = Result<Pipe, WsError>>,
{
let ordered = self.ordered_candidates(peer, set);
let session = try_candidates_with(&ordered, &self.config, dialer).await?;
self.record_winner(peer, ordered[session.1].kind);
Ok(session.0)
}
pub fn ordered_candidates(&self, key: &str, set: &EndpointSet) -> Vec<Endpoint> {
let cached = {
let cache = self.cache.lock().expect("peer cache lock");
cache.get(key).copied()
};
candidates(set, cached, Instant::now(), &self.config)
}
pub fn record_winner(&self, key: &str, kind: TransportKind) {
let mut cache = self.cache.lock().expect("peer cache lock");
cache.insert(key.to_string(), (kind, Instant::now()));
}
pub fn attempt_timeout(&self) -> Duration {
self.config.attempt_timeout
}
#[cfg(not(all(target_arch = "wasm32", target_os = "wasi")))]
pub async fn dial_candidate(&self, endpoint: &Endpoint) -> Result<Pipe, WsError> {
match endpoint.kind {
TransportKind::Unix => dial_unix(&endpoint.address).await,
TransportKind::WebSocket => {
let (pipe, initiator) = unb_transport::ws::dial(&endpoint.address).await?;
Ok(Pipe::Piped { pipe, initiator })
}
#[cfg(all(feature = "webtransport", not(target_arch = "wasm32")))]
TransportKind::WebTransport => {
let (pipe, initiator, bodies) = self
.webtransport
.dial(&endpoint.address, endpoint.cert_hash)
.await?;
Ok(Pipe::piped_with_streams(pipe, initiator, bodies))
}
#[cfg(all(feature = "webtransport", target_arch = "wasm32"))]
TransportKind::WebTransport => {
let (pipe, initiator) =
unb_transport::webtransport::dial(&endpoint.address, endpoint.cert_hash)
.await?;
Ok(Pipe::Piped { pipe, initiator })
}
#[cfg(not(feature = "webtransport"))]
TransportKind::WebTransport => Err(WsError::Connect(
"webtransport support is not compiled into this client".into(),
)),
}
}
}
impl Default for Peers {
fn default() -> Peers {
Peers::new()
}
}
#[cfg(not(all(target_arch = "wasm32", target_os = "wasi")))]
pub async fn dial_endpoints(set: &EndpointSet) -> Result<Arc<ClientSession>, WsError> {
let ordered = candidates(set, None, Instant::now(), &DialConfig::default());
let dialer = |endpoint: Endpoint| async move { dial_candidate(&endpoint).await };
Ok(
try_candidates_with(&ordered, &DialConfig::default(), dialer)
.await?
.0,
)
}
fn candidates(
set: &EndpointSet,
cached: Option<(TransportKind, Instant)>,
now: Instant,
config: &DialConfig,
) -> Vec<Endpoint> {
let mut ordered: Vec<Endpoint> = set
.0
.iter()
.filter(|endpoint| config.kind_supported(endpoint.kind))
.cloned()
.collect();
ordered.sort_by_key(|endpoint| endpoint.kind.rank());
if let Some((kind, recorded_at)) = cached {
let advertised = ordered.iter().any(|endpoint| endpoint.kind == kind);
let preferred = ordered
.first()
.map(|endpoint| endpoint.kind == kind)
.unwrap_or(false);
let fresh = now.duration_since(recorded_at) < CACHE_TTL;
if advertised && (preferred || fresh) {
ordered.sort_by_key(|endpoint| (endpoint.kind != kind, endpoint.kind.rank()));
}
}
ordered
}
async fn try_candidates_with<D, Fut>(
ordered: &[Endpoint],
config: &DialConfig,
dialer: D,
) -> Result<(Arc<ClientSession>, usize), WsError>
where
D: Fn(Endpoint) -> Fut,
Fut: Future<Output = Result<Pipe, WsError>>,
{
let mut last_error = WsError::Connect("no supported endpoint in set".into());
for (index, endpoint) in ordered.iter().enumerate() {
match n0_future::time::timeout(config.attempt_timeout, async {
ready_session(dialer(endpoint.clone()).await?).await
})
.await
{
Ok(Ok(session)) => return Ok((session, index)),
Ok(Err(error)) => last_error = error,
Err(_) => {
last_error = WsError::Connect(format!("dial timed out for {:?}", endpoint.kind));
}
}
}
Err(last_error)
}
async fn ready_session(pipe: Pipe) -> Result<Arc<ClientSession>, WsError> {
let wire = Wire::open(pipe);
match wire.session_outcome().await? {
SessionOutcome::Established => Ok(Arc::new(wire.client_session())),
SessionOutcome::Retired(reason) => {
Err(WsError::Connect(format!("session retired: {reason:?}")))
}
}
}
#[cfg(all(feature = "unix", unix))]
async fn dial_unix(address: &str) -> Result<Pipe, WsError> {
let (pipe, bodies) =
unb_transport::unix::connect_with_bodies(std::path::Path::new(address)).await?;
Ok(Pipe::piped_with_streams(pipe, true, bodies))
}
#[cfg(not(all(feature = "unix", unix)))]
async fn dial_unix(_address: &str) -> Result<Pipe, WsError> {
Err(WsError::Connect(
"unix transport is not compiled into this client".into(),
))
}
#[cfg(not(all(target_arch = "wasm32", target_os = "wasi")))]
pub async fn dial_candidate(endpoint: &Endpoint) -> Result<Pipe, WsError> {
match endpoint.kind {
TransportKind::Unix => dial_unix(&endpoint.address).await,
TransportKind::WebSocket => {
let (pipe, initiator) = unb_transport::ws::dial(&endpoint.address).await?;
Ok(Pipe::Piped { pipe, initiator })
}
#[cfg(all(feature = "webtransport", not(target_arch = "wasm32")))]
TransportKind::WebTransport => {
let (pipe, initiator, bodies) =
unb_transport::webtransport::dial(&endpoint.address, endpoint.cert_hash).await?;
Ok(Pipe::piped_with_streams(pipe, initiator, bodies))
}
#[cfg(all(feature = "webtransport", target_arch = "wasm32"))]
TransportKind::WebTransport => {
let (pipe, initiator) =
unb_transport::webtransport::dial(&endpoint.address, endpoint.cert_hash).await?;
Ok(Pipe::Piped { pipe, initiator })
}
#[cfg(not(feature = "webtransport"))]
TransportKind::WebTransport => Err(WsError::Connect(
"webtransport support is not compiled into this client".into(),
)),
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::pair;
fn set(kinds: &[TransportKind]) -> EndpointSet {
kinds
.iter()
.map(|kind| Endpoint {
kind: *kind,
address: String::new(),
cert_hash: None,
})
.collect()
}
#[test]
fn webtransport_is_preferred_when_supported() {
let ordered = candidates(
&set(&[TransportKind::WebSocket, TransportKind::WebTransport]),
None,
Instant::now(),
&DialConfig::default(),
);
if TransportKind::WebTransport.supported() {
assert_eq!(ordered[0].kind, TransportKind::WebTransport);
assert_eq!(ordered[1].kind, TransportKind::WebSocket);
} else {
assert_eq!(ordered.len(), 1);
assert_eq!(ordered[0].kind, TransportKind::WebSocket);
}
}
#[test]
fn a_fresh_non_preferred_winner_is_tried_first() {
let now = Instant::now();
let ordered = candidates(
&set(&[TransportKind::WebTransport, TransportKind::WebSocket]),
Some((TransportKind::WebSocket, now)),
now,
&DialConfig::default(),
);
assert_eq!(ordered[0].kind, TransportKind::WebSocket);
}
#[test]
fn an_expired_non_preferred_winner_falls_back_to_preference_order() {
let now = Instant::now();
let stale = now.checked_sub(CACHE_TTL + Duration::from_secs(1)).unwrap();
let ordered = candidates(
&set(&[TransportKind::WebTransport, TransportKind::WebSocket]),
Some((TransportKind::WebSocket, stale)),
now,
&DialConfig::default(),
);
if TransportKind::WebTransport.supported() {
assert_eq!(ordered[0].kind, TransportKind::WebTransport);
}
}
#[test]
fn a_cached_kind_absent_from_the_set_is_ignored() {
let now = Instant::now();
let ordered = candidates(
&set(&[TransportKind::WebTransport]),
Some((TransportKind::WebSocket, now)),
now,
&DialConfig::default(),
);
assert!(ordered
.iter()
.all(|endpoint| endpoint.kind == TransportKind::WebTransport));
}
#[test]
fn injected_platform_support_overrides_the_compiled_intersection() {
let config = DialConfig {
supported: Some(vec![TransportKind::WebSocket]),
..DialConfig::default()
};
let ordered = candidates(
&set(&[TransportKind::WebTransport, TransportKind::WebSocket]),
None,
Instant::now(),
&config,
);
assert_eq!(ordered.len(), 1);
assert_eq!(ordered[0].kind, TransportKind::WebSocket);
let none_supported = DialConfig {
supported: Some(Vec::new()),
..DialConfig::default()
};
let empty = candidates(
&set(&[TransportKind::WebTransport, TransportKind::WebSocket]),
None,
Instant::now(),
&none_supported,
);
assert!(empty.is_empty());
}
#[test]
fn unsupported_kinds_are_never_candidates() {
let ordered = candidates(
&set(&[TransportKind::WebTransport, TransportKind::WebSocket]),
None,
Instant::now(),
&DialConfig::default(),
);
assert!(ordered.iter().all(|endpoint| endpoint.kind.supported()));
}
#[cfg(not(target_arch = "wasm32"))]
#[tokio::test]
async fn an_empty_intersection_is_an_error() {
let result = dial_endpoints(&EndpointSet::new()).await;
assert!(matches!(result, Err(WsError::Connect(_))));
}
#[cfg(not(target_arch = "wasm32"))]
#[tokio::test]
async fn a_handshake_failure_falls_back_to_the_next_endpoint() {
let set: EndpointSet = vec![
Endpoint {
kind: TransportKind::WebTransport,
address: "failed".into(),
cert_hash: None,
},
Endpoint {
kind: TransportKind::WebSocket,
address: "ready".into(),
cert_hash: None,
},
]
.into();
let peers = Peers::with_config(DialConfig {
attempt_timeout: Duration::from_millis(100),
supported: Some(vec![TransportKind::WebTransport, TransportKind::WebSocket]),
});
let session = peers
.dial_with("owner", &set, |endpoint| async move {
if endpoint.address == "failed" {
return Err(WsError::Connect("failed".into()));
}
let (client, server) = pair();
let _server = Wire::open(server);
Ok(client)
})
.await
.unwrap();
assert_eq!(Arc::strong_count(&session), 1);
assert_eq!(
peers
.cache
.lock()
.unwrap()
.get("owner")
.map(|entry| entry.0),
Some(TransportKind::WebSocket)
);
}
#[cfg(not(target_arch = "wasm32"))]
#[tokio::test]
async fn handshake_timeout_falls_back_without_caching_the_timed_out_candidate() {
let set: EndpointSet = vec![
Endpoint {
kind: TransportKind::WebTransport,
address: "stalled".into(),
cert_hash: None,
},
Endpoint {
kind: TransportKind::WebSocket,
address: "ready".into(),
cert_hash: None,
},
]
.into();
let peers = Peers::with_config(DialConfig {
attempt_timeout: Duration::from_millis(25),
supported: Some(vec![TransportKind::WebTransport, TransportKind::WebSocket]),
});
peers
.dial_with("owner", &set, |endpoint| async move {
let (client, server) = pair();
if endpoint.address == "stalled" {
std::mem::forget(server);
} else {
let _server = Wire::open(server);
}
Ok(client)
})
.await
.unwrap();
assert_eq!(
peers
.cache
.lock()
.unwrap()
.get("owner")
.map(|entry| entry.0),
Some(TransportKind::WebSocket)
);
}
#[cfg(not(target_arch = "wasm32"))]
#[tokio::test]
async fn a_ready_winner_is_preferred_on_the_next_dial() {
let set = set(&[TransportKind::WebTransport, TransportKind::WebSocket]);
let peers = Peers::with_config(DialConfig {
attempt_timeout: Duration::from_millis(100),
supported: Some(vec![TransportKind::WebTransport, TransportKind::WebSocket]),
});
peers
.dial_with("owner", &set, |endpoint| async move {
if endpoint.kind == TransportKind::WebTransport {
return Err(WsError::Connect("failed".into()));
}
let (client, server) = pair();
let _server = Wire::open(server);
Ok(client)
})
.await
.unwrap();
assert_eq!(
peers.ordered_candidates("owner", &set)[0].kind,
TransportKind::WebSocket
);
}
}