axnet/smoltcp_impl/
dns.rs

1use alloc::vec::Vec;
2use axerrno::{AxError, AxResult, ax_err_type};
3use core::net::IpAddr;
4
5use smoltcp::iface::SocketHandle;
6use smoltcp::socket::dns::{self, GetQueryResultError, StartQueryError};
7use smoltcp::wire::DnsQueryType;
8
9use super::{ETH0, SOCKET_SET, SocketSetWrapper};
10
11/// A DNS socket.
12struct DnsSocket {
13    handle: Option<SocketHandle>,
14}
15
16impl DnsSocket {
17    #[allow(clippy::new_without_default)]
18    /// Creates a new DNS socket.
19    pub fn new() -> Self {
20        let socket = SocketSetWrapper::new_dns_socket();
21        let handle = Some(SOCKET_SET.add(socket));
22        Self { handle }
23    }
24
25    #[allow(dead_code)]
26    /// Update the list of DNS servers, will replace all existing servers.
27    pub fn update_servers(self, servers: &[smoltcp::wire::IpAddress]) {
28        SOCKET_SET.with_socket_mut::<dns::Socket, _, _>(self.handle.unwrap(), |socket| {
29            socket.update_servers(servers)
30        });
31    }
32
33    /// Query a address with given DNS query type.
34    pub fn query(&self, name: &str, query_type: DnsQueryType) -> AxResult<Vec<IpAddr>> {
35        // let local_addr = self.local_addr.unwrap_or_else(f);
36        let handle = self.handle.ok_or_else(|| ax_err_type!(InvalidInput))?;
37        let iface = &ETH0.iface;
38        let query_handle = SOCKET_SET
39            .with_socket_mut::<dns::Socket, _, _>(handle, |socket| {
40                socket.start_query(iface.lock().context(), name, query_type)
41            })
42            .map_err(|e| match e {
43                StartQueryError::NoFreeSlot => {
44                    ax_err_type!(ResourceBusy, "socket query() failed: no free slot")
45                }
46                StartQueryError::InvalidName => {
47                    ax_err_type!(InvalidInput, "socket query() failed: invalid name")
48                }
49                StartQueryError::NameTooLong => {
50                    ax_err_type!(InvalidInput, "socket query() failed: too long name")
51                }
52            })?;
53        loop {
54            SOCKET_SET.poll_interfaces();
55            match SOCKET_SET.with_socket_mut::<dns::Socket, _, _>(handle, |socket| {
56                socket.get_query_result(query_handle).map_err(|e| match e {
57                    GetQueryResultError::Pending => AxError::WouldBlock,
58                    GetQueryResultError::Failed => {
59                        ax_err_type!(ConnectionRefused, "socket query() failed")
60                    }
61                })
62            }) {
63                Ok(n) => {
64                    let mut res = Vec::with_capacity(n.capacity());
65                    for ip in n {
66                        res.push(IpAddr::from(ip))
67                    }
68                    return Ok(res);
69                }
70                Err(AxError::WouldBlock) => axtask::yield_now(),
71                Err(e) => return Err(e),
72            }
73        }
74    }
75}
76
77impl Drop for DnsSocket {
78    fn drop(&mut self) {
79        if let Some(handle) = self.handle {
80            SOCKET_SET.remove(handle);
81        }
82    }
83}
84
85/// Public function for DNS query.
86pub fn dns_query(name: &str) -> AxResult<alloc::vec::Vec<IpAddr>> {
87    let socket = DnsSocket::new();
88    socket.query(name, DnsQueryType::A)
89}