diff --git a/src/connect.rs b/src/connect.rs index e944469..6a4d792 100644 --- a/src/connect.rs +++ b/src/connect.rs @@ -8,6 +8,7 @@ use futures::{future::BoxFuture, Stream}; use hyper_util::rt::TokioIo; use rustls::pki_types::PrivateKeyDer; use rustls::{ClientConfig, ServerConfig}; +use std::collections::{BTreeMap, BTreeSet}; use std::io; use std::net::IpAddr; use std::pin::Pin; @@ -97,33 +98,24 @@ impl IntermeshClient { } /// Connect to a mesh name. + /// + /// If the name resolves to multiple IMIDs, the IMID with the minimum + /// value (per `Ord`) is selected deterministically. pub(crate) async fn connect_name( &mut self, name: &Name, port: u16, ) -> Result> { - let imid = self - .te - .derivation() - .name_to_imid - .get(name) - .and_then(|s| s.iter().next()) - .cloned() - .ok_or_else(|| anyhow!("failed to resolve name to IMID: {name}"))?; - + let imid = resolve_name_to_imid(name, &self.te.derivation().name_to_imid)?; self.connect_imid(&imid, port).await } /// Connect to an IMID. + /// + /// If the IMID resolves to multiple IPs, the IP with the minimum value + /// (per `Ord`) is selected deterministically. async fn connect_imid(&mut self, imid: &Imid, port: u16) -> Result> { - // Trust engine is authoritative; bootstrap hint is fallback for bootstrap. - let te_ip = self - .te - .derivation() - .imid_to_ip - .get(imid) - .and_then(|s| s.iter().next()) - .copied(); + let te_ip = resolve_imid_to_ip(imid, &self.te.derivation().imid_to_ip).ok(); let hint_ip = self .bootstrap_hint .take_if(|(hint_imid, _)| hint_imid == imid) @@ -176,6 +168,32 @@ impl Service for IntermeshClient { } } +/// Resolve a mesh name to an IMID, deterministically selecting the minimum +/// IMID (per `Ord`) when the name maps to multiple candidates. +fn resolve_name_to_imid( + name: &Name, + name_to_imid: &BTreeMap>, +) -> Result { + name_to_imid + .get(name) + .and_then(|s| s.first()) + .cloned() + .ok_or_else(|| anyhow!("failed to resolve name to IMID: {name}")) +} + +/// Resolve an IMID to an IP address, deterministically selecting the minimum +/// IP (per `Ord`) when the IMID maps to multiple candidates. +fn resolve_imid_to_ip( + imid: &Imid, + imid_to_ip: &BTreeMap>, +) -> Result { + imid_to_ip + .get(imid) + .and_then(|s| s.first()) + .copied() + .ok_or_else(|| anyhow!("failed to find IP for IMID: {imid}")) +} + /// Creates a stream of TLS-wrapped TCP connections for serving gRPC requests. /// /// This function creates a `TlsStream` that uses the `IntermeshVerifier` to @@ -294,7 +312,7 @@ mod tests { use crate::proto::intermesh::GossipUpdate; use crate::test_utils::TestFixture; use futures::TryStreamExt; - use std::collections::BTreeMap; + use std::collections::{BTreeMap, BTreeSet}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::sync::mpsc; use tokio::time::{timeout, Duration}; @@ -474,6 +492,67 @@ mod tests { server_task.await.assert(); } + #[test] + fn test_resolve_name_to_imid_first() { + let a = ImidKeypair::test_keypair("a").to_imid(); + let b = ImidKeypair::test_keypair("b").to_imid(); + let (min, max) = if a < b { + (a.clone(), b.clone()) + } else { + (b.clone(), a.clone()) + }; + + let map = BTreeMap::from([( + "test.local".parse().assert(), + BTreeSet::from([min.clone(), max]), + )]); + + let result = resolve_name_to_imid(&"test.local".parse().assert(), &map).assert(); + assert_eq!( + result, min, + "should select the minimum IMID per Ord (BTreeSet::first)" + ); + } + + #[test] + fn test_resolve_name_to_imid_not_found() { + let a = ImidKeypair::test_keypair("a").to_imid(); + let map = BTreeMap::from([("test.local".parse().assert(), BTreeSet::from([a]))]); + + let err = resolve_name_to_imid(&"other.local".parse().assert(), &map).unwrap_err(); + assert!( + err.to_string().contains("failed to resolve name to IMID"), + "unexpected error: {err}" + ); + } + + #[test] + fn test_resolve_imid_to_ip_first() { + let imid = ImidKeypair::test_keypair("test").to_imid(); + let ip: IpAddr = "127.0.0.1".parse().assert(); + + let map = BTreeMap::from([(imid.clone(), BTreeSet::from([ip]))]); + + let result = resolve_imid_to_ip(&imid, &map).assert(); + assert_eq!(result, ip); + } + + #[test] + fn test_resolve_imid_to_ip_not_found() { + let imid = ImidKeypair::test_keypair("test").to_imid(); + let other = ImidKeypair::test_keypair("other").to_imid(); + let map = BTreeMap::from([( + imid, + BTreeSet::from(["127.0.0.1".parse::().assert()]), + )]); + + let err = resolve_imid_to_ip(&other, &map).unwrap_err(); + assert!( + err.to_string().contains("failed to find IP for IMID"), + "unexpected error: {err}" + ); + } + #[tokio::test(flavor = "multi_thread")] async fn test_tonic_integration() { timeout(Duration::from_secs(30), async {