nomilo/src/dns/client.rs

use std::{future::Future, pin::Pin, task::{Context, Poll}};
use std::ops::{Deref, DerefMut};

use rocket::{Request, State, http::Status, request::{FromRequest, Outcome}};
use tokio::{net::TcpStream as TokioTcpStream, task};
use trust_dns_client::{client::AsyncClient, error::ClientError, op::DnsResponse, tcp::TcpClientStream};
use trust_dns_proto::error::ProtoError;
use trust_dns_proto::iocompat::AsyncIoTokioAsStd;

use crate::config::Config;



pub struct DnsClient(AsyncClient);

impl Deref for DnsClient {
    type Target = AsyncClient;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl DerefMut for DnsClient {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}


#[rocket::async_trait]
impl<'r> FromRequest<'r> for DnsClient {
    type Error = ();
    async fn from_request(request: &'r Request<'_>) -> Outcome<Self, Self::Error> {
        let config = try_outcome!(request.guard::<State<Config>>().await);
        let (stream, handle) = TcpClientStream::<AsyncIoTokioAsStd<TokioTcpStream>>::new(config.dns.server);
        let client = AsyncClient::with_timeout(
            stream,
            handle,
            std::time::Duration::from_secs(5),
            None);
        let (client, bg) = match client.await {
            Err(e) => {
                println!("Failed to connect to DNS server {:#?}", e);
                return Outcome::Failure((Status::InternalServerError, ()))
            },
            Ok(c) => c
        };
        task::spawn(bg);
        Outcome::Success(DnsClient(client))
    }
}

// Reimplement this type here as ClientReponse in trust-dns crate have private fields
pub struct ClientResponse<R>(pub(crate) R)
where
    R: Future<Output = Result<DnsResponse, ProtoError>> + Send + Unpin + 'static;

impl<R> Future for ClientResponse<R>
where
    R: Future<Output = Result<DnsResponse, ProtoError>> + Send + Unpin + 'static,
{
    type Output = Result<DnsResponse, ClientError>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        // This is from the future_utils crate, we simply reuse the reexport from Rocket
        rocket::futures::FutureExt::poll_unpin(&mut self.0, cx).map_err(ClientError::from)
    }
}