Skip to content

Commit

Permalink
fix and test after rebase
Browse files Browse the repository at this point in the history
  • Loading branch information
Frando committed Mar 25, 2024
1 parent 5ce8142 commit 67629cb
Show file tree
Hide file tree
Showing 4 changed files with 70 additions and 8 deletions.
54 changes: 49 additions & 5 deletions iroh-net/src/discovery.rs
Original file line number Diff line number Diff line change
Expand Up @@ -563,7 +563,7 @@ mod test_dns_pkarr {
relay::{RelayMap, RelayMode},
test_utils::{
dns_server::{run_dns_server, Resolver},
run_relay_server,
run_relay_server, run_relay_server_with_pkarr,
},
AddrInfo, MagicEndpoint, NodeAddr,
};
Expand Down Expand Up @@ -639,7 +639,7 @@ mod test_dns_pkarr {
let (nameserver, pkarr_url, state, task) =
spawn_dns_and_pkarr(origin.clone(), cancel.clone()).await?;

let (relay_map, _relay_url, _relay_guard) = run_relay_server().await.unwrap();
let (relay_map, _relay_url, _relay_guard) = run_relay_server().await?;
let ep1 = ep_with_discovery(relay_map.clone(), nameserver, &origin, &pkarr_url).await?;
let ep2 = ep_with_discovery(relay_map, nameserver, &origin, &pkarr_url).await?;

Expand All @@ -655,6 +655,51 @@ mod test_dns_pkarr {
Ok(())
}

#[tokio::test]
async fn relay_pkarr_publish_dns_discover() -> Result<()> {
let _logging_guard = iroh_test::logging::setup();

let cancel = CancellationToken::new();
let origin = "testdns.example".to_string();
let (nameserver, pkarr_url, state, task) =
spawn_dns_and_pkarr(origin.clone(), cancel.clone()).await?;

let (relay_map, _relay_url, _relay_guard) =
run_relay_server_with_pkarr(Some(pkarr_url)).await?;
let ep1 = ep_with_discovery_publish_relay(relay_map.clone(), nameserver, &origin).await?;
let ep2 = ep_with_discovery_publish_relay(relay_map, nameserver, &origin).await?;

// wait until our shared state received the update from pkarr publishing
state.on_update().await;

// we connect only by node id!
let ep2_node_id = ep2.node_id();
let res = ep1.connect(ep2_node_id.into(), TEST_ALPN).await;
assert!(res.is_ok(), "connection established");
cancel.cancel();
task.await??;
Ok(())
}

async fn ep_with_discovery_publish_relay(
relay_map: RelayMap,
nameserver: SocketAddr,
node_origin: &str,
) -> Result<MagicEndpoint> {
let secret_key = SecretKey::generate();
let resolver = dns_resolver(nameserver)?;
let discovery = DnsDiscovery::new(node_origin.to_string());
MagicEndpoint::builder()
.relay_mode(RelayMode::Custom(relay_map))
.secret_key(secret_key)
.dns_resolver(resolver)
.pkarr_announce()
.alpns(vec![TEST_ALPN.to_vec()])
.discovery(Box::new(discovery))
.bind(0)
.await
}

async fn ep_with_discovery(
relay_map: RelayMap,
nameserver: SocketAddr,
Expand All @@ -670,15 +715,14 @@ mod test_dns_pkarr {
pkarr_relay.clone(),
)),
]);
let ep = MagicEndpoint::builder()
MagicEndpoint::builder()
.relay_mode(RelayMode::Custom(relay_map))
.secret_key(secret_key)
.dns_resolver(resolver)
.alpns(vec![TEST_ALPN.to_vec()])
.discovery(Box::new(discovery))
.bind(0)
.await?;
Ok(ep)
.await
}

fn dns_resolver(nameserver: SocketAddr) -> Result<TokioAsyncResolver> {
Expand Down
2 changes: 1 addition & 1 deletion iroh-net/src/magicsock/relay_actor.rs
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@ use tokio_util::sync::CancellationToken;
use tracing::{debug, info, info_span, trace, warn, Instrument};

use crate::{
discovery::pkarr_relay_publish::DEFAULT_PKARR_TTL,
discovery::pkarr_publish::DEFAULT_PKARR_TTL,
dns::node_info::NodeInfo,
key::{PublicKey, PUBLIC_KEY_LENGTH},
relay::{self, http::ClientError, ReceivedMessage, RelayUrl, MAX_PACKET_SIZE},
Expand Down
2 changes: 1 addition & 1 deletion iroh-net/src/relay/server.rs
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@ use tokio_util::sync::CancellationToken;
use tracing::{debug, info_span, trace, warn, Instrument};
use url::Url;

use crate::discovery::pkarr_relay_publish::PkarrRelayClient as PkarrClient;
use crate::discovery::pkarr_publish::PkarrRelayClient as PkarrClient;
use crate::key::{PublicKey, SecretKey, SharedSecret};

use super::{
Expand Down
20 changes: 19 additions & 1 deletion iroh-net/src/test_utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -24,12 +24,30 @@ pub(crate) struct CleanupDropGuard(pub(crate) oneshot::Sender<()>);
///
/// [`MagicEndpoint::connect`]: crate::magic_endpoint::MagicEndpoint
pub(crate) async fn run_relay_server() -> Result<(RelayMap, RelayUrl, CleanupDropGuard)> {
run_relay_server_with_pkarr(None).await
}

/// Runs a relay server with STUN enabled suitable for tests.
///
/// The returned `Url` is the url of the relay server in the returned [`RelayMap`], it
/// is always `Some` as that is how the [`MagicEndpoint::connect`] API expects it.
///
/// [`MagicEndpoint::connect`]: crate::magic_endpoint::MagicEndpoint
pub(crate) async fn run_relay_server_with_pkarr(
pkarr_relay: Option<url::Url>,
) -> Result<(RelayMap, RelayUrl, CleanupDropGuard)> {
let server_key = SecretKey::generate();
let me = server_key.public().fmt_short();
let tls_config = crate::relay::http::make_tls_config();
let server = crate::relay::http::ServerBuilder::new("127.0.0.1:0".parse().unwrap())
.secret_key(Some(server_key))
.tls_config(Some(tls_config))
.tls_config(Some(tls_config));
let server = if let Some(pkarr_relay) = pkarr_relay {
server.pkarr_relay(pkarr_relay)
} else {
server
};
let server = server
.spawn()
.instrument(error_span!("relay server", %me))
.await?;
Expand Down

0 comments on commit 67629cb

Please sign in to comment.