[Fix] Connection Management

other-iota/Cargo.toml

@@ -4,3 +4,17 @@ version = "0.1.0"
 edition = "2024"
 
 [dependencies]
+async-trait = "0.1.89"
+iota-connection = { path = "../iota-connection" }
+iota-identity = { path = "../iota-identity" }
+iota-util = { path = "../iota-util" }
+mtp = { git = "https://git.methanium.net/Methanium/mtp.git", rev = "1f19a0d897c265d1e3f590a876f95e766ff99318", features = ["client", "crypto", "pipes", "web-server"] }
+reqwest = { version = "0.13", features = ["json"] }
+serde = "1"
+tokio = { version = "1.50.0", features = ["full"] }
+web-server = { path = "../web-server" }
+
+[dev-dependencies]
+iota-storage = { path = "../iota-storage" }
+rcgen = "0.14"
+tempfile = "3"

other-iota/src/relay_router.rs

@@ -0,0 +1,413 @@
+use async_trait::async_trait;
+use iota_connection::connection_handler::{PeerRouter, RouteDestination, RouteError, RouteOutcome};
+use iota_connection::relay_service::{IngressSource, RelayService};
+use iota_identity::{
+    AuthorityLocator, IotaNodeId, LocalNodeIdentity, PublicKeyBundle, RelayRouterId,
+};
+use iota_util::mtp_compat::OptionalDataValueExt;
+use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue, TypeMap};
+use std::collections::HashMap;
+use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering};
+use std::sync::{Arc, RwLock};
+
+use crate::{LocalDeliverySink, route_outcome};
+
+const ROUTER_PROTOCOL: &str = "tensamin-relay-router:v1";
+static NEXT_ROUTER_FRAME_ID: AtomicU32 = AtomicU32::new(1);
+
+#[async_trait]
+trait RouterPeerConnection: Send + Sync {
+    async fn request(&self, frame: CommunicationValue) -> Result<CommunicationValue, RouteError>;
+}
+
+type RegisteredRouterPeer = (u64, Arc<dyn RouterPeerConnection>);
+
+#[derive(Default)]
+pub struct RelayRouterService {
+    peers: RwLock<HashMap<IotaNodeId, RegisteredRouterPeer>>,
+    next_registration: AtomicU64,
+}
+
+impl RelayRouterService {
+    pub fn has_connection(&self, node: &IotaNodeId) -> bool {
+        self.peers
+            .read()
+            .is_ok_and(|peers| peers.contains_key(node))
+    }
+
+    fn register(&self, node: IotaNodeId, connection: Arc<dyn RouterPeerConnection>) -> u64 {
+        let registration = self.next_registration.fetch_add(1, Ordering::Relaxed) + 1;
+        if let Ok(mut peers) = self.peers.write() {
+            peers.insert(node, (registration, connection));
+        }
+        registration
+    }
+
+    fn unregister(&self, node: &IotaNodeId, registration: u64) {
+        if let Ok(mut peers) = self.peers.write()
+            && peers
+                .get(node)
+                .is_some_and(|current| current.0 == registration)
+        {
+            peers.remove(node);
+        }
+    }
+
+    async fn forward(
+        &self,
+        source: &IotaNodeId,
+        destination: &IotaNodeId,
+        frame: CommunicationValue,
+    ) -> Result<CommunicationValue, RouteError> {
+        if source == destination {
+            return Err(RouteError::Delivery(
+                "relay router source and destination are identical".into(),
+            ));
+        }
+        let peer = self
+            .peers
+            .read()
+            .ok()
+            .and_then(|peers| peers.get(destination).map(|peer| peer.1.clone()))
+            .ok_or_else(|| RouteError::NoRoute(destination.clone()))?;
+        peer.request(frame).await
+    }
+}
+
+#[async_trait]
+impl web_server::MtpConnectionHandler for RelayRouterService {
+    async fn accept(&self, connection: mtp::webserver::WebMTPConnection) {
+        if connection.auth_state != mtp::host::AuthState::Authenticated {
+            return;
+        }
+        let Some(public_key) = connection.client_public_key.clone() else {
+            return;
+        };
+        let Ok(node) = IotaNodeId::from_public_keys(&public_key) else {
+            return;
+        };
+        let connection = Arc::new(connection);
+        let pending = Arc::new(tokio::sync::Mutex::new(HashMap::new()));
+        let registration = self.register(
+            node.clone(),
+            Arc::new(IncomingRouterPeer {
+                connection: connection.clone(),
+                pending: pending.clone(),
+            }),
+        );
+        while let Ok(frame) = connection.receive().await {
+            if let Some(frame_id) = frame.id()
+                && let Some(response) = pending.lock().await.remove(&frame_id)
+            {
+                let _ = response.send(frame);
+                continue;
+            }
+            let Some(frame_id) = frame.id() else {
+                continue;
+            };
+            let response = match decode_router_request(&frame) {
+                Ok((destination, relay)) => match self.forward(&node, &destination, relay).await {
+                    Ok(response) => response.with_id(frame_id),
+                    Err(RouteError::NoRoute(_)) => {
+                        CommunicationValue::new(CommunicationType::ErrorNoIota).with_id(frame_id)
+                    }
+                    Err(RouteError::Delivery(message)) => {
+                        CommunicationValue::new(CommunicationType::ErrorInternal)
+                            .with_id(frame_id)
+                            .add_typed_default(DataType::ErrorMessage, DataValue::Str(message))
+                    }
+                },
+                Err(message) => CommunicationValue::new(CommunicationType::ErrorInvalidData)
+                    .with_id(frame_id)
+                    .add_typed_default(DataType::ErrorMessage, DataValue::Str(message)),
+            };
+            let _ = connection.sender.send(&response).await;
+        }
+        self.unregister(&node, registration);
+    }
+}
+
+struct IncomingRouterPeer {
+    connection: Arc<mtp::webserver::WebMTPConnection>,
+    pending:
+        Arc<tokio::sync::Mutex<HashMap<u32, tokio::sync::oneshot::Sender<CommunicationValue>>>>,
+}
+
+#[async_trait]
+impl RouterPeerConnection for IncomingRouterPeer {
+    async fn request(
+        &self,
+        mut frame: CommunicationValue,
+    ) -> Result<CommunicationValue, RouteError> {
+        let frame_id = next_frame_id();
+        frame = frame.with_id(frame_id);
+        let (send, receive) = tokio::sync::oneshot::channel();
+        if self.pending.lock().await.insert(frame_id, send).is_some() {
+            return Err(RouteError::Delivery(
+                "relay router request ID is already pending".into(),
+            ));
+        }
+        if let Err(error) = self.connection.sender.send(&frame).await {
+            self.pending.lock().await.remove(&frame_id);
+            return Err(RouteError::Delivery(error.to_string()));
+        }
+        tokio::time::timeout(std::time::Duration::from_secs(30), receive)
+            .await
+            .map_err(|_| RouteError::Delivery("relay router request timed out".into()))?
+            .map_err(|_| RouteError::Delivery("relay router connection closed".into()))
+    }
+}
+
+pub struct RelayRouterClient {
+    router_id: RelayRouterId,
+    connection: Arc<mtp::client::MTPConnection>,
+    connected: AtomicBool,
+    disconnected: tokio::sync::Notify,
+}
+
+#[derive(Default)]
+pub struct RelayRouterSet {
+    routers: RwLock<Vec<Arc<RelayRouterClient>>>,
+}
+
+impl RelayRouterSet {
+    pub fn add(&self, router: Arc<RelayRouterClient>) {
+        if let Ok(mut routers) = self.routers.write() {
+            routers.retain(|current| current.router_id() != router.router_id());
+            routers.push(router);
+        }
+    }
+}
+
+#[async_trait]
+impl PeerRouter for RelayRouterSet {
+    async fn route(
+        &self,
+        destination: &RouteDestination,
+        frame: CommunicationValue,
+    ) -> Result<RouteOutcome, RouteError> {
+        let routers = self
+            .routers
+            .read()
+            .map(|routers| routers.clone())
+            .unwrap_or_default();
+        if routers.is_empty() {
+            return Ok(RouteOutcome::Retryable {
+                reason: "no relay router is configured".into(),
+            });
+        }
+        let mut last_reason = "all relay routers are unavailable".to_string();
+        for router in routers {
+            if !router.is_connected() {
+                continue;
+            }
+            match router.route(destination, frame.clone()).await {
+                outcome @ Ok(RouteOutcome::Accepted { .. } | RouteOutcome::Rejected { .. }) => {
+                    return outcome;
+                }
+                Ok(RouteOutcome::Retryable { reason }) => last_reason = reason,
+                Err(error) => last_reason = error.to_string(),
+            }
+        }
+        Ok(RouteOutcome::Retryable {
+            reason: last_reason,
+        })
+    }
+}
+
+impl RelayRouterClient {
+    pub async fn connect(
+        identity: LocalNodeIdentity,
+        endpoint: &AuthorityLocator,
+        pinned_certificate: Vec<u8>,
+        router_key: PublicKeyBundle,
+        relay: Arc<RelayService>,
+        deliveries: Arc<dyn LocalDeliverySink>,
+    ) -> Result<Arc<Self>, RouteError> {
+        let router_id = RelayRouterId::from_public_keys(&router_key)
+            .map_err(|error| RouteError::Delivery(error.to_string()))?;
+        let config = mtp::client::ClientConfig::new(format!("https://{}/", endpoint.as_str()))
+            .with_pinned_pem(pinned_certificate)
+            .with_client_id(identity.node_id().mtp_authentication_hint())
+            .with_description(format!("iota-router: {}", identity.node_id().as_str()));
+        let connection =
+            mtp::client::MTPClient::auth_connect(config, &identity.keyring(), &router_key)
+                .await
+                .map_err(|error| RouteError::Delivery(error.to_string()))?;
+        if connection.auth_state != mtp::client::AuthState::Authenticated {
+            return Err(RouteError::Delivery(
+                "relay router did not complete authentication".into(),
+            ));
+        }
+        let client = Arc::new(Self {
+            router_id,
+            connection: Arc::new(connection),
+            connected: AtomicBool::new(true),
+            disconnected: tokio::sync::Notify::new(),
+        });
+        let reader = client.clone();
+        tokio::spawn(async move {
+            while let Ok(frame) = reader.connection.receive().await {
+                let outcome = relay
+                    .accept_relay(
+                        IngressSource::RelayRouter {
+                            router_id: reader.router_id.clone(),
+                        },
+                        frame,
+                    )
+                    .await;
+                let Ok(outcome) = outcome else {
+                    continue;
+                };
+                let (response, local_deliveries) = outcome.into_parts();
+                for delivery in local_deliveries {
+                    deliveries.deliver(delivery.recipient, delivery.frame).await;
+                }
+                if let Some(response) = response {
+                    let _ = reader.connection.sender.send(&response).await;
+                }
+            }
+            reader.connected.store(false, Ordering::Release);
+            reader.disconnected.notify_waiters();
+        });
+        Ok(client)
+    }
+
+    pub fn router_id(&self) -> &RelayRouterId {
+        &self.router_id
+    }
+
+    pub fn is_connected(&self) -> bool {
+        self.connected.load(Ordering::Acquire)
+    }
+
+    pub async fn wait_disconnected(&self) {
+        loop {
+            let disconnected = self.disconnected.notified();
+            if !self.is_connected() {
+                return;
+            }
+            disconnected.await;
+        }
+    }
+}
+
+#[async_trait]
+impl PeerRouter for RelayRouterClient {
+    async fn route(
+        &self,
+        destination: &RouteDestination,
+        frame: CommunicationValue,
+    ) -> Result<RouteOutcome, RouteError> {
+        let RouteDestination::Iota(destination) = destination else {
+            return Ok(RouteOutcome::Retryable {
+                reason: "relay router cannot route a legacy Omega destination".into(),
+            });
+        };
+        let request = encode_router_request(destination, &frame)?;
+        let response = self
+            .connection
+            .request(&request, None)
+            .await
+            .map_err(|error| RouteError::Delivery(error.to_string()))?;
+        if response.is_type(CommunicationType::ErrorNoIota)
+            || response.is_type(CommunicationType::ErrorInternal)
+        {
+            return Ok(RouteOutcome::Retryable {
+                reason: response
+                    .get_data(DataType::ErrorMessage)
+                    .as_str()
+                    .unwrap_or("relay router destination is unavailable")
+                    .to_owned(),
+            });
+        }
+        route_outcome(response)
+    }
+}
+
+pub struct DirectThenRelayRouter {
+    direct: Arc<dyn PeerRouter>,
+    relay: Arc<dyn PeerRouter>,
+}
+
+impl DirectThenRelayRouter {
+    pub fn new(direct: Arc<dyn PeerRouter>, relay: Arc<dyn PeerRouter>) -> Self {
+        Self { direct, relay }
+    }
+}
+
+#[async_trait]
+impl PeerRouter for DirectThenRelayRouter {
+    async fn route(
+        &self,
+        destination: &RouteDestination,
+        frame: CommunicationValue,
+    ) -> Result<RouteOutcome, RouteError> {
+        match self.direct.route(destination, frame.clone()).await {
+            outcome @ Ok(RouteOutcome::Accepted { .. } | RouteOutcome::Rejected { .. }) => outcome,
+            Ok(RouteOutcome::Retryable { .. }) | Err(_) => {
+                self.relay.route(destination, frame).await
+            }
+        }
+    }
+}
+
+fn encode_router_request(
+    destination: &IotaNodeId,
+    frame: &CommunicationValue,
+) -> Result<CommunicationValue, RouteError> {
+    let type_map = frame
+        .type_map()
+        .ok_or_else(|| RouteError::Delivery("relay frame has no type map".into()))?;
+    let payload = frame
+        .to_bytes()
+        .map_err(|error| RouteError::Delivery(error.to_string()))?;
+    Ok(CommunicationValue::new(CommunicationType::Relay)
+        .with_id(next_frame_id())
+        .add_typed_default(DataType::About, DataValue::Str(ROUTER_PROTOCOL.into()))
+        .add_typed_default(
+            DataType::Description,
+            DataValue::Str(destination.as_str().into()),
+        )
+        .add_typed_default(
+            DataType::Version,
+            DataValue::Str(type_map.version.to_string()),
+        )
+        .add_typed_default(DataType::SecurePayload, DataValue::Bytes(payload)))
+}
+
+fn decode_router_request(
+    frame: &CommunicationValue,
+) -> Result<(IotaNodeId, CommunicationValue), String> {
+    if !frame.is_type(CommunicationType::Relay)
+        || frame.get_data(DataType::About).as_str() != Some(ROUTER_PROTOCOL)
+    {
+        return Err("frame is not a relay-router request".into());
+    }
+    let destination = frame
+        .get_data(DataType::Description)
+        .as_str()
+        .ok_or_else(|| "relay-router destination is missing".to_string())
+        .and_then(|destination| IotaNodeId::new(destination).map_err(|error| error.to_string()))?;
+    let version = frame
+        .get_data(DataType::Version)
+        .as_str()
+        .and_then(mtp::type_map::Version::parse)
+        .ok_or_else(|| "relay-router type-map version is invalid".to_string())?;
+    let payload = match frame.get_data(DataType::SecurePayload) {
+        Some(DataValue::Bytes(payload)) => payload,
+        _ => return Err("relay-router payload is missing".into()),
+    };
+    let relay = CommunicationValue::from_bytes_with(payload, &TypeMap::new(version))
+        .map_err(|error| error.to_string())?;
+    Ok((destination, relay))
+}
+
+fn next_frame_id() -> u32 {
+    loop {
+        let id = NEXT_ROUTER_FRAME_ID.fetch_add(1, Ordering::Relaxed);
+        if id != 0 {
+            return id;
+        }
+    }
+}

