jive-dns/src/lib.rs

use futures_io::{AsyncRead, AsyncWrite};
use hickory_resolver::config::{ResolverConfig, ResolverOpts};
use std::{
    future::Future,
    pin::Pin,
    task::{Context, Poll},
};

pub fn resolver(config: ResolverConfig, options: ResolverOpts) -> JiveResolver {
    JiveResolver::new(
        config,
        options,
        hickory_resolver::name_server::GenericConnector::new(JiveRuntime),
    )
}

pub fn resolver_from_system_conf() -> Result<JiveResolver, hickory_resolver::error::ResolveError> {
    JiveResolver::from_system_conf(hickory_resolver::name_server::GenericConnector::new(
        JiveRuntime,
    ))
}

#[derive(Clone)]
pub struct JiveRuntime;

pub struct JiveTcpStream {
    pub(crate) io: jive::io::Async<jive::net::TcpStream>,
}

pub struct JiveUdpSocket {
    io: jive::io::Async<jive::net::UdpSocket>,
}

pub type JiveResolver = hickory_resolver::AsyncResolver<JiveConnectionProvider>;

pub type JiveConnection = hickory_resolver::name_server::GenericConnection;
pub type JiveConnectionProvider = hickory_resolver::name_server::GenericConnector<JiveRuntime>;

#[derive(Clone, Copy)]
pub struct JiveRuntimeHandle;

pub struct JiveTimer;

impl JiveTcpStream {
    pub fn into_inner(self) -> jive::io::Async<jive::net::TcpStream> {
        self.io
    }

    pub fn from_inner(io: jive::io::Async<jive::net::TcpStream>) -> Self {
        JiveTcpStream { io }
    }
}

impl hickory_resolver::name_server::RuntimeProvider for JiveRuntime {
    type Handle = JiveRuntimeHandle;
    type Tcp = JiveTcpStream;
    type Timer = JiveTimer;
    type Udp = JiveUdpSocket;

    fn create_handle(&self) -> Self::Handle {
        JiveRuntimeHandle
    }

    fn connect_tcp(
        &self,
        server_addr: std::net::SocketAddr,
    ) -> Pin<Box<dyn Send + Future<Output = std::io::Result<Self::Tcp>>>> {
        Box::pin(async move {
            Ok(JiveTcpStream {
                io: jive::io::Async::<jive::net::TcpStream>::connect(server_addr).await?,
            })
        })
    }

    fn bind_udp(
        &self,
        local_addr: std::net::SocketAddr,
        _server_addr: std::net::SocketAddr,
    ) -> Pin<Box<dyn Send + Future<Output = std::io::Result<Self::Udp>>>> {
        Box::pin(async move {
            Ok(JiveUdpSocket {
                io: jive::io::Async::<jive::net::UdpSocket>::bind(local_addr).await?,
            })
        })
    }
}

impl hickory_resolver::name_server::Spawn for JiveRuntimeHandle {
    fn spawn_bg<F>(&mut self, future: F)
    where
        F: Future<Output = Result<(), hickory_proto::error::ProtoError>> + Send + 'static,
    {
        let _join = jive::spawn(future);
    }
}

#[async_trait::async_trait]
impl hickory_proto::tcp::Connect for JiveTcpStream {
    async fn connect_with_bind(
        addr: std::net::SocketAddr,
        bind_addr: Option<std::net::SocketAddr>,
    ) -> std::io::Result<Self> {
        let io = if let Some(_bind_addr) = bind_addr {
            jive::io::Async::<jive::net::TcpStream>::connect_with_bind(addr, bind_addr).await?
        } else {
            jive::io::Async::<jive::net::TcpStream>::connect(addr).await?
        };

        Ok(JiveTcpStream { io })
    }
}

impl hickory_proto::tcp::DnsTcpStream for JiveTcpStream {
    type Time = JiveTimer;
}

impl AsyncRead for JiveTcpStream {
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut [u8],
    ) -> Poll<std::io::Result<usize>> {
        Pin::new(&mut self.get_mut().io).poll_read(cx, buf)
    }
}

impl AsyncWrite for JiveTcpStream {
    fn poll_write(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<std::io::Result<usize>> {
        Pin::new(&mut self.get_mut().io).poll_write(cx, buf)
    }

    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
        Pin::new(&mut self.get_mut().io).poll_flush(cx)
    }

    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
        Pin::new(&mut self.get_mut().io).poll_close(cx)
    }
}

#[async_trait::async_trait]
impl hickory_proto::udp::DnsUdpSocket for JiveUdpSocket {
    type Time = JiveTimer;

    fn poll_send_to(
        &self,
        cx: &mut Context<'_>,
        buf: &[u8],
        target: std::net::SocketAddr,
    ) -> Poll<std::io::Result<usize>> {
        self.io.poll_send_to(cx, buf, target)
    }

    fn poll_recv_from(
        &self,
        cx: &mut Context<'_>,
        buf: &mut [u8],
    ) -> Poll<std::io::Result<(usize, std::net::SocketAddr)>> {
        self.io.poll_recv_from(cx, buf)
    }
}

#[async_trait::async_trait]
impl hickory_proto::udp::UdpSocket for JiveUdpSocket {
    async fn connect(addr: std::net::SocketAddr) -> std::io::Result<Self> {
        let bind_addr: std::net::SocketAddr = match addr {
            std::net::SocketAddr::V4(_) => (std::net::Ipv4Addr::UNSPECIFIED, 0).into(),
            std::net::SocketAddr::V6(_) => (std::net::Ipv6Addr::UNSPECIFIED, 0).into(),
        };

        Self::bind(bind_addr).await
    }

    async fn connect_with_bind(
        _addr: std::net::SocketAddr,
        bind_addr: std::net::SocketAddr,
    ) -> std::io::Result<Self> {
        Self::bind(bind_addr).await
    }

    async fn bind(addr: std::net::SocketAddr) -> std::io::Result<Self> {
        let io = jive::io::Async::<jive::net::UdpSocket>::bind(addr).await?;

        Ok(JiveUdpSocket { io })
    }
}

#[async_trait::async_trait]
impl hickory_proto::Time for JiveTimer {
    async fn delay_for(duration: std::time::Duration) {
        jive::time::sleep(duration).await
    }

    async fn timeout<F: 'static + Future + Send>(
        duration: std::time::Duration,
        future: F,
    ) -> Result<F::Output, std::io::Error> {
        jive::time::timeout(duration, Box::pin(future)).await
    }
}

impl hickory_proto::Executor for JiveRuntime {
    fn new() -> Self {
        Self
    }

    fn block_on<F: Future>(&mut self, future: F) -> F::Output {
        jive::block_on(future)
    }
}