dnsclientx 0.5.0

use async_std::{
    net::{TcpStream, UdpSocket},
    prelude::*,
};
use dnssector::DNS_MAX_COMPRESSED_SIZE;
use std::{io, net::SocketAddr, time::Duration};

pub async fn query_raw_udp(
    local: &SocketAddr,
    upstream: &SocketAddr,
    packet: &[u8],
    timeout: Duration,
) -> io::Result<Vec<u8>> {
    async_std::io::timeout(timeout, async {
        let socket = UdpSocket::bind(local).await?;
        socket.connect(upstream).await?;
        socket.send(packet).await?;
        let mut response = vec![0; DNS_MAX_COMPRESSED_SIZE];
        let response_len = socket.recv(&mut response).await?;
        response.truncate(response_len);
        Ok(response)
    })
    .await
}

pub async fn query_raw_tcp(
    upstream: &SocketAddr,
    packet: &[u8],
    timeout: Duration,
) -> io::Result<Vec<u8>> {
    async_std::io::timeout(timeout, async {
        let mut stream = TcpStream::connect(upstream).await?;
        let _ = stream.set_nodelay(true);
        stream.write_all(packet).await?;
        let mut response_len_bytes = [0u8; 2];
        stream.read_exact(&mut response_len_bytes).await?;
        let response_len =
            ((response_len_bytes[0] as usize) << 8) | (response_len_bytes[1] as usize);
        let mut response = vec![0; response_len];
        stream.read_exact(&mut response).await?;
        Ok(response)
    })
    .await
}