Rust DoH DNS over HTTP Resolver for reqwest

doh.rs

//! A DoH client for the sync crate
//!
//! To be used in ClientBuilder.dns_resolver

use once_cell::sync::Lazy;
use std::{
    cell::RefCell,
    collections::HashMap,
    net::SocketAddr,
    time::{Duration, Instant},
};

use hyper::client::connect::dns::Name;
use reqwest::{
    dns::{Addrs, Resolve, Resolving},
    Client,
};
use serde::Deserialize;

static DNS_CACHE: Lazy<DNSCache> = Lazy::new(DNSCache::new);
static DNS_HTTP_CLIENT: Lazy<Client> = Lazy::new(|| {
    Client::builder()
        .timeout(Duration::from_secs(5))
        .http2_prior_knowledge()
        .http2_keep_alive_interval(Duration::from_secs(5))
        .http2_keep_alive_timeout(Duration::from_secs(5))
        .build()
        .unwrap()
});

pub struct DoHResolver;

impl Resolve for DoHResolver {
    fn resolve(&self, name: Name) -> Resolving {
        // also 1.1.1.1
        let url = format!("https://1.12.12.12/dns-query?name={}&type=A", name);
        Box::pin(async move {
            let mut ret = vec![];
            if let Some(mut addrs) = DNS_CACHE.get(&name) {
                ret.append(&mut addrs);
            } else {
                let resp = DNS_HTTP_CLIENT.get(&url).send().await?;
                let resp = resp.json::<DnsResponse>().await?;

                if let Some(ans) = resp.answer {
                    for record in ans {
                        if record.typ == 1 {
                            // A record
                            let addr = SocketAddr::new(record.data.parse().unwrap(), 443);
                            DNS_CACHE.insert(name.clone(), addr, record.ttl);
                            ret.push(addr);
                        }
                    }
                }
                log::debug!("DoH resolved: {}: {:?}", name, ret);
            }

            let addrs: Addrs = Box::new(ret.into_iter());
            Ok(addrs)
        })
    }
}

#[derive(Debug)]
pub struct Entry {
    addr: SocketAddr,
    expiration: Option<Instant>,
}

#[derive(Debug, Default)]
pub struct DNSCache {
    cache: RefCell<HashMap<Name, Vec<Entry>>>,
}

unsafe impl Send for DNSCache {}
unsafe impl Sync for DNSCache {}

impl DNSCache {
    pub fn new() -> Self {
        Self {
            cache: Default::default(),
        }
    }

    // if ttl is expired, remove it from cache
    pub fn get(&self, name: &Name) -> Option<Vec<SocketAddr>> {
        let mut cache = self.cache.borrow_mut();
        let now = Instant::now();
        let mut addrs = Vec::new();
        if let Some(entries) = cache.get_mut(name) {
            entries.retain(|e| {
                if let Some(expiration) = e.expiration {
                    if expiration < now {
                        return false;
                    }
                }
                addrs.push(e.addr);
                true
            });
            if entries.is_empty() {
                cache.remove(name);
            }
        }
        if addrs.is_empty() {
            None
        } else {
            Some(addrs)
        }
    }

    pub fn insert(&self, name: Name, addr: SocketAddr, ttl: Option<u32>) {
        let mut cache = self.cache.borrow_mut();
        let now = Instant::now();
        let expiration = ttl.map(|ttl| now + Duration::from_secs(ttl.into()));
        let entry = Entry { addr, expiration };
        cache.entry(name).or_default().push(entry);
    }
}

#[derive(Deserialize, Debug)]
pub struct DnsResponse {
    #[serde(rename = "Status")]
    pub status: u8,
    #[serde(rename = "TC")]
    pub truncated: bool,

    // "Always true for Google Public DNS"
    #[serde(rename = "RD")]
    pub recursion_desired: bool,
    #[serde(rename = "RA")]
    pub recursion_available: bool,

    #[serde(rename = "AD")]
    pub dnssec_validated: bool,
    #[serde(rename = "CD")]
    pub dnssec_disabled: bool,

    #[serde(rename = "Question")]
    pub question: Vec<DnsQuestion>,

    #[serde(rename = "Answer")]
    pub answer: Option<Vec<DnsAnswer>>,

    #[serde(rename = "Comment")]
    pub comment: Option<String>,
}

#[derive(Deserialize, Debug)]
pub struct DnsQuestion {
    pub name: String,
    #[serde(rename = "type")]
    pub typ: u16,
}

#[derive(Deserialize, Debug)]
pub struct DnsAnswer {
    pub name: String,
    #[serde(rename = "type")]
    pub typ: u16,
    #[serde(rename = "TTL")]
    pub ttl: Option<u32>,
    pub data: String,
}