Skip to content

Commit 9e5759d

Browse files
FrandoFrando
committed
cleanup
1 parent 536e3bd commit 9e5759d

File tree

3 files changed

+102
-106
lines changed

3 files changed

+102
-106
lines changed

‎iroh-net/src/discovery.rs

Lines changed: 14 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -559,16 +559,16 @@ mod test_dns_pkarr {
559559

560560
use crate::{
561561
discovery::pkarr_publish,
562-
dns::{
563-
node_info::{lookup_by_id, parse_hickory_node_info_name, NodeInfo},
564-
tests::dns_server::{self, Resolver},
565-
},
562+
dns::node_info::{lookup_by_id, parse_hickory_node_info_name, NodeInfo},
566563
relay::{RelayMap, RelayMode},
567-
test_utils::run_relay_server,
564+
test_utils::{
565+
dns_server::{run_dns_server, Resolver},
566+
run_relay_server,
567+
},
568568
AddrInfo, MagicEndpoint, NodeAddr,
569569
};
570570

571-
use self::state::State;
571+
use self::{pkarr_relay::run_pkarr_relay, state::State};
572572

573573
use super::{dns::DnsDiscovery, ConcurrentDiscovery};
574574

@@ -578,13 +578,13 @@ mod test_dns_pkarr {
578578
let cancel = CancellationToken::new();
579579
let origin = "testdns.example".to_string();
580580
let state = State::new(origin.clone());
581-
let (dns_addr, dns_task) = dns_server::spawn(state.clone(), cancel.clone()).await?;
581+
let (nameserver, dns_task) = run_dns_server(state.clone(), cancel.clone()).await?;
582582

583583
let node_secret = SecretKey::generate();
584584
let (node_info, signed_packet) = generate_node_info(&node_secret);
585585
state.upsert(signed_packet)?;
586586

587-
let resolver = dns_resolver(dns_addr)?;
587+
let resolver = dns_resolver(nameserver)?;
588588
let resolved = lookup_by_id(&resolver, &node_info.node_id, &origin).await?;
589589

590590
assert_eq!(resolved, node_info.into());
@@ -600,7 +600,7 @@ mod test_dns_pkarr {
600600

601601
let cancel = CancellationToken::new();
602602
let origin = "testdns.example".to_string();
603-
let (dns_addr, pkarr_url, _state, task) =
603+
let (nameserver, pkarr_url, _state, task) =
604604
spawn_dns_and_pkarr(origin.clone(), cancel.clone()).await?;
605605

606606
let secret_key = SecretKey::generate();
@@ -613,7 +613,7 @@ mod test_dns_pkarr {
613613
};
614614
publisher.publish_addr_info(&addr_info).await?;
615615

616-
let resolver = dns_resolver(dns_addr)?;
616+
let resolver = dns_resolver(nameserver)?;
617617
let resolved = lookup_by_id(&resolver, &node_id, &origin).await?;
618618

619619
let expected = NodeAddr {
@@ -708,14 +708,14 @@ mod test_dns_pkarr {
708708
cancel: CancellationToken,
709709
) -> Result<(SocketAddr, Url, State, JoinHandle<Result<()>>)> {
710710
let state = State::new(origin);
711-
let (dns_addr, dns_task) = dns_server::spawn(state.clone(), cancel.clone()).await?;
712-
let (pkarr_url, pkarr_task) = pkarr_relay::spawn(state.clone(), cancel.clone()).await?;
711+
let (nameserver, dns_task) = run_dns_server(state.clone(), cancel.clone()).await?;
712+
let (pkarr_url, pkarr_task) = run_pkarr_relay(state.clone(), cancel.clone()).await?;
713713
let join_handle = tokio::task::spawn(async move {
714714
dns_task.await??;
715715
pkarr_task.await??;
716716
Ok(())
717717
});
718-
Ok((dns_addr, pkarr_url, state, join_handle))
718+
Ok((nameserver, pkarr_url, state, join_handle))
719719
}
720720

721721
mod state {
@@ -792,7 +792,6 @@ mod test_dns_pkarr {
792792
const TTL: u32 = 30;
793793
let this = self.clone();
794794
async move {
795-
println!("QUERY {:?}", query.queries());
796795
for query in query.queries() {
797796
let Some(node_id) = parse_hickory_node_info_name(query.name()) else {
798797
continue;
@@ -828,7 +827,7 @@ mod test_dns_pkarr {
828827

829828
use super::State as AppState;
830829

831-
pub async fn spawn(
830+
pub async fn run_pkarr_relay(
832831
state: AppState,
833832
cancel: CancellationToken,
834833
) -> Result<(Url, JoinHandle<Result<()>>)> {

‎iroh-net/src/dns.rs

Lines changed: 0 additions & 91 deletions
Original file line numberDiff line numberDiff line change
@@ -176,95 +176,4 @@ pub(crate) mod tests {
176176
assert!(!res.is_empty());
177177
dbg!(res);
178178
}
179-
180-
pub(crate) mod dns_server {
181-
use std::net::{Ipv4Addr, SocketAddr};
182-
183-
use anyhow::{ensure, Result};
184-
use futures::{future::BoxFuture, Future};
185-
use hickory_proto::{
186-
op::{header::MessageType, Message},
187-
serialize::binary::BinDecodable,
188-
};
189-
use tokio::{net::UdpSocket, task::JoinHandle};
190-
use tokio_util::sync::CancellationToken;
191-
use tracing::{debug, warn};
192-
193-
pub trait Resolver: Send + Sync + 'static {
194-
fn resolve(
195-
&self,
196-
query: &Message,
197-
reply: &mut Message,
198-
) -> impl Future<Output = Result<()>> + Send;
199-
}
200-
201-
pub type ResolveCallback = Box<
202-
dyn Fn(&Message, &mut Message) -> BoxFuture<'static, Result<()>>
203-
+ Send
204-
+ Sync
205-
+ 'static,
206-
>;
207-
impl Resolver for ResolveCallback {
208-
fn resolve(
209-
&self,
210-
query: &Message,
211-
reply: &mut Message,
212-
) -> impl Future<Output = Result<()>> + Send {
213-
(&self)(query, reply)
214-
}
215-
}
216-
217-
pub async fn spawn(
218-
resolver: impl Resolver,
219-
cancel: CancellationToken,
220-
) -> Result<(SocketAddr, JoinHandle<Result<()>>)> {
221-
let bind_addr = SocketAddr::from((Ipv4Addr::LOCALHOST, 0));
222-
let socket = UdpSocket::bind(bind_addr).await?;
223-
let bound_addr = socket.local_addr()?;
224-
let s = TestDnsServer {
225-
socket,
226-
cancel,
227-
resolver,
228-
};
229-
let join_handle = tokio::task::spawn(async move { s.run().await });
230-
Ok((bound_addr, join_handle))
231-
}
232-
233-
struct TestDnsServer<R> {
234-
resolver: R,
235-
socket: UdpSocket,
236-
cancel: CancellationToken,
237-
}
238-
239-
impl<R: Resolver> TestDnsServer<R> {
240-
async fn run(self) -> Result<()> {
241-
let mut buf = [0; 1450];
242-
loop {
243-
tokio::select! {
244-
_ = self.cancel.cancelled() => break,
245-
res = self.socket.recv_from(&mut buf) => {
246-
let (len, from) = res?;
247-
if let Err(err) = self.handle_datagram(from, &buf[..len]).await {
248-
warn!(?err, %from, "failed to handle incoming datagram");
249-
}
250-
}
251-
};
252-
}
253-
Ok(())
254-
}
255-
256-
async fn handle_datagram(&self, from: SocketAddr, buf: &[u8]) -> Result<()> {
257-
let packet = Message::from_bytes(buf)?;
258-
debug!(queries = ?packet.queries(), %from, "received query");
259-
let mut reply = packet.clone();
260-
reply.set_message_type(MessageType::Response);
261-
self.resolver.resolve(&packet, &mut reply).await?;
262-
debug!(?reply, %from, "send reply");
263-
let buf = reply.to_vec()?;
264-
let len = self.socket.send_to(&buf, from).await?;
265-
ensure!(len == buf.len(), "failed to send complete packet");
266-
Ok(())
267-
}
268-
}
269-
}
270179
}

‎iroh-net/src/test_utils.rs

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,3 +62,91 @@ pub(crate) async fn run_relay_server() -> Result<(RelayMap, RelayUrl, CleanupDro
6262

6363
Ok((m, url, CleanupDropGuard(tx)))
6464
}
65+
66+
pub mod dns_server {
67+
use std::net::{Ipv4Addr, SocketAddr};
68+
69+
use anyhow::{ensure, Result};
70+
use futures::{future::BoxFuture, Future};
71+
use hickory_proto::{
72+
op::{header::MessageType, Message},
73+
serialize::binary::BinDecodable,
74+
};
75+
use tokio::{net::UdpSocket, task::JoinHandle};
76+
use tokio_util::sync::CancellationToken;
77+
use tracing::{debug, warn};
78+
79+
pub trait Resolver: Send + Sync + 'static {
80+
fn resolve(
81+
&self,
82+
query: &Message,
83+
reply: &mut Message,
84+
) -> impl Future<Output = Result<()>> + Send;
85+
}
86+
87+
pub type ResolveCallback = Box<
88+
dyn Fn(&Message, &mut Message) -> BoxFuture<'static, Result<()>> + Send + Sync + 'static,
89+
>;
90+
impl Resolver for ResolveCallback {
91+
fn resolve(
92+
&self,
93+
query: &Message,
94+
reply: &mut Message,
95+
) -> impl Future<Output = Result<()>> + Send {
96+
(&self)(query, reply)
97+
}
98+
}
99+
100+
pub async fn run_dns_server(
101+
resolver: impl Resolver,
102+
cancel: CancellationToken,
103+
) -> Result<(SocketAddr, JoinHandle<Result<()>>)> {
104+
let bind_addr = SocketAddr::from((Ipv4Addr::LOCALHOST, 0));
105+
let socket = UdpSocket::bind(bind_addr).await?;
106+
let bound_addr = socket.local_addr()?;
107+
let s = TestDnsServer {
108+
socket,
109+
cancel,
110+
resolver,
111+
};
112+
let join_handle = tokio::task::spawn(async move { s.run().await });
113+
Ok((bound_addr, join_handle))
114+
}
115+
116+
struct TestDnsServer<R> {
117+
resolver: R,
118+
socket: UdpSocket,
119+
cancel: CancellationToken,
120+
}
121+
122+
impl<R: Resolver> TestDnsServer<R> {
123+
async fn run(self) -> Result<()> {
124+
let mut buf = [0; 1450];
125+
loop {
126+
tokio::select! {
127+
_ = self.cancel.cancelled() => break,
128+
res = self.socket.recv_from(&mut buf) => {
129+
let (len, from) = res?;
130+
if let Err(err) = self.handle_datagram(from, &buf[..len]).await {
131+
warn!(?err, %from, "failed to handle incoming datagram");
132+
}
133+
}
134+
};
135+
}
136+
Ok(())
137+
}
138+
139+
async fn handle_datagram(&self, from: SocketAddr, buf: &[u8]) -> Result<()> {
140+
let packet = Message::from_bytes(buf)?;
141+
debug!(queries = ?packet.queries(), %from, "received query");
142+
let mut reply = packet.clone();
143+
reply.set_message_type(MessageType::Response);
144+
self.resolver.resolve(&packet, &mut reply).await?;
145+
debug!(?reply, %from, "send reply");
146+
let buf = reply.to_vec()?;
147+
let len = self.socket.send_to(&buf, from).await?;
148+
ensure!(len == buf.len(), "failed to send complete packet");
149+
Ok(())
150+
}
151+
}
152+
}

0 commit comments

Comments
 (0)