1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120

use std::net::SocketAddr;

use adns_proto::{Header, Packet, PacketParseError, Question};
use rand::{thread_rng, Rng};
use thiserror::Error;
use tokio::{
    io::{AsyncReadExt, AsyncWriteExt},
    net::{TcpStream, ToSocketAddrs, UdpSocket},
};

pub struct DnsClient {
    udp: UdpSocket,
}

#[derive(Error, Debug)]
pub enum DnsQueryError {
    #[error("packet ID mismatch")]
    IDMismatch,
    #[error("packet too large >64KB")]
    PacketTooLarge,
    #[error("{0}")]
    IoError(#[from] std::io::Error),
    #[error("dns parse error {0}")]
    PacketParseError(#[from] PacketParseError),
}

impl DnsClient {
    pub async fn new() -> Result<Self, DnsQueryError> {
        Ok(Self {
            udp: UdpSocket::bind("0.0.0.0:0".parse::<SocketAddr>().unwrap()).await?,
        })
    }

    pub async fn query(
        &mut self,
        servers: impl ToSocketAddrs,
        questions: Vec<Question>,
    ) -> Result<Packet, DnsQueryError> {
        let id: u16 = thread_rng().gen();
        let packet = Packet {
            header: Header {
                id,
                recursion_desired: true,
                recursion_available: true,
                ..Default::default()
            },
            questions,
            ..Default::default()
        };
        let serialized = packet.serialize(usize::MAX);
        if serialized.len() > 512 {
            self.query_tcp(&servers, id, &serialized).await
        } else {
            self.udp.send_to(&serialized, &servers).await?;
            let mut response = [0u8; 512];
            let mut size;
            loop {
                size = self.udp.recv(&mut response).await?;
                if size < 2 || u16::from_be_bytes(response[..2].try_into().unwrap()) != id {
                    continue;
                }
                break;
            }
            match Packet::parse(&response[..size]) {
                Ok(packet) => Ok(packet.0),
                Err(PacketParseError::Truncated) => self.query_tcp(&servers, id, &serialized).await,
                Err(e) => Err(e.into()),
            }
        }
    }

    async fn query_tcp(
        &mut self,
        servers: impl ToSocketAddrs,
        id: u16,
        packet: &[u8],
    ) -> Result<Packet, DnsQueryError> {
        let mut client = TcpStream::connect(servers).await?;
        client
            .write_u16(
                packet
                    .len()
                    .try_into()
                    .map_err(|_| DnsQueryError::PacketTooLarge)?,
            )
            .await?;
        client.write_all(packet).await?;
        let len = client.read_u16().await?;
        let mut response = vec![0u8; len as usize];
        client.read_exact(&mut response).await?;

        let packet = Packet::parse(&response)?.0;
        if packet.header.id != id {
            return Err(DnsQueryError::IDMismatch);
        }
        Ok(packet)
    }
}

#[cfg(test)]
mod tests {
    use adns_proto::Type;

    use super::*;

    #[tokio::test]
    async fn test_query() {
        let mut client = DnsClient::new().await.unwrap();
        let response = client
            .query(
                "8.8.8.8:53",
                vec![Question::new(Type::A, "google.com").unwrap()],
            )
            .await
            .unwrap();
        for answer in &response.answers {
            println!("{answer}");
        }
    }
}