any_dns/
server.rs

1#![allow(unused)]
2
3use std::{net::SocketAddr, sync::mpsc::channel};
4
5use crate::{
6    custom_handler::{CustomHandler, EmptyHandler, HandlerHolder},
7    dns_socket::DnsSocket,
8};
9
10pub struct Builder {
11    icann_resolver: SocketAddr,
12    listen: SocketAddr,
13    handler: HandlerHolder,
14}
15
16impl Builder {
17    pub fn new() -> Self {
18        Self {
19            icann_resolver: SocketAddr::from(([192, 168, 1, 1], 53)),
20            listen: SocketAddr::from(([0, 0, 0, 0], 53)),
21            handler: HandlerHolder::new(EmptyHandler::new()),
22        }
23    }
24
25    /// Set the DNS resolver for normal ICANN domains. Defaults to 192.168.1.1:53
26    pub fn icann_resolver(mut self, icann_resolver: SocketAddr) -> Self {
27        self.icann_resolver = icann_resolver;
28        self
29    }
30
31    /// Set socket the server should listen on. Defaults to 0.0.0.0:53
32    pub fn listen(mut self, listen: SocketAddr) -> Self {
33        self.listen = listen;
34        self
35    }
36
37    /** Set handler to process the dns packet. `Ok()`` should include a dns packet with answers. `Err()` will fallback to ICANN. */
38    pub fn handler(mut self, handler: impl CustomHandler + 'static) -> Self {
39        self.handler = HandlerHolder::new(handler);
40        self
41    }
42
43    // /** Build and start server. */
44    pub async fn build(self) -> tokio::io::Result<AnyDNS> {
45        AnyDNS::new(self.listen, self.icann_resolver, self.handler).await
46    }
47}
48
49#[derive(Debug)]
50pub struct AnyDNS {
51    join_handle: tokio::task::JoinHandle<()>,
52}
53
54impl AnyDNS {
55    pub async fn new(
56        listener: SocketAddr,
57        icann_fallback: SocketAddr,
58        handler: HandlerHolder,
59    ) -> tokio::io::Result<Self> {
60        let mut socket = DnsSocket::new(listener, icann_fallback, handler).await?;
61        let join_handle = tokio::spawn(async move {
62            socket.receive_loop().await;
63        });
64
65        let server = Self { join_handle };
66
67        Ok(server)
68    }
69
70    /**
71     * Stops the server and consumes the instance.
72     */
73    pub fn stop(self) {
74        self.join_handle.abort();
75    }
76
77    /**
78     * Waits on CTRL+C
79     */
80    pub async fn wait_on_ctrl_c(&self) {
81        match tokio::signal::ctrl_c().await {
82            Ok(()) => {}
83            Err(err) => {
84                eprintln!("Unable to listen for shutdown signal Ctrl+C: {}", err);
85            }
86        }
87    }
88}
89
90#[cfg(test)]
91mod tests {
92    use std::{net::SocketAddr, thread::sleep, time::Duration};
93
94    use crate::server::Builder;
95
96    #[tokio::test]
97    async fn run() {
98        let listening: SocketAddr = "0.0.0.0:34255".parse().unwrap();
99        let dns = Builder::new().listen(listening).build().await.unwrap();
100        println!("Started");
101        sleep(Duration::from_secs(5));
102        println!("Stop");
103        dns.stop();
104        println!("Stopped");
105
106        // let mut query = Packet::new_query(0);
107        // let qname = Name::new("google.ch").unwrap();
108        // let qtype = simple_dns::QTYPE::TYPE(simple_dns::TYPE::A);
109        // let qclass = simple_dns::QCLASS::CLASS(simple_dns::CLASS::IN);
110        // let question = Question::new(qname, qtype, qclass, false);
111        // query.questions = vec![question];
112
113        // let query = query.build_bytes_vec_compressed().unwrap();
114        // let to: SocketAddr = "8.8.8.8:53".parse().unwrap();
115        // let result = socket.request(&query, &to, Duration::from_secs(5)).await.unwrap();
116        // let reply = Packet::parse(&result).unwrap();
117        // dbg!(reply);
118    }
119}
any_dns/server.rs

any_dns/
server.rs
