use std::collections::HashSet;
use std::sync::Arc;
use affinidi_tdk::common::TDKSharedState;
use affinidi_tdk::common::config::TDKConfig;
use affinidi_tdk::messaging::ATM;
use affinidi_tdk::messaging::config::ATMConfig;
use affinidi_tdk::messaging::profiles::ATMProfile;
use affinidi_tdk::secrets_resolver::SecretsResolver;
use affinidi_tdk::secrets_resolver::secrets::Secret;
use tokio::sync::Mutex;
use tracing::{debug, warn};
pub struct SessionHub {
tdk: Arc<TDKSharedState>,
atm: Arc<ATM>,
attached: Mutex<HashSet<String>>,
}
#[derive(Debug)]
pub enum HubError {
AlreadyAttached(String),
Attach(String),
}
impl std::fmt::Display for HubError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::AlreadyAttached(did) => write!(
f,
"`{did}` already has a session on this hub — the mediator permits one \
websocket per DID, so shut that session down before opening another"
),
Self::Attach(msg) => write!(f, "could not attach the identity: {msg}"),
}
}
}
impl std::error::Error for HubError {}
pub(crate) async fn shared_tdk_config() -> Result<TDKConfig, String> {
let resolver = crate::resolver::shared_did_resolver(None)
.await
.map_err(|e| format!("DID resolver init failed: {e}"))?;
TDKConfig::builder()
.with_did_resolver(resolver)
.build()
.map_err(|e| format!("TDK config build failed: {e}"))
}
impl SessionHub {
pub async fn new() -> Result<Arc<Self>, Box<dyn std::error::Error>> {
Self::with_configs(shared_tdk_config().await?, ATMConfig::builder().build()?).await
}
pub async fn with_configs(
tdk_config: TDKConfig,
atm_config: ATMConfig,
) -> Result<Arc<Self>, Box<dyn std::error::Error>> {
let tdk = Arc::new(TDKSharedState::new(tdk_config).await?);
let atm = Arc::new(ATM::new(atm_config, Arc::clone(&tdk)).await?);
debug!("session hub initialised (one TDK + one ATM)");
Ok(Arc::new(Self {
tdk,
atm,
attached: Mutex::new(HashSet::new()),
}))
}
pub(crate) fn atm(&self) -> &Arc<ATM> {
&self.atm
}
pub(crate) fn tdk(&self) -> &Arc<TDKSharedState> {
&self.tdk
}
pub async fn identities(&self) -> Vec<String> {
let mut dids: Vec<String> = self.attached.lock().await.iter().cloned().collect();
dids.sort();
dids
}
pub(crate) async fn attach(
self: &Arc<Self>,
did: &str,
secrets: Vec<Secret>,
mediator_did: &str,
) -> Result<AttachedIdentity, HubError> {
let secret_ids: Vec<String> = secrets.iter().map(|s| s.id.clone()).collect();
if !self.attached.lock().await.insert(did.to_string()) {
return Err(HubError::AlreadyAttached(did.to_string()));
}
for secret in secrets {
self.tdk().secrets_resolver().insert(secret).await;
}
let profile = match self.register_profile(did, mediator_did).await {
Ok(profile) => profile,
Err(e) => {
self.attached.lock().await.remove(did);
self.evict_secrets(&secret_ids).await;
return Err(e);
}
};
Ok(AttachedIdentity {
hub: Arc::clone(self),
profile,
did: did.to_string(),
secret_ids,
})
}
async fn register_profile(
&self,
did: &str,
mediator_did: &str,
) -> Result<Arc<ATMProfile>, HubError> {
let profile = ATMProfile::new(
self.atm.as_ref(),
None,
did.to_string(),
Some(mediator_did.to_string()),
)
.await
.map_err(|e| HubError::Attach(format!("build profile for `{did}`: {e}")))?;
self.atm
.profile_add(&profile, false)
.await
.map_err(|e| HubError::Attach(format!("register profile for `{did}`: {e}")))
}
pub(crate) async fn detach(&self, did: &str, secret_ids: &[String]) {
let was_attached = self.attached.lock().await.remove(did);
if !was_attached {
return;
}
match self.atm.profile_remove(did).await {
Ok(true) => debug!(did, "identity detached from the hub"),
Ok(false) => warn!(did, "identity was not registered with the ATM at detach"),
Err(e) => warn!(did, "could not detach identity cleanly: {e}"),
}
self.evict_secrets(secret_ids).await;
}
async fn evict_secrets(&self, secret_ids: &[String]) {
for id in secret_ids {
self.tdk().secrets_resolver().remove_secret(id).await;
}
}
pub async fn shutdown(&self) {
let remaining = self.attached.lock().await.len();
if remaining > 0 {
debug!(
remaining,
"hub shutdown with identities still attached — stopping their transports"
);
}
self.atm.graceful_shutdown().await;
self.attached.lock().await.clear();
}
}
pub(crate) struct AttachedIdentity {
pub(crate) hub: Arc<SessionHub>,
pub(crate) profile: Arc<ATMProfile>,
pub(crate) did: String,
pub(crate) secret_ids: Vec<String>,
}
impl std::fmt::Debug for AttachedIdentity {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("AttachedIdentity")
.field("did", &self.did)
.field("secrets", &self.secret_ids.len())
.finish()
}
}
impl AttachedIdentity {
pub(crate) async fn detach(&self) {
self.hub.detach(&self.did, &self.secret_ids).await;
}
}
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub(crate) enum HubOwnership {
Exclusive,
Shared,
}
#[cfg(test)]
mod tests {
use super::*;
fn identity(seed_byte: u8) -> (String, Vec<Secret>) {
let seed = [seed_byte; 32];
let sk = ed25519_dalek::SigningKey::from_bytes(&seed);
let did = format!(
"did:key:{}",
crate::did_key::ed25519_multibase_pubkey(&sk.verifying_key().to_bytes())
);
let secrets = crate::did_key::secrets_from_did_key(&did, &seed).expect("build secrets");
(did, vec![secrets.signing, secrets.key_agreement])
}
fn mediator_placeholder() -> String {
identity(0xEE).0
}
#[tokio::test]
async fn attaching_registers_the_identity_and_its_secrets() {
let hub = SessionHub::new().await.expect("build hub");
let (did, secrets) = identity(0x01);
let secret_ids: Vec<String> = secrets.iter().map(|s| s.id.clone()).collect();
let attached = hub
.attach(&did, secrets, &mediator_placeholder())
.await
.expect("attach");
assert_eq!(hub.identities().await, vec![did.clone()]);
assert!(
hub.atm().find_profile(&did).await.is_some(),
"the ATM must know about the profile, not just the session"
);
for id in &secret_ids {
assert!(
hub.tdk().secrets_resolver().get_secret(id).await.is_some(),
"secret {id} must be resolvable while the identity is attached"
);
}
attached.detach().await;
assert!(hub.identities().await.is_empty());
assert!(
hub.atm().find_profile(&did).await.is_none(),
"detach must remove the profile from the ATM"
);
for id in &secret_ids {
assert!(
hub.tdk().secrets_resolver().get_secret(id).await.is_none(),
"detach must evict {id} — a torn-down identity's keys must not \
stay reachable to whatever else runs on the hub"
);
}
hub.shutdown().await;
}
#[tokio::test]
async fn a_second_session_for_the_same_did_is_refused() {
let hub = SessionHub::new().await.expect("build hub");
let mediator = mediator_placeholder();
let (did, secrets) = identity(0x02);
let first = hub.attach(&did, secrets, &mediator).await.expect("attach");
let (_, again) = identity(0x02);
let err = hub
.attach(&did, again, &mediator)
.await
.expect_err("the same DID must not attach twice");
assert!(matches!(err, HubError::AlreadyAttached(ref d) if *d == did));
assert_eq!(hub.identities().await, vec![did.clone()]);
assert!(hub.atm().find_profile(&did).await.is_some());
first.detach().await;
hub.shutdown().await;
}
#[tokio::test]
async fn identities_are_independent_and_detach_is_idempotent() {
let hub = SessionHub::new().await.expect("build hub");
let mediator = mediator_placeholder();
let (alice, alice_secrets) = identity(0x03);
let (bob, bob_secrets) = identity(0x04);
let bob_secret_ids: Vec<String> = bob_secrets.iter().map(|s| s.id.clone()).collect();
let alice_id = hub
.attach(&alice, alice_secrets, &mediator)
.await
.expect("attach alice");
let bob_id = hub
.attach(&bob, bob_secrets, &mediator)
.await
.expect("attach bob");
assert_eq!(hub.identities().await.len(), 2);
bob_id.detach().await;
assert_eq!(hub.identities().await, vec![alice.clone()]);
assert!(hub.atm().find_profile(&alice).await.is_some());
bob_id.detach().await;
assert_eq!(hub.identities().await, vec![alice.clone()]);
for id in &bob_secret_ids {
assert!(hub.tdk().secrets_resolver().get_secret(id).await.is_none());
}
alice_id.detach().await;
hub.shutdown().await;
}
#[tokio::test]
async fn a_detached_did_can_attach_again() {
let hub = SessionHub::new().await.expect("build hub");
let mediator = mediator_placeholder();
let (did, secrets) = identity(0x05);
hub.attach(&did, secrets, &mediator)
.await
.expect("first attach")
.detach()
.await;
let (_, secrets) = identity(0x05);
let second = hub
.attach(&did, secrets, &mediator)
.await
.expect("re-attach after detach");
assert_eq!(hub.identities().await, vec![did]);
second.detach().await;
hub.shutdown().await;
}
#[test]
fn already_attached_names_the_did_and_the_fix() {
let msg = HubError::AlreadyAttached("did:key:zAlice".into()).to_string();
assert!(
msg.contains("did:key:zAlice"),
"names the offending DID: {msg}"
);
assert!(
msg.contains("one websocket per DID"),
"explains why a second session is not extra capacity: {msg}"
);
}
#[test]
fn ownership_distinguishes_a_borrowed_hub_from_an_owned_one() {
assert_ne!(HubOwnership::Exclusive, HubOwnership::Shared);
}
#[tokio::test]
async fn shared_tdk_config_carries_the_shared_resolver() {
let config = shared_tdk_config().await.expect("TDK config");
assert!(
config.did_resolver().is_some(),
"the TDK would build its own PublicOnly resolver"
);
}
#[tokio::test]
async fn opt_in_lets_the_tdk_dial_a_loopback_webvh_host() {
let exe = std::env::current_exe().expect("path of this test binary");
let output = std::process::Command::new(exe)
.args([
"--exact",
"session_hub::tests::child_tdk_dials_loopback_with_opt_in",
"--ignored",
"--nocapture",
])
.env("VTA_ALLOW_PRIVATE_ENDPOINTS", "1")
.env_remove("PNM_RESOLVER_URL")
.output()
.expect("re-run this test binary");
assert!(
output.status.success(),
"child failed ({}):\n--- stdout ---\n{}\n--- stderr ---\n{}",
output.status,
String::from_utf8_lossy(&output.stdout),
String::from_utf8_lossy(&output.stderr),
);
}
#[tokio::test]
#[ignore = "driven by opt_in_lets_the_tdk_dial_a_loopback_webvh_host, which sets the env"]
async fn child_tdk_dials_loopback_with_opt_in() {
use std::sync::atomic::{AtomicUsize, Ordering};
assert_eq!(
std::env::var("VTA_ALLOW_PRIVATE_ENDPOINTS").ok().as_deref(),
Some("1"),
"run through opt_in_lets_the_tdk_dial_a_loopback_webvh_host"
);
let v4 = tokio::net::TcpListener::bind("127.0.0.1:0")
.await
.expect("bind loopback listener");
let port = v4.local_addr().expect("listener address").port();
let v6 = tokio::net::TcpListener::bind(("::1", port)).await.ok();
let accepted = Arc::new(AtomicUsize::new(0));
for listener in std::iter::once(v4).chain(v6) {
let counter = Arc::clone(&accepted);
tokio::spawn(async move {
while let Ok((socket, _)) = listener.accept().await {
counter.fetch_add(1, Ordering::SeqCst);
drop(socket);
}
});
}
let tdk = TDKSharedState::new(shared_tdk_config().await.expect("TDK config"))
.await
.expect("TDK on the shared resolver");
let did = format!("did:webvh:QmScidNotResolvable:localhost%3A{port}");
let outcome = tdk
.did_resolver()
.resolve(&did)
.await
.map(|_| ())
.map_err(|e| e.to_string());
if let Err(message) = &outcome {
assert!(
!message.contains("BlockedHost"),
"opt-in set, yet the TDK's resolver refused {did}: {message}"
);
}
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
assert!(
accepted.load(Ordering::SeqCst) > 0,
"opt-in set, yet {did} was never dialled (outcome: {outcome:?})"
);
}
}