[Fix] Connection Management

This commit is contained in:
Alex Emmet 2026-09-13 20:58:41 +02:00
commit 3f2ac18333
No known key found for this signature in database
122 changed files with 19970 additions and 5263 deletions

View file

@ -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();
}