other-iota/tests/direct_federation.rs

@@ -0,0 +1,582 @@
+use async_trait::async_trait;
+use iota_connection::connection_handler::{PeerRouter, RouteDestination, RouteError, RouteOutcome};
+use iota_connection::federated_relay::FederatedRelayV2;
+use iota_connection::relay_service::{RelayNodeIdentity, RelayService};
+use iota_identity::{
+    AuthorityLocator, IdentityError, IdentityResolver, IotaNodeId, LocalNodeIdentity,
+    LocalUserDescriptor, LocalUserId, LocalUserStore, NodeDirectory, NodeIdentityResolver,
+    PrincipalDescriptor, PrincipalHandle, PrincipalHome, PrincipalId, PrincipalStore,
+    PublicKeyBundle, ResolutionContext, ResolvedNodeIdentity, ResolvedPrincipal,
+    SignedPrincipalDescriptor, UserAddress,
+};
+use iota_storage::identity::SqlitePrincipalStore;
+use iota_storage::util::{chat_files, e2ee_storage};
+use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue};
+use mtp::crypto::Keyring;
+use mtp::host::{AuthenticationPolicy, HostConfig};
+use mtp::webserver::{MTPWebServer, WebServerConfig};
+use other_iota::{DirectPeerRouter, LocalDeliverySink, PeerManager};
+use std::collections::HashMap;
+use std::net::{IpAddr, Ipv4Addr};
+use std::sync::{Arc, Mutex, RwLock};
+
+struct StoreResolver;
+
+#[async_trait]
+impl IdentityResolver for StoreResolver {
+    async fn resolve_address(
+        &self,
+        _: &UserAddress,
+        _: &ResolutionContext,
+    ) -> Result<ResolvedPrincipal, IdentityError> {
+        Err(IdentityError::NotFound)
+    }
+
+    async fn resolve_principal(
+        &self,
+        principal: &PrincipalId,
+    ) -> Result<ResolvedPrincipal, IdentityError> {
+        SqlitePrincipalStore
+            .get_by_canonical_id(principal)?
+            .ok_or(IdentityError::NotFound)
+    }
+
+    async fn signing_keys(
+        &self,
+        principal: &PrincipalId,
+        _: &ResolutionContext,
+    ) -> Result<Vec<PublicKeyBundle>, IdentityError> {
+        SqlitePrincipalStore.signing_keys(principal)
+    }
+}
+
+struct HostedRecipient {
+    principal: PrincipalHandle,
+}
+
+impl LocalUserStore for HostedRecipient {
+    fn get_local_user(&self, _: LocalUserId) -> Result<Option<LocalUserDescriptor>, IdentityError> {
+        Ok(None)
+    }
+
+    fn get_local_user_by_username(
+        &self,
+        _: &str,
+    ) -> Result<Option<LocalUserDescriptor>, IdentityError> {
+        Ok(None)
+    }
+
+    fn is_hosted_here(&self, id: LocalUserId) -> Result<bool, IdentityError> {
+        Ok(id == LocalUserId(1))
+    }
+
+    fn local_user_for_principal(
+        &self,
+        principal: PrincipalHandle,
+    ) -> Result<Option<LocalUserId>, IdentityError> {
+        Ok((principal == self.principal).then_some(LocalUserId(1)))
+    }
+
+    fn principal_for_local_user(
+        &self,
+        user: LocalUserId,
+    ) -> Result<Option<PrincipalHandle>, IdentityError> {
+        Ok((user == LocalUserId(1)).then_some(self.principal))
+    }
+}
+
+struct TestNodeIdentity(LocalNodeIdentity);
+
+#[async_trait]
+impl RelayNodeIdentity for TestNodeIdentity {
+    async fn keyring(&self) -> Option<Arc<Keyring>> {
+        Some(self.0.keyring())
+    }
+
+    fn node_id(&self) -> Option<IotaNodeId> {
+        Some(self.0.node_id().clone())
+    }
+}
+
+struct NoRoute;
+
+#[async_trait]
+impl PeerRouter for NoRoute {
+    async fn route(
+        &self,
+        destination: &RouteDestination,
+        _: CommunicationValue,
+    ) -> Result<RouteOutcome, RouteError> {
+        match destination {
+            RouteDestination::Iota(node) => Err(RouteError::NoRoute(node.clone())),
+            RouteDestination::LegacyOmegaIota { .. } => Ok(RouteOutcome::Retryable {
+                reason: "legacy route is unavailable".into(),
+            }),
+        }
+    }
+}
+
+#[derive(Default)]
+struct MemoryNodes {
+    nodes: RwLock<HashMap<IotaNodeId, ResolvedNodeIdentity>>,
+    certificates: RwLock<HashMap<AuthorityLocator, Vec<u8>>>,
+}
+
+#[async_trait]
+impl NodeIdentityResolver for MemoryNodes {
+    async fn resolve_node(&self, node: &IotaNodeId) -> Result<ResolvedNodeIdentity, IdentityError> {
+        self.nodes
+            .read()
+            .map_err(|_| IdentityError::Storage("node lock is poisoned".into()))?
+            .get(node)
+            .cloned()
+            .ok_or(IdentityError::NotFound)
+    }
+}
+
+impl NodeDirectory for MemoryNodes {
+    fn upsert_node_descriptor(
+        &self,
+        _: &iota_identity::VerifiedNodeDescriptor,
+    ) -> Result<(), IdentityError> {
+        Ok(())
+    }
+
+    fn pin_endpoint(
+        &self,
+        _: &AuthorityLocator,
+        _: &IotaNodeId,
+        _: i64,
+    ) -> Result<(), IdentityError> {
+        Ok(())
+    }
+
+    fn pinned_node(&self, _: &AuthorityLocator) -> Result<Option<IotaNodeId>, IdentityError> {
+        Ok(None)
+    }
+
+    fn pin_endpoint_certificate(
+        &self,
+        endpoint: &AuthorityLocator,
+        certificate: &[u8],
+    ) -> Result<(), IdentityError> {
+        self.certificates
+            .write()
+            .map_err(|_| IdentityError::Storage("certificate lock is poisoned".into()))?
+            .insert(endpoint.clone(), certificate.to_vec());
+        Ok(())
+    }
+
+    fn pinned_endpoint_certificate(
+        &self,
+        endpoint: &AuthorityLocator,
+    ) -> Result<Option<Vec<u8>>, IdentityError> {
+        Ok(self
+            .certificates
+            .read()
+            .map_err(|_| IdentityError::Storage("certificate lock is poisoned".into()))?
+            .get(endpoint)
+            .cloned())
+    }
+
+    fn record_successful_endpoint(
+        &self,
+        _: &IotaNodeId,
+        _: &AuthorityLocator,
+    ) -> Result<(), IdentityError> {
+        Ok(())
+    }
+}
+
+#[derive(Default)]
+struct Deliveries(Mutex<Vec<(PrincipalHandle, CommunicationValue)>>);
+
+#[async_trait]
+impl LocalDeliverySink for Deliveries {
+    async fn deliver(&self, recipient: PrincipalHandle, frame: CommunicationValue) {
+        self.0.lock().unwrap().push((recipient, frame));
+    }
+}
+
+fn verified_principal(
+    authority: &LocalNodeIdentity,
+    user: &Keyring,
+    user_id: u64,
+) -> iota_identity::VerifiedPrincipalDescriptor {
+    let now = iota_storage::util::sync::now_millis();
+    SignedPrincipalDescriptor::sign(
+        PrincipalDescriptor {
+            principal: PrincipalId {
+                authority: authority.authority_id().clone(),
+                user_id,
+            },
+            authority_kind: iota_identity::AuthorityKind::Iota,
+            username: Some("alice".into()),
+            display_name: None,
+            public_keys: vec![user.public_key_bundle()],
+            home: PrincipalHome::Iota(authority.node_id().clone()),
+            revision: 1,
+            valid_until: Some(now + 60_000),
+            issued_at: now,
+        },
+        &authority.keyring(),
+    )
+    .unwrap()
+    .verify(
+        authority.authority_id(),
+        &authority.public_keys(),
+        None,
+        now,
+    )
+    .unwrap()
+}
+
+fn resolved_node(
+    identity: &LocalNodeIdentity,
+    direct_endpoints: Vec<AuthorityLocator>,
+) -> ResolvedNodeIdentity {
+    ResolvedNodeIdentity {
+        node: identity.node_id().clone(),
+        public_keys: vec![identity.public_keys()],
+        descriptor_revision: 1,
+        valid_until: None,
+        direct_endpoints,
+        relay_hints: Vec::new(),
+        last_successful_endpoint: None,
+        resolved_at: iota_storage::util::sync::now_millis(),
+    }
+}
+
+fn certificate() -> (Vec<u8>, Vec<u8>) {
+    let key_pair = rcgen::KeyPair::generate().unwrap();
+    let params =
+        rcgen::CertificateParams::new(vec!["localhost".into(), "127.0.0.1".into()]).unwrap();
+    let certificate = params.self_signed(&key_pair).unwrap();
+    (
+        certificate.pem().into_bytes(),
+        key_pair.serialize_pem().into_bytes(),
+    )
+}
+
+async fn start_peer_server(
+    port: u16,
+    certificate: Vec<u8>,
+    private_key: Vec<u8>,
+    client: &LocalNodeIdentity,
+    host: &LocalNodeIdentity,
+    manager: Arc<PeerManager>,
+) -> (
+    AuthorityLocator,
+    tokio::sync::oneshot::Sender<()>,
+    tokio::task::JoinHandle<()>,
+) {
+    let client_key = client.public_keys();
+    let client_hint = client.node_id().mtp_authentication_hint();
+    let host_keyring = Keyring::from_bytes(&host.keyring().try_to_bytes().unwrap()).unwrap();
+    let host = HostConfig::new(
+        IpAddr::V4(Ipv4Addr::LOCALHOST),
+        port,
+        certificate,
+        private_key,
+    )
+    .with_authentication(
+        host_keyring,
+        Box::new(move |client_id, _| {
+            let public_key = (client_id == client_hint).then(|| client_key.clone());
+            Box::pin(async move { public_key })
+        }),
+        Box::new(|_, _| Box::pin(async { 0 })),
+    )
+    .with_authentication_policy(AuthenticationPolicy::ForceAuthentication);
+    let mut server = MTPWebServer::new(host, WebServerConfig::new())
+        .await
+        .unwrap();
+    let endpoint = AuthorityLocator::new(server.local_addr().to_string()).unwrap();
+    let (shutdown, mut shutdown_requested) = tokio::sync::oneshot::channel();
+    let task = tokio::spawn(async move {
+        loop {
+            tokio::select! {
+                _ = &mut shutdown_requested => {
+                    server.shutdown().await;
+                    break;
+                }
+                accepted = server.accept() => match accepted {
+                    Ok(Some(connection)) => {
+                        let manager = manager.clone();
+                        tokio::spawn(async move {
+                            web_server::MtpConnectionHandler::accept(manager.as_ref(), connection)
+                                .await;
+                        });
+                    }
+                    Ok(None) => break,
+                    Err(_) => continue,
+                }
+            }
+        }
+    });
+    (endpoint, shutdown, task)
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+async fn authenticated_peer_delivers_relay_v2_to_same_numeric_user_on_another_iota() {
+    let storage = tempfile::tempdir().unwrap();
+    iota_util::file_util::configure_storage_directory(storage.path().to_owned());
+    iota_storage::util::db::initialize_database().unwrap();
+
+    let iota_a = LocalNodeIdentity::from_keyring(Keyring::generate()).unwrap();
+    let iota_b = LocalNodeIdentity::from_keyring(Keyring::generate()).unwrap();
+    let user_a = Keyring::generate();
+    let user_b = Keyring::generate();
+    let principal_a = verified_principal(&iota_a, &user_a, 1);
+    let principal_b = verified_principal(&iota_b, &user_b, 1);
+    let handle_a = SqlitePrincipalStore
+        .upsert_remote_descriptor(&principal_a)
+        .unwrap();
+    let handle_b = SqlitePrincipalStore
+        .upsert_remote_descriptor(&principal_b)
+        .unwrap();
+    assert_ne!(handle_a, handle_b);
+
+    let chat_id = e2ee_storage::principal_chat_id(handle_a, handle_b).unwrap();
+    e2ee_storage::put_chat_secret(e2ee_storage::StoredChatSecret {
+        user_id: "1".into(),
+        secret_id: e2ee_storage::principal_secret_id(&chat_id),
+        chat_id,
+        version: 1,
+        encrypted_secret: vec![1],
+        kem_ciphertext: vec![2],
+        wrapping_scheme: "test".into(),
+        created_at: 1,
+        updated_at: 1,
+    })
+    .unwrap();
+
+    let resolver: Arc<dyn IdentityResolver> = Arc::new(StoreResolver);
+    let local_users: Arc<dyn LocalUserStore> = Arc::new(HostedRecipient {
+        principal: handle_b,
+    });
+    let relay_b = Arc::new(
+        RelayService::new(local_users, Arc::new(NoRoute), None)
+            .with_federation(resolver, Arc::new(TestNodeIdentity(iota_b.clone()))),
+    );
+    let deliveries = Arc::new(Deliveries::default());
+    let nodes_b = Arc::new(MemoryNodes::default());
+    nodes_b
+        .nodes
+        .write()
+        .unwrap()
+        .insert(iota_a.node_id().clone(), resolved_node(&iota_a, Vec::new()));
+    let router_b = Arc::new(DirectPeerRouter::default());
+    let manager_b = Arc::new(PeerManager::new(
+        iota_b.clone(),
+        nodes_b.clone(),
+        router_b,
+        relay_b.clone(),
+        deliveries.clone(),
+    ));
+
+    let (certificate, private_key) = certificate();
+    let (endpoint, shutdown, server_task) = start_peer_server(
+        0,
+        certificate.clone(),
+        private_key.clone(),
+        &iota_a,
+        &iota_b,
+        manager_b.clone(),
+    )
+    .await;
+
+    let wrong_nodes = Arc::new(MemoryNodes::default());
+    wrong_nodes.nodes.write().unwrap().insert(
+        iota_b.node_id().clone(),
+        ResolvedNodeIdentity {
+            public_keys: vec![Keyring::generate().public_key_bundle()],
+            ..resolved_node(&iota_b, vec![endpoint.clone()])
+        },
+    );
+    let wrong_manager = Arc::new(PeerManager::new(
+        iota_a.clone(),
+        wrong_nodes,
+        Arc::new(DirectPeerRouter::default()),
+        relay_b.clone(),
+        Arc::new(Deliveries::default()),
+    ));
+    assert!(
+        wrong_manager
+            .connect(
+                iota_b.node_id().clone(),
+                &endpoint,
+                Some(certificate.clone()),
+            )
+            .await
+            .is_err()
+    );
+
+    let nodes_a = Arc::new(MemoryNodes::default());
+    nodes_a.nodes.write().unwrap().insert(
+        iota_b.node_id().clone(),
+        resolved_node(&iota_b, vec![endpoint.clone()]),
+    );
+    nodes_a
+        .pin_endpoint_certificate(&endpoint, &certificate)
+        .unwrap();
+    let router_a = Arc::new(DirectPeerRouter::default());
+    let manager_a = Arc::new(PeerManager::new(
+        iota_a.clone(),
+        nodes_a.clone(),
+        router_a.clone(),
+        relay_b,
+        Arc::new(Deliveries::default()),
+    ));
+    router_a.attach_manager(&manager_a);
+    manager_a.connect_best(iota_b.node_id()).await.unwrap();
+
+    let content = CommunicationValue::new(CommunicationType::MessageSend)
+        .add_typed_default(DataType::Content, DataValue::Str("ciphertext".into()))
+        .add_typed_default(DataType::SendTime, DataValue::SignedNumber(10))
+        .add_typed_default(DataType::VersionNumber, DataValue::SignedNumber(1));
+    let frame = FederatedRelayV2::sign(
+        principal_a.descriptor().principal.clone(),
+        principal_b.descriptor().principal.clone(),
+        iota_b.node_id().clone(),
+        "a-to-b".into(),
+        10,
+        content,
+        &user_a,
+    )
+    .unwrap()
+    .into_frame(7)
+    .unwrap();
+    let outcome = tokio::time::timeout(
+        std::time::Duration::from_secs(5),
+        router_a.route(&RouteDestination::Iota(iota_b.node_id().clone()), frame),
+    )
+    .await
+    .unwrap()
+    .unwrap();
+    assert!(matches!(
+        outcome,
+        RouteOutcome::Accepted {
+            ref relay_message_id,
+            ..
+        } if relay_message_id == "a-to-b"
+    ));
+
+    let messages = chat_files::get_messages_for_principal(1, handle_a, 0, 10).unwrap();
+    assert_eq!(messages.len(), 1);
+    assert_eq!(messages[0].content, "ciphertext");
+    {
+        let delivered = deliveries.0.lock().unwrap();
+        assert_eq!(delivered.len(), 1);
+        assert_eq!(delivered[0].0, handle_b);
+        assert!(delivered[0].1.is_type(CommunicationType::MessageLive));
+    }
+
+    shutdown.send(()).unwrap();
+    server_task.await.unwrap();
+    tokio::time::timeout(std::time::Duration::from_secs(5), async {
+        while router_a.has_connection(iota_b.node_id()) {
+            tokio::task::yield_now().await;
+        }
+    })
+    .await
+    .unwrap();
+
+    let restarted_relay_b = Arc::new(
+        RelayService::new(
+            Arc::new(HostedRecipient {
+                principal: handle_b,
+            }),
+            Arc::new(NoRoute),
+            None,
+        )
+        .with_federation(
+            Arc::new(StoreResolver),
+            Arc::new(TestNodeIdentity(iota_b.clone())),
+        ),
+    );
+    let restarted_manager_b = Arc::new(PeerManager::new(
+        iota_b.clone(),
+        nodes_b,
+        Arc::new(DirectPeerRouter::default()),
+        restarted_relay_b,
+        deliveries,
+    ));
+    let (restarted_endpoint, shutdown, server_task) = start_peer_server(
+        0,
+        certificate.clone(),
+        private_key,
+        &iota_a,
+        &iota_b,
+        restarted_manager_b,
+    )
+    .await;
+    nodes_a
+        .nodes
+        .write()
+        .unwrap()
+        .get_mut(iota_b.node_id())
+        .unwrap()
+        .direct_endpoints = vec![restarted_endpoint.clone()];
+    nodes_a
+        .pin_endpoint_certificate(&restarted_endpoint, &certificate)
+        .unwrap();
+    let router_a = Arc::new(DirectPeerRouter::default());
+    let restarted_relay_a = Arc::new(
+        RelayService::new(
+            Arc::new(HostedRecipient {
+                principal: handle_a,
+            }),
+            Arc::new(NoRoute),
+            None,
+        )
+        .with_federation(
+            Arc::new(StoreResolver),
+            Arc::new(TestNodeIdentity(iota_a.clone())),
+        ),
+    );
+    let manager_a = Arc::new(PeerManager::new(
+        iota_a.clone(),
+        nodes_a,
+        router_a.clone(),
+        restarted_relay_a,
+        Arc::new(Deliveries::default()),
+    ));
+    router_a.attach_manager(&manager_a);
+
+    let second = FederatedRelayV2::sign(
+        principal_a.descriptor().principal.clone(),
+        principal_b.descriptor().principal.clone(),
+        iota_b.node_id().clone(),
+        "a-to-b-after-reconnect".into(),
+        11,
+        CommunicationValue::new(CommunicationType::MessageSend)
+            .add_typed_default(DataType::Content, DataValue::Str("second".into()))
+            .add_typed_default(DataType::SendTime, DataValue::SignedNumber(11))
+            .add_typed_default(DataType::VersionNumber, DataValue::SignedNumber(1)),
+        &user_a,
+    )
+    .unwrap()
+    .into_frame(8)
+    .unwrap();
+    let outcome = router_a
+        .route(&RouteDestination::Iota(iota_b.node_id().clone()), second)
+        .await
+        .unwrap();
+    assert!(matches!(
+        outcome,
+        RouteOutcome::Accepted {
+            ref relay_message_id,
+            ..
+        } if relay_message_id == "a-to-b-after-reconnect"
+    ));
+    assert_eq!(
+        chat_files::get_messages_for_principal(1, handle_a, 0, 10)
+            .unwrap()
+            .len(),
+        2
+    );
+
+    shutdown.send(()).unwrap();
+    server_task.await.unwrap();
+}

other-iota/tests/relay_router.rs

@@ -0,0 +1,344 @@
+use async_trait::async_trait;
+use iota_connection::connection_handler::{PeerRouter, RouteDestination, RouteError, RouteOutcome};
+use iota_connection::federated_relay::FederatedRelayV2;
+use iota_connection::relay_service::{RelayNodeIdentity, RelayService};
+use iota_identity::{
+    AuthorityLocator, IdentityError, IdentityResolver, IotaNodeId, LocalNodeIdentity,
+    LocalUserDescriptor, LocalUserId, LocalUserStore, PrincipalDescriptor, PrincipalHandle,
+    PrincipalHome, PrincipalId, PrincipalStore, PublicKeyBundle, ResolutionContext,
+    ResolvedPrincipal, SignedPrincipalDescriptor, UserAddress,
+};
+use iota_storage::identity::SqlitePrincipalStore;
+use iota_storage::util::{chat_files, e2ee_storage};
+use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue};
+use mtp::crypto::Keyring;
+use mtp::host::{AuthenticationPolicy, HostConfig};
+use mtp::webserver::{MTPWebServer, WebServerConfig};
+use other_iota::LocalDeliverySink;
+use other_iota::relay_router::{RelayRouterClient, RelayRouterService};
+use std::net::{IpAddr, Ipv4Addr};
+use std::sync::{Arc, Mutex};
+
+struct StoreResolver;
+
+#[async_trait]
+impl IdentityResolver for StoreResolver {
+    async fn resolve_address(
+        &self,
+        _: &UserAddress,
+        _: &ResolutionContext,
+    ) -> Result<ResolvedPrincipal, IdentityError> {
+        Err(IdentityError::NotFound)
+    }
+
+    async fn resolve_principal(
+        &self,
+        principal: &PrincipalId,
+    ) -> Result<ResolvedPrincipal, IdentityError> {
+        SqlitePrincipalStore
+            .get_by_canonical_id(principal)?
+            .ok_or(IdentityError::NotFound)
+    }
+
+    async fn signing_keys(
+        &self,
+        principal: &PrincipalId,
+        _: &ResolutionContext,
+    ) -> Result<Vec<PublicKeyBundle>, IdentityError> {
+        SqlitePrincipalStore.signing_keys(principal)
+    }
+}
+
+struct HostedRecipient(PrincipalHandle);
+
+impl LocalUserStore for HostedRecipient {
+    fn get_local_user(&self, _: LocalUserId) -> Result<Option<LocalUserDescriptor>, IdentityError> {
+        Ok(None)
+    }
+
+    fn get_local_user_by_username(
+        &self,
+        _: &str,
+    ) -> Result<Option<LocalUserDescriptor>, IdentityError> {
+        Ok(None)
+    }
+
+    fn is_hosted_here(&self, user: LocalUserId) -> Result<bool, IdentityError> {
+        Ok(user == LocalUserId(1))
+    }
+
+    fn local_user_for_principal(
+        &self,
+        principal: PrincipalHandle,
+    ) -> Result<Option<LocalUserId>, IdentityError> {
+        Ok((principal == self.0).then_some(LocalUserId(1)))
+    }
+
+    fn principal_for_local_user(
+        &self,
+        user: LocalUserId,
+    ) -> Result<Option<PrincipalHandle>, IdentityError> {
+        Ok((user == LocalUserId(1)).then_some(self.0))
+    }
+}
+
+struct TestNodeIdentity(LocalNodeIdentity);
+
+#[async_trait]
+impl RelayNodeIdentity for TestNodeIdentity {
+    async fn keyring(&self) -> Option<Arc<Keyring>> {
+        Some(self.0.keyring())
+    }
+
+    fn node_id(&self) -> Option<IotaNodeId> {
+        Some(self.0.node_id().clone())
+    }
+}
+
+struct NoRoute;
+
+#[async_trait]
+impl PeerRouter for NoRoute {
+    async fn route(
+        &self,
+        destination: &RouteDestination,
+        _: CommunicationValue,
+    ) -> Result<RouteOutcome, RouteError> {
+        match destination {
+            RouteDestination::Iota(node) => Err(RouteError::NoRoute(node.clone())),
+            RouteDestination::LegacyOmegaIota { .. } => Ok(RouteOutcome::Retryable {
+                reason: "legacy route is unavailable".into(),
+            }),
+        }
+    }
+}
+
+#[derive(Default)]
+struct Deliveries(Mutex<Vec<(PrincipalHandle, CommunicationValue)>>);
+
+#[async_trait]
+impl LocalDeliverySink for Deliveries {
+    async fn deliver(&self, recipient: PrincipalHandle, frame: CommunicationValue) {
+        self.0.lock().unwrap().push((recipient, frame));
+    }
+}
+
+fn verified_principal(
+    authority: &LocalNodeIdentity,
+    user: &Keyring,
+) -> iota_identity::VerifiedPrincipalDescriptor {
+    let now = iota_storage::util::sync::now_millis();
+    SignedPrincipalDescriptor::sign(
+        PrincipalDescriptor {
+            principal: PrincipalId {
+                authority: authority.authority_id().clone(),
+                user_id: 1,
+            },
+            authority_kind: iota_identity::AuthorityKind::Iota,
+            username: Some("alice".into()),
+            display_name: None,
+            public_keys: vec![user.public_key_bundle()],
+            home: PrincipalHome::Iota(authority.node_id().clone()),
+            revision: 1,
+            valid_until: Some(now + 60_000),
+            issued_at: now,
+        },
+        &authority.keyring(),
+    )
+    .unwrap()
+    .verify(
+        authority.authority_id(),
+        &authority.public_keys(),
+        None,
+        now,
+    )
+    .unwrap()
+}
+
+fn certificate() -> (Vec<u8>, Vec<u8>) {
+    let key_pair = rcgen::KeyPair::generate().unwrap();
+    let params =
+        rcgen::CertificateParams::new(vec!["localhost".into(), "127.0.0.1".into()]).unwrap();
+    let certificate = params.self_signed(&key_pair).unwrap();
+    (
+        certificate.pem().into_bytes(),
+        key_pair.serialize_pem().into_bytes(),
+    )
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+async fn authenticated_router_forwards_opaque_relay_v2() {
+    let storage = tempfile::tempdir().unwrap();
+    iota_util::file_util::configure_storage_directory(storage.path().to_owned());
+    iota_storage::util::db::initialize_database().unwrap();
+
+    let iota_a = LocalNodeIdentity::from_keyring(Keyring::generate()).unwrap();
+    let iota_b = LocalNodeIdentity::from_keyring(Keyring::generate()).unwrap();
+    let user_a = Keyring::generate();
+    let user_b = Keyring::generate();
+    let principal_a = verified_principal(&iota_a, &user_a);
+    let principal_b = verified_principal(&iota_b, &user_b);
+    let handle_a = SqlitePrincipalStore
+        .upsert_remote_descriptor(&principal_a)
+        .unwrap();
+    let handle_b = SqlitePrincipalStore
+        .upsert_remote_descriptor(&principal_b)
+        .unwrap();
+    let chat_id = e2ee_storage::principal_chat_id(handle_a, handle_b).unwrap();
+    e2ee_storage::put_chat_secret(e2ee_storage::StoredChatSecret {
+        user_id: "1".into(),
+        secret_id: e2ee_storage::principal_secret_id(&chat_id),
+        chat_id,
+        version: 1,
+        encrypted_secret: vec![1],
+        kem_ciphertext: vec![2],
+        wrapping_scheme: "test".into(),
+        created_at: 1,
+        updated_at: 1,
+    })
+    .unwrap();
+
+    let relay = Arc::new(
+        RelayService::new(Arc::new(HostedRecipient(handle_b)), Arc::new(NoRoute), None)
+            .with_federation(
+                Arc::new(StoreResolver),
+                Arc::new(TestNodeIdentity(iota_b.clone())),
+            ),
+    );
+    let deliveries = Arc::new(Deliveries::default());
+    let router = Arc::new(RelayRouterService::default());
+    let router_keyring = Keyring::generate();
+    let router_public_key = router_keyring.public_key_bundle();
+    let key_a = iota_a.public_keys();
+    let key_b = iota_b.public_keys();
+    let hint_a = iota_a.node_id().mtp_authentication_hint();
+    let hint_b = iota_b.node_id().mtp_authentication_hint();
+    let (certificate, private_key) = certificate();
+    let host = HostConfig::new(
+        IpAddr::V4(Ipv4Addr::LOCALHOST),
+        0,
+        certificate.clone(),
+        private_key,
+    )
+    .with_authentication(
+        router_keyring,
+        Box::new(move |client_id, _| {
+            let key = if client_id == hint_a {
+                Some(key_a.clone())
+            } else if client_id == hint_b {
+                Some(key_b.clone())
+            } else {
+                None
+            };
+            Box::pin(async move { key })
+        }),
+        Box::new(|_, _| Box::pin(async { 0 })),
+    )
+    .with_authentication_policy(AuthenticationPolicy::ForceAuthentication);
+    let mut server = MTPWebServer::new(host, WebServerConfig::new())
+        .await
+        .unwrap();
+    let endpoint = AuthorityLocator::new(server.local_addr().to_string()).unwrap();
+    let server_task = {
+        let router = router.clone();
+        tokio::spawn(async move {
+            loop {
+                match server.accept().await {
+                    Ok(Some(connection)) => {
+                        let router = router.clone();
+                        tokio::spawn(async move {
+                            web_server::MtpConnectionHandler::accept(router.as_ref(), connection)
+                                .await;
+                        });
+                    }
+                    Ok(None) => break,
+                    Err(_) => continue,
+                }
+            }
+        })
+    };
+
+    let iota_a_router = RelayRouterClient::connect(
+        iota_a.clone(),
+        &endpoint,
+        certificate.clone(),
+        router_public_key.clone(),
+        relay.clone(),
+        deliveries.clone(),
+    )
+    .await
+    .unwrap();
+    tokio::time::timeout(std::time::Duration::from_secs(5), async {
+        while !router.has_connection(iota_a.node_id()) {
+            tokio::task::yield_now().await;
+        }
+    })
+    .await
+    .unwrap();
+
+    let frame = FederatedRelayV2::sign(
+        principal_a.descriptor().principal.clone(),
+        principal_b.descriptor().principal.clone(),
+        iota_b.node_id().clone(),
+        "through-router".into(),
+        10,
+        CommunicationValue::new(CommunicationType::MessageSend)
+            .add_typed_default(DataType::Content, DataValue::Str("opaque".into()))
+            .add_typed_default(DataType::SendTime, DataValue::SignedNumber(10))
+            .add_typed_default(DataType::VersionNumber, DataValue::SignedNumber(1)),
+        &user_a,
+    )
+    .unwrap()
+    .into_frame(7)
+    .unwrap();
+    assert!(matches!(
+        iota_a_router
+            .route(
+                &RouteDestination::Iota(iota_b.node_id().clone()),
+                frame.clone(),
+            )
+            .await
+            .unwrap(),
+        RouteOutcome::Retryable { .. }
+    ));
+
+    let _iota_b_router = RelayRouterClient::connect(
+        iota_b.clone(),
+        &endpoint,
+        certificate,
+        router_public_key,
+        relay,
+        deliveries.clone(),
+    )
+    .await
+    .unwrap();
+    tokio::time::timeout(std::time::Duration::from_secs(5), async {
+        while !router.has_connection(iota_b.node_id()) {
+            tokio::task::yield_now().await;
+        }
+    })
+    .await
+    .unwrap();
+    let outcome = iota_a_router
+        .route(&RouteDestination::Iota(iota_b.node_id().clone()), frame)
+        .await
+        .unwrap();
+    assert!(
+        matches!(
+        &outcome,
+        RouteOutcome::Accepted {
+            relay_message_id,
+            ..
+        } if relay_message_id == "through-router"
+        ),
+        "unexpected router outcome: {outcome:?}"
+    );
+    let messages = chat_files::get_messages_for_principal(1, handle_a, 0, 10).unwrap();
+    assert_eq!(messages.len(), 1);
+    assert_eq!(messages[0].content, "opaque");
+    let delivered = deliveries.0.lock().unwrap();
+    assert_eq!(delivered.len(), 1);
+    assert_eq!(delivered[0].0, handle_b);
+
+    server_task.abort();
+}