dnscrypt 0.1.0

//! `reqwest` drop-in DNS resolver backed by DNSCrypt.
//!
//! [`DnscryptResolver`] implements [`reqwest::dns::Resolve`], allowing it
//! to be plugged directly into any `reqwest::Client`:
//!
//! ```no_run
//! use dnscrypt::DnscryptResolver;
//! use std::sync::Arc;
//!
//! let resolver = Arc::new(DnscryptResolver::new());
//! let client = reqwest::Client::builder()
//!     .dns_resolver(resolver)
//!     .build()
//!     .unwrap();
//! ```
//!
//! The resolver **caches the DNSCrypt session** across lookups.  A new session
//! is established automatically when the cached one expires (per the certificate
//! `ts_end` field).

use std::net::SocketAddr;
use std::sync::{Arc, Mutex, PoisonError};

use reqwest::dns::{Name, Resolve, Resolving};

use crate::resolver::{DnscryptClientState, HARDCODED_RESOLVERS, HardcodedResolver};

#[cfg(not(feature = "tokio"))]
use crate::resolver::{establish_dnscrypt_session, resolve_domain_via_dnscrypt_session};

/// A `reqwest`-compatible DNS resolver that routes every lookup through DNSCrypt.
///
/// Construct one and pass it to `reqwest::ClientBuilder::dns_resolver`:
///
/// ```no_run
/// # use dnscrypt::DnscryptResolver;
/// # use std::sync::Arc;
/// let resolver = Arc::new(DnscryptResolver::new());
/// let client = reqwest::Client::builder()
///     .dns_resolver(resolver)
///     .build()
///     .unwrap();
/// ```
///
/// The resolver is `Clone`-safe; all clones share the same cached session via
/// an `Arc<Mutex<…>>`.
#[derive(Debug)]
pub struct DnscryptResolver {
    session: Arc<Mutex<Option<Vec<DnscryptClientState>>>>,
    resolvers: &'static [HardcodedResolver],
}

impl Clone for DnscryptResolver {
    fn clone(&self) -> Self {
        Self {
            session: self.session.clone(),
            resolvers: self.resolvers,
        }
    }
}

impl DnscryptResolver {
    /// Create a new resolver with default Quad9 resolvers and no pre-established session.
    pub fn new() -> Self {
        Self {
            session: Arc::new(Mutex::new(None)),
            resolvers: HARDCODED_RESOLVERS,
        }
    }

    /// Create a new resolver with a custom static list of resolvers and no pre-established session.
    pub fn new_with_resolvers(resolvers: &'static [HardcodedResolver]) -> Self {
        Self {
            session: Arc::new(Mutex::new(None)),
            resolvers,
        }
    }

    /// Return the cached session if it is still valid, otherwise establish a
    /// fresh one and cache it (synchronous version).
    #[cfg(not(feature = "tokio"))]
    fn get_or_establish_session(&self) -> Result<Vec<DnscryptClientState>, String> {
        let mut lock = self.session.lock().unwrap_or_else(PoisonError::into_inner);

        if let Some(ref s_vec) = *lock {
            if let Some(s) = s_vec.first() {
                let now = std::time::SystemTime::now()
                    .duration_since(std::time::UNIX_EPOCH)
                    .unwrap_or_default()
                    .as_secs();
                if now < s.ts_end as u64 {
                    return Ok(s_vec.clone());
                }
            }
        }

        let new_sessions = establish_dnscrypt_session(self.resolvers).map_err(|e| e.to_string())?;
        *lock = Some(new_sessions.clone());
        Ok(new_sessions)
    }

    /// Return the cached session if it is still valid, otherwise establish a
    /// fresh one and cache it (asynchronous version).
    #[cfg(feature = "tokio")]
    async fn get_or_establish_session_async(&self) -> Result<Vec<DnscryptClientState>, String> {
        {
            let lock = self.session.lock().unwrap_or_else(PoisonError::into_inner);
            if let Some(ref s_vec) = *lock {
                if let Some(s) = s_vec.first() {
                    let now = std::time::SystemTime::now()
                        .duration_since(std::time::UNIX_EPOCH)
                        .unwrap_or_default()
                        .as_secs();
                    if now < s.ts_end as u64 {
                        return Ok(s_vec.clone());
                    }
                }
            }
        }

        let new_sessions = crate::resolver::establish_dnscrypt_session_async(self.resolvers)
            .await
            .map_err(|e| e.to_string())?;

        let mut lock = self.session.lock().unwrap_or_else(PoisonError::into_inner);
        if let Some(ref s_vec) = *lock {
            if let Some(s) = s_vec.first() {
                let now = std::time::SystemTime::now()
                    .duration_since(std::time::UNIX_EPOCH)
                    .unwrap_or_default()
                    .as_secs();
                if now < s.ts_end as u64 {
                    return Ok(s_vec.clone());
                }
            }
        }
        *lock = Some(new_sessions.clone());
        Ok(new_sessions)
    }
}

impl Default for DnscryptResolver {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(not(feature = "tokio"))]
impl Resolve for DnscryptResolver {
    fn resolve(&self, name: Name) -> Resolving {
        let domain = name.as_str().to_string();
        let session_arc = self.session.clone();
        let resolvers = self.resolvers;

        let res = (|| -> Result<
            Box<dyn Iterator<Item = SocketAddr> + Send>,
            Box<dyn std::error::Error + Send + Sync>,
        > {
            let proxy = DnscryptResolver { session: session_arc, resolvers };
            let sessions = proxy.get_or_establish_session().map_err(|e| {
                Box::new(std::io::Error::other(e)) as Box<dyn std::error::Error + Send + Sync>
            })?;

            #[cfg(not(feature = "max-security"))]
            let ips = {
                let mut rb = [0u8; 4];
                getrandom::fill(&mut rb).expect("getrandom failed");
                let start_idx = (u32::from_ne_bytes(rb) as usize) % sessions.len();

                let mut last_err = None;
                let mut resolved_ips = None;
                for i in 0..sessions.len() {
                    let idx = (start_idx + i) % sessions.len();
                    let session = &sessions[idx];
                    match resolve_domain_via_dnscrypt_session(session, &domain) {
                        Ok(ips) => { resolved_ips = Some(ips); break; },
                        Err(e) => last_err = Some(e),
                    }
                }
                resolved_ips.ok_or_else(|| last_err.unwrap_or_else(|| crate::error::Error::Protocol("no sessions".into())))
            };

            #[cfg(feature = "max-security")]
            let ips = {
                let mut results = std::collections::HashMap::new();
                for session in &sessions {
                    if let Ok(resolved_ips) = resolve_domain_via_dnscrypt_session(session, &domain) {
                        for ip in resolved_ips {
                            *results.entry(ip).or_insert(0) += 1;
                        }
                    }
                }
                if results.is_empty() {
                    Err(crate::error::Error::Protocol("max-security: all providers failed".into()))
                } else {
                    let mut max_count = 0;
                    for count in results.values() {
                        if *count > max_count {
                            max_count = *count;
                        }
                    }
                    let mut sorted_ips: Vec<_> = results.into_iter().collect();
                    sorted_ips.sort_by(|a, b| b.1.cmp(&a.1));
                    Ok(sorted_ips
                        .into_iter()
                        .filter(|(_, count)| *count == max_count)
                        .map(|(ip, _)| ip)
                        .collect::<Vec<std::net::IpAddr>>())
                }
            };

            let ips = ips.map_err(|e| {
                Box::new(std::io::Error::new(
                    std::io::ErrorKind::NotFound,
                    e.to_string(),
                )) as Box<dyn std::error::Error + Send + Sync>
            })?;

            Ok(Box::new(ips.into_iter().map(|ip| SocketAddr::new(ip, 0))))
        })();

        Box::pin(std::future::ready(res))
    }
}

#[cfg(feature = "tokio")]
impl Resolve for DnscryptResolver {
    fn resolve(&self, name: Name) -> Resolving {
        let domain = name.as_str().to_string();
        let session_arc = self.session.clone();
        let resolvers = self.resolvers;

        Box::pin(async move {
            let proxy = DnscryptResolver {
                session: session_arc,
                resolvers,
            };
            let sessions = proxy.get_or_establish_session_async().await.map_err(|e| {
                Box::new(std::io::Error::other(e)) as Box<dyn std::error::Error + Send + Sync>
            })?;

            #[cfg(not(feature = "max-security"))]
            let ips = {
                let mut rb = [0u8; 4];
                getrandom::fill(&mut rb).expect("getrandom failed");
                let start_idx = (u32::from_ne_bytes(rb) as usize) % sessions.len();

                let mut last_err = None;
                let mut resolved_ips = None;
                for i in 0..sessions.len() {
                    let idx = (start_idx + i) % sessions.len();
                    let session = &sessions[idx];
                    match crate::resolver::resolve_domain_via_dnscrypt_session_async(
                        session, &domain,
                    )
                    .await
                    {
                        Ok(ips) => {
                            resolved_ips = Some(ips);
                            break;
                        }
                        Err(e) => last_err = Some(e),
                    }
                }
                resolved_ips.ok_or_else(|| {
                    last_err.unwrap_or_else(|| crate::error::Error::Protocol("no sessions".into()))
                })
            };

            #[cfg(feature = "max-security")]
            let ips = {
                let mut results = std::collections::HashMap::new();
                let mut tasks = Vec::new();
                for session in &sessions {
                    let session_clone = session.clone();
                    let domain_clone = domain.clone();
                    tasks.push(tokio::spawn(async move {
                        crate::resolver::resolve_domain_via_dnscrypt_session_async(
                            &session_clone,
                            &domain_clone,
                        )
                        .await
                    }));
                }
                for task in tasks {
                    if let Ok(Ok(resolved_ips)) = task.await {
                        for ip in resolved_ips {
                            *results.entry(ip).or_insert(0) += 1;
                        }
                    }
                }
                if results.is_empty() {
                    Err(crate::error::Error::Protocol(
                        "max-security: all providers failed".into(),
                    ))
                } else {
                    let mut max_count = 0;
                    for count in results.values() {
                        if *count > max_count {
                            max_count = *count;
                        }
                    }
                    let mut sorted_ips: Vec<_> = results.into_iter().collect();
                    sorted_ips.sort_by(|a, b| b.1.cmp(&a.1));
                    Ok(sorted_ips
                        .into_iter()
                        .filter(|(_, count)| *count == max_count)
                        .map(|(ip, _)| ip)
                        .collect::<Vec<std::net::IpAddr>>())
                }
            };

            let ips = ips.map_err(|e| {
                Box::new(std::io::Error::new(
                    std::io::ErrorKind::NotFound,
                    e.to_string(),
                )) as Box<dyn std::error::Error + Send + Sync>
            })?;

            let addrs: Box<dyn Iterator<Item = SocketAddr> + Send> =
                Box::new(ips.into_iter().map(|ip| SocketAddr::new(ip, 0)));
            Ok(addrs)
        })
    }
}