dnscrypt 0.1.1

//! Core `DNSCrypt` v2 session management and resolution logic.
//!
//! This module owns:
//! - The hardcoded resolver list with pinned Ed25519 public keys.
//! - Session establishment (certificate fetch → X25519 key agreement).
//! - Encrypted query dispatch and response decryption.
//!
//! Both synchronous and asynchronous variations are supported. Async variants
//! require the `tokio` feature.

use std::net::{IpAddr, Ipv4Addr, SocketAddr};

use std::collections::HashMap;
use std::sync::{Mutex, OnceLock};
use x25519_dalek::{PublicKey as X25519PublicKey, StaticSecret};

use zeroize::{Zeroize, ZeroizeOnDrop};

use crate::cert::{DnscryptCert, verify_cert};
use crate::crypto::{
    run_hchacha20, xchacha20_djb_poly1305_decrypt, xchacha20_djb_poly1305_encrypt,
};
use crate::error::Error;
use crate::net::send_dns_query;
use crate::packet::{
    build_a_record_query, build_txt_record_query, pad_query, parse_dns_response,
    parse_txt_response, unpad_response, validate_domain,
};

// ── Constants ─────────────────────────────────────────────────────────────────

/// Magic bytes expected at the start of every `DNSCrypt` resolver response.
const RESOLVER_MAGIC: [u8; 8] = [0x72, 0x36, 0x66, 0x6e, 0x76, 0x57, 0x6a, 0x38];

/// Minimum padded query length in bytes (`DNSCrypt` spec recommends ≥ 256 for
/// queries and ≥ 1024 for the initial padding target).
const QUERY_MIN_PAD_LEN: usize = 1024;

/// In the default (non `max-security`) build, [`establish_dnscrypt_session`]
/// stops opening new sessions once it has this many, instead of always
/// contacting every hardcoded resolver. This keeps session establishment
/// latency low while still leaving fallback resolvers available if the
/// randomly chosen session fails at query time. The `max-security` feature
/// always establishes sessions with every distinct provider, since it needs
/// them all for majority-vote aggregation.
#[cfg(not(feature = "max-security"))]
const SESSION_ESTABLISH_TARGET: usize = 3;

// ── Shared selection helpers ─────────────────────────────────────────────────

/// Pick a random starting index into a slice of length `len`, used to spread
/// load across sessions/resolvers instead of always preferring the first one.
#[cfg(not(feature = "max-security"))]
fn random_start_index(len: usize) -> usize {
    let rb: [u8; 4] = crate::rng::random_bytes();
    (u32::from_ne_bytes(rb) as usize) % len
}

/// Reduce a per-IP vote count (as gathered by the `max-security` aggregation
/// across multiple providers) down to the IP(s) with the highest count.
#[cfg(feature = "max-security")]
pub(crate) fn majority_ips(results: HashMap<IpAddr, u32>) -> Vec<IpAddr> {
    let max_count = results.values().copied().max().unwrap_or(0);
    let mut ips: Vec<IpAddr> = results
        .into_iter()
        .filter(|(_, count)| *count == max_count)
        .map(|(ip, _)| ip)
        .collect();
    ips.sort();
    ips
}

// ── Resolver registry & Cache ────────────────────────────────────────────────

fn get_serial_cache() -> &'static Mutex<HashMap<[u8; 32], u32>> {
    static CACHE: OnceLock<Mutex<HashMap<[u8; 32], u32>>> = OnceLock::new();
    CACHE.get_or_init(|| Mutex::new(HashMap::new()))
}

// ── Resolver registry ─────────────────────────────────────────────────────────

/// A statically configured `DNSCrypt` resolver with its provider public key pinned
/// at compile time.
#[derive(Clone, Copy, Debug)]
pub struct HardcodedResolver {
    /// IP address and port for `DNSCrypt` traffic (typically port 8443).
    pub ip: SocketAddr,
    /// Provider name used to query the certificate TXT record.
    pub provider_name: &'static str,
    /// Pinned Ed25519 provider public key used to verify the certificate.
    pub provider_pk: [u8; 32],
}

/// The built-in list of trusted `DNSCrypt` resolvers.
pub const HARDCODED_RESOLVERS: &[HardcodedResolver] = &[
    // Quad9 (1)
    HardcodedResolver {
        ip: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(9, 9, 9, 9)), 8443),
        provider_name: "2.dnscrypt-cert.quad9.net",
        provider_pk: [
            0x67, 0xc8, 0x47, 0xb8, 0xc8, 0x75, 0x8c, 0xd1, 0x20, 0x24, 0x55, 0x43, 0xbe, 0x75,
            0x67, 0x46, 0xdf, 0x34, 0xdf, 0x1d, 0x84, 0xc0, 0x0b, 0x8c, 0x47, 0x03, 0x68, 0xdf,
            0x82, 0x1d, 0x86, 0x3e,
        ],
    },
    // Quad9 (2)
    HardcodedResolver {
        ip: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(149, 112, 112, 112)), 8443),
        provider_name: "2.dnscrypt-cert.quad9.net",
        provider_pk: [
            0x67, 0xc8, 0x47, 0xb8, 0xc8, 0x75, 0x8c, 0xd1, 0x20, 0x24, 0x55, 0x43, 0xbe, 0x75,
            0x67, 0x46, 0xdf, 0x34, 0xdf, 0x1d, 0x84, 0xc0, 0x0b, 0x8c, 0x47, 0x03, 0x68, 0xdf,
            0x82, 0x1d, 0x86, 0x3e,
        ],
    },
    // Quad9 (3)
    HardcodedResolver {
        ip: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(149, 112, 112, 9)), 8443),
        provider_name: "2.dnscrypt-cert.quad9.net",
        provider_pk: [
            0x67, 0xc8, 0x47, 0xb8, 0xc8, 0x75, 0x8c, 0xd1, 0x20, 0x24, 0x55, 0x43, 0xbe, 0x75,
            0x67, 0x46, 0xdf, 0x34, 0xdf, 0x1d, 0x84, 0xc0, 0x0b, 0x8c, 0x47, 0x03, 0x68, 0xdf,
            0x82, 0x1d, 0x86, 0x3e,
        ],
    },
    // cs-norway
    HardcodedResolver {
        ip: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(91, 219, 215, 227)), 443),
        provider_name: "2.dnscrypt-cert.cryptostorm.is",
        provider_pk: [
            0x31, 0x33, 0x72, 0xad, 0x59, 0x56, 0x32, 0xc2, 0x41, 0x6b, 0x87, 0x2f, 0x09, 0x8f,
            0x85, 0x1b, 0xdd, 0xb9, 0x65, 0x28, 0x4c, 0x6c, 0xbe, 0x9a, 0x4f, 0x19, 0x09, 0x64,
            0x30, 0xdb, 0xa9, 0x5a,
        ],
    },
    // dorriere-fr
    HardcodedResolver {
        ip: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(90, 46, 206, 248)), 443),
        provider_name: "2.dnscrypt-cert.dnscrypt-recursive-to-root-udp-only.nicolas-dorriere.fr",
        provider_pk: [
            0x19, 0x62, 0xa8, 0x25, 0xde, 0x11, 0x77, 0xa8, 0xba, 0xcd, 0x58, 0x45, 0xad, 0x53,
            0xdc, 0x02, 0x17, 0x32, 0x96, 0x26, 0x2b, 0xe0, 0xb3, 0xe3, 0x26, 0x8a, 0xeb, 0x60,
            0x32, 0x42, 0x23, 0x9c,
        ],
    },
];

// ── Session ───────────────────────────────────────────────────────────────────

/// An active `DNSCrypt` session with a specific resolver.
#[derive(Clone, Debug, Zeroize, ZeroizeOnDrop)]
pub struct DnscryptClientState {
    /// Resolver address this session is bound to.
    #[zeroize(skip)]
    pub resolver_ip: SocketAddr,
    /// Client ephemeral X25519 public key (sent in every query packet).
    pub client_pk_bytes: [u8; 32],
    /// 8-byte magic prefix from the resolver certificate.
    pub client_magic: [u8; 8],
    /// Symmetric session key derived via `HChaCha20` from the DH shared secret.
    pub shared_key: [u8; 32],
    /// Certificate expiry (Unix timestamp).
    #[zeroize(skip)]
    pub ts_end: u32,
}

// ── Shared session-establishment helpers ─────────────────────────────────────

/// Pick the highest-serial valid certificate out of a raw TXT response.
fn best_cert_from_txt_response(resp_data: &[u8], provider_pk: &[u8; 32]) -> Option<DnscryptCert> {
    parse_txt_response(resp_data)
        .into_iter()
        .filter_map(|raw| verify_cert(&raw, provider_pk))
        .reduce(|a, b| if b.serial > a.serial { b } else { a })
}

/// Reject a certificate whose serial number is not newer than the highest one
/// already seen for this provider key, and record `cert`'s serial otherwise.
///
/// This defends against a compromised or misbehaving resolver rolling back to
/// a previously superseded (and potentially leaked) certificate.
fn check_and_record_serial(provider_pk: [u8; 32], cert: &DnscryptCert) -> Result<(), String> {
    let mut cache = get_serial_cache()
        .lock()
        .unwrap_or_else(std::sync::PoisonError::into_inner);
    if let Some(&highest) = cache.get(&provider_pk) {
        if cert.serial < highest {
            return Err(format!("cert serial {} < cached {highest}", cert.serial));
        }
    }
    cache.insert(provider_pk, cert.serial);
    drop(cache);
    Ok(())
}

/// Derive a fresh `DNSCrypt` session (ephemeral X25519 keypair + `HChaCha20`
/// shared key) against `resolver`'s verified `cert`.
fn derive_session(resolver: &HardcodedResolver, cert: &DnscryptCert) -> DnscryptClientState {
    let secret: [u8; 32] = crate::rng::random_bytes();
    let secret_key = StaticSecret::from(secret);
    let public_key = X25519PublicKey::from(&secret_key);

    let mut client_pk_bytes = [0u8; 32];
    client_pk_bytes.copy_from_slice(public_key.as_bytes());

    let peer_pk = X25519PublicKey::from(cert.resolver_pk);
    let shared_secret = secret_key.diffie_hellman(&peer_pk);
    let shared_key = run_hchacha20(shared_secret.as_bytes(), &[0u8; 16]);

    DnscryptClientState {
        resolver_ip: resolver.ip,
        client_pk_bytes,
        client_magic: cert.client_magic,
        shared_key,
        ts_end: cert.ts_end,
    }
}

// ── Public Sync API ──────────────────────────────────────────────────────────

/// Establish a `DNSCrypt` session against the first available resolver in the list.
///
/// # Errors
///
/// Returns [`Error::CertificateError`] if every resolver in `resolvers` fails
/// to produce a valid, unexpired certificate (network failure, signature
/// failure, or a replayed/stale serial number).
pub fn establish_dnscrypt_session(
    resolvers: &'static [HardcodedResolver],
) -> Result<Vec<DnscryptClientState>, Error> {
    let mut errors = Vec::new();
    let mut active_sessions = Vec::new();
    let mut seen_pks = std::collections::HashSet::new();

    for resolver in resolvers {
        #[cfg(not(feature = "max-security"))]
        if active_sessions.len() >= SESSION_ESTABLISH_TARGET {
            break;
        }

        if seen_pks.contains(&resolver.provider_pk) {
            continue;
        }

        let txid: [u8; 2] = crate::rng::random_bytes();
        let txt_query = build_txt_record_query(resolver.provider_name, txid);
        let resp_data = match send_dns_query(resolver.ip, &txt_query, false) {
            Ok(d) => d,
            Err(e) => {
                errors.push(format!("{}: {e}", resolver.ip));
                continue;
            }
        };

        let Some(cert) = best_cert_from_txt_response(&resp_data, &resolver.provider_pk) else {
            errors.push(format!("{}: no valid certificate", resolver.ip));
            continue;
        };

        if let Err(reason) = check_and_record_serial(resolver.provider_pk, &cert) {
            errors.push(format!("{}: {reason}", resolver.ip));
            continue;
        }

        active_sessions.push(derive_session(resolver, &cert));
        seen_pks.insert(resolver.provider_pk);
    }

    if active_sessions.is_empty() {
        Err(Error::CertificateError(format!(
            "all resolvers failed: {}",
            errors.join("; ")
        )))
    } else {
        Ok(active_sessions)
    }
}

/// Send encrypted DNS A and AAAA queries using an existing `session`.
///
/// # Errors
///
/// Returns [`Error::Protocol`] if `domain` is not a valid DNS name, or
/// [`Error::NoRecord`] if neither the A nor the AAAA query returned any
/// address.
pub fn resolve_domain_via_dnscrypt_session(
    session: &DnscryptClientState,
    domain: &str,
) -> Result<Vec<IpAddr>, Error> {
    validate_domain(domain).map_err(Error::Protocol)?;

    let mut all_ips = Vec::new();

    // A Query
    let txid_a: [u8; 2] = crate::rng::random_bytes();
    let query_a = build_a_record_query(domain, txid_a);
    if let Ok(ips) = query_via_dnscrypt_session(session, &query_a, txid_a) {
        all_ips.extend(ips);
    }

    // AAAA Query
    let txid_aaaa: [u8; 2] = crate::rng::random_bytes();
    let query_aaaa = crate::packet::build_aaaa_record_query(domain, txid_aaaa);
    if let Ok(ips) = query_via_dnscrypt_session(session, &query_aaaa, txid_aaaa) {
        all_ips.extend(ips);
    }

    if all_ips.is_empty() {
        return Err(Error::NoRecord);
    }
    Ok(all_ips)
}

fn decrypt_response(
    session: &DnscryptClientState,
    response_packet: &[u8],
    client_nonce: &[u8],
) -> Result<Vec<u8>, Error> {
    if response_packet.len() < 32 {
        return Err(Error::Protocol("response packet too short".to_string()));
    }
    if response_packet[..8] != RESOLVER_MAGIC {
        return Err(Error::Protocol("invalid resolver magic".to_string()));
    }

    let mut resp_nonce = [0u8; 24];
    resp_nonce.copy_from_slice(&response_packet[8..32]);
    #[cfg(feature = "zeroize")]
    let resp_nonce_z = zeroize::Zeroizing::new(resp_nonce);
    #[cfg(not(feature = "zeroize"))]
    let resp_nonce_z = resp_nonce;

    if resp_nonce_z[..12] != client_nonce[..] {
        return Err(Error::Protocol(
            "client nonce mismatch in resolver response".to_string(),
        ));
    }

    let encrypted_response = &response_packet[32..];
    #[cfg(feature = "zeroize")]
    let rn_ref = &*resp_nonce_z;
    #[cfg(not(feature = "zeroize"))]
    let rn_ref = &resp_nonce_z;
    let decrypted = xchacha20_djb_poly1305_decrypt(&session.shared_key, rn_ref, encrypted_response)
        .map_err(|_| Error::AuthenticationFailure)?;
    
    Ok(decrypted)
}

fn query_via_dnscrypt_session(
    session: &DnscryptClientState,
    client_query: &[u8],
    txid: [u8; 2],
) -> Result<Vec<IpAddr>, Error> {
    #[cfg(feature = "zeroize")]
    let client_query = zeroize::Zeroizing::new(client_query.to_vec());
    #[cfg(feature = "zeroize")]
    let padded_query = pad_query(&client_query, QUERY_MIN_PAD_LEN);
    #[cfg(not(feature = "zeroize"))]
    let padded_query = pad_query(client_query, QUERY_MIN_PAD_LEN);
    #[cfg(feature = "zeroize")]
    let padded_query = zeroize::Zeroizing::new(padded_query);

    let client_nonce: [u8; 12] = crate::rng::random_bytes();
    #[cfg(feature = "zeroize")]
    let client_nonce_z = zeroize::Zeroizing::new(client_nonce);
    #[cfg(not(feature = "zeroize"))]
    let client_nonce_z = client_nonce;

    let mut query_nonce = [0u8; 24];
    query_nonce[..12].copy_from_slice(&client_nonce_z[..]);
    #[cfg(feature = "zeroize")]
    let query_nonce_z = zeroize::Zeroizing::new(query_nonce);
    #[cfg(not(feature = "zeroize"))]
    let query_nonce_z = query_nonce;

    #[cfg(feature = "zeroize")]
    let qn_ref = &*query_nonce_z;
    #[cfg(not(feature = "zeroize"))]
    let qn_ref = &query_nonce_z;
    let encrypted_query =
        xchacha20_djb_poly1305_encrypt(&session.shared_key, qn_ref, &padded_query);
    #[cfg(feature = "zeroize")]
    let encrypted_query = zeroize::Zeroizing::new(encrypted_query);

    let mut request_packet = Vec::with_capacity(8 + 32 + 12 + encrypted_query.len());
    request_packet.extend_from_slice(&session.client_magic);
    request_packet.extend_from_slice(&session.client_pk_bytes);
    request_packet.extend_from_slice(&client_nonce_z[..]);
    request_packet.extend_from_slice(&encrypted_query);
    #[cfg(feature = "zeroize")]
    let request_packet = zeroize::Zeroizing::new(request_packet);

    let response_packet =
        send_dns_query(session.resolver_ip, &request_packet, false).map_err(Error::Network)?;
    #[cfg(feature = "zeroize")]
    let response_packet = zeroize::Zeroizing::new(response_packet);

    let decrypted = decrypt_response(session, &response_packet, &client_nonce_z[..])?;
    #[cfg(feature = "zeroize")]
    let decrypted = zeroize::Zeroizing::new(decrypted);

    let mut unpadded = unpad_response(&decrypted).map_err(Error::Protocol)?;

    if unpadded.len() >= 3 && (unpadded[2] & 0x02) != 0 {
        // Response is truncated! Fallback to TCP.
        let tcp_response_packet =
            send_dns_query(session.resolver_ip, &request_packet, true).map_err(Error::Network)?;
        #[cfg(feature = "zeroize")]
        let tcp_response_packet = zeroize::Zeroizing::new(tcp_response_packet);

        let decrypted_tcp = decrypt_response(session, &tcp_response_packet, &client_nonce_z[..])?;
        #[cfg(feature = "zeroize")]
        let decrypted_tcp = zeroize::Zeroizing::new(decrypted_tcp);

        unpadded = unpad_response(&decrypted_tcp).map_err(Error::Protocol)?;
    }
    #[cfg(feature = "zeroize")]
    let unpadded = zeroize::Zeroizing::new(unpadded);

    if unpadded.len() < 2 || unpadded[0..2] != txid {
        return Err(Error::Protocol("inner txid mismatch".to_string()));
    }

    let ips = parse_dns_response(&unpadded);
    if ips.is_empty() {
        return Err(Error::NoRecord);
    }
    Ok(ips)
}

// ── High-level resolution with caching ───────────────────────────────────────
//
// These wrappers back the convenience entry points (`resolve`, `resolve_async`,
// and `DnscryptResolver`). They consult the fixed-TTL answer cache first and,
// on a miss, round-robin across the available sessions starting from a random
// index so repeated calls don't always hammer the same resolver.

/// Resolve `domain` using one of `sessions`, consulting and populating the
/// shared answer cache. Used by the non-`max-security` convenience APIs.
#[cfg(not(feature = "max-security"))]
pub(crate) fn resolve_with_cache(
    sessions: &[DnscryptClientState],
    domain: &str,
) -> Result<Vec<IpAddr>, Error> {
    if let Some(ips) = crate::cache::get(domain) {
        return Ok(ips);
    }

    let start_idx = random_start_index(sessions.len());
    let mut last_err = None;
    for i in 0..sessions.len() {
        let session = &sessions[(start_idx + i) % sessions.len()];
        match resolve_domain_via_dnscrypt_session(session, domain) {
            Ok(ips) => {
                crate::cache::put(domain, ips.clone());
                return Ok(ips);
            }
            Err(e) => last_err = Some(e),
        }
    }
    Err(last_err.unwrap_or_else(|| Error::Protocol("no sessions available".into())))
}

/// Async counterpart of [`resolve_with_cache`].
#[cfg(all(feature = "tokio", not(feature = "max-security")))]
pub(crate) async fn resolve_with_cache_async(
    sessions: &[DnscryptClientState],
    domain: &str,
) -> Result<Vec<IpAddr>, Error> {
    if let Some(ips) = crate::cache::get(domain) {
        return Ok(ips);
    }

    let start_idx = random_start_index(sessions.len());
    let mut last_err = None;
    for i in 0..sessions.len() {
        let session = &sessions[(start_idx + i) % sessions.len()];
        match resolve_domain_via_dnscrypt_session_async(session, domain).await {
            Ok(ips) => {
                crate::cache::put(domain, ips.clone());
                return Ok(ips);
            }
            Err(e) => last_err = Some(e),
        }
    }
    Err(last_err.unwrap_or_else(|| Error::Protocol("no sessions available".into())))
}

// ── Public Async API ─────────────────────────────────────────────────────────

/// Establish a `DNSCrypt` session against the first available resolver in the list asynchronously.
///
/// Only available when the `tokio` feature is enabled.
///
/// # Errors
///
/// Returns [`Error::CertificateError`] if every resolver in `resolvers` fails
/// to produce a valid, unexpired certificate (network failure, signature
/// failure, or a replayed/stale serial number).
#[cfg(feature = "tokio")]
pub async fn establish_dnscrypt_session_async(
    resolvers: &'static [HardcodedResolver],
) -> Result<Vec<DnscryptClientState>, Error> {
    let mut errors = Vec::new();
    let mut active_sessions = Vec::new();
    let mut seen_pks = std::collections::HashSet::new();

    for resolver in resolvers {
        #[cfg(not(feature = "max-security"))]
        if active_sessions.len() >= SESSION_ESTABLISH_TARGET {
            break;
        }

        if seen_pks.contains(&resolver.provider_pk) {
            continue;
        }

        let txid: [u8; 2] = crate::rng::random_bytes();
        let txt_query = build_txt_record_query(resolver.provider_name, txid);
        let resp_data = match crate::net::send_dns_query_async(resolver.ip, &txt_query, false).await
        {
            Ok(d) => d,
            Err(e) => {
                errors.push(format!("{}: {e}", resolver.ip));
                continue;
            }
        };

        let Some(cert) = best_cert_from_txt_response(&resp_data, &resolver.provider_pk) else {
            errors.push(format!("{}: no valid certificate", resolver.ip));
            continue;
        };

        if let Err(reason) = check_and_record_serial(resolver.provider_pk, &cert) {
            errors.push(format!("{}: {reason}", resolver.ip));
            continue;
        }

        active_sessions.push(derive_session(resolver, &cert));
        seen_pks.insert(resolver.provider_pk);
    }

    if active_sessions.is_empty() {
        Err(Error::CertificateError(format!(
            "all resolvers failed: {}",
            errors.join("; ")
        )))
    } else {
        Ok(active_sessions)
    }
}

/// Send encrypted DNS A and AAAA queries using an existing `session` asynchronously.
///
/// Only available when the `tokio` feature is enabled.
///
/// # Errors
///
/// Returns [`Error::Protocol`] if `domain` is not a valid DNS name, or
/// [`Error::NoRecord`] if neither the A nor the AAAA query returned any
/// address.
#[cfg(feature = "tokio")]
pub async fn resolve_domain_via_dnscrypt_session_async(
    session: &DnscryptClientState,
    domain: &str,
) -> Result<Vec<IpAddr>, Error> {
    validate_domain(domain).map_err(Error::Protocol)?;

    let txid_a: [u8; 2] = crate::rng::random_bytes();
    let query_a = build_a_record_query(domain, txid_a);

    let txid_aaaa: [u8; 2] = crate::rng::random_bytes();
    let query_aaaa = crate::packet::build_aaaa_record_query(domain, txid_aaaa);

    // We can run both queries concurrently
    let (res_a, res_aaaa) = tokio::join!(
        query_via_dnscrypt_session_async(session, &query_a, txid_a),
        query_via_dnscrypt_session_async(session, &query_aaaa, txid_aaaa)
    );

    let mut all_ips = Vec::new();
    if let Ok(ips) = res_a {
        all_ips.extend(ips);
    }
    if let Ok(ips) = res_aaaa {
        all_ips.extend(ips);
    }

    if all_ips.is_empty() {
        return Err(Error::NoRecord);
    }
    Ok(all_ips)
}

#[cfg(feature = "tokio")]
async fn query_via_dnscrypt_session_async(
    session: &DnscryptClientState,
    client_query: &[u8],
    txid: [u8; 2],
) -> Result<Vec<IpAddr>, Error> {
    #[cfg(feature = "zeroize")]
    let client_query = zeroize::Zeroizing::new(client_query.to_vec());
    #[cfg(feature = "zeroize")]
    let padded_query = pad_query(&client_query, QUERY_MIN_PAD_LEN);
    #[cfg(not(feature = "zeroize"))]
    let padded_query = pad_query(client_query, QUERY_MIN_PAD_LEN);
    #[cfg(feature = "zeroize")]
    let padded_query = zeroize::Zeroizing::new(padded_query);

    let client_nonce: [u8; 12] = crate::rng::random_bytes();
    #[cfg(feature = "zeroize")]
    let client_nonce_z = zeroize::Zeroizing::new(client_nonce);
    #[cfg(not(feature = "zeroize"))]
    let client_nonce_z = client_nonce;

    let mut query_nonce = [0u8; 24];
    query_nonce[..12].copy_from_slice(&client_nonce_z[..]);
    #[cfg(feature = "zeroize")]
    let query_nonce_z = zeroize::Zeroizing::new(query_nonce);
    #[cfg(not(feature = "zeroize"))]
    let query_nonce_z = query_nonce;

    #[cfg(feature = "zeroize")]
    let qn_ref = &*query_nonce_z;
    #[cfg(not(feature = "zeroize"))]
    let qn_ref = &query_nonce_z;
    let encrypted_query =
        xchacha20_djb_poly1305_encrypt(&session.shared_key, qn_ref, &padded_query);
    #[cfg(feature = "zeroize")]
    let encrypted_query = zeroize::Zeroizing::new(encrypted_query);

    let mut request_packet = Vec::with_capacity(8 + 32 + 12 + encrypted_query.len());
    request_packet.extend_from_slice(&session.client_magic);
    request_packet.extend_from_slice(&session.client_pk_bytes);
    request_packet.extend_from_slice(&client_nonce_z[..]);
    request_packet.extend_from_slice(&encrypted_query);
    #[cfg(feature = "zeroize")]
    let request_packet = zeroize::Zeroizing::new(request_packet);

    let response_packet =
        crate::net::send_dns_query_async(session.resolver_ip, &request_packet, false)
            .await
            .map_err(Error::Network)?;
    #[cfg(feature = "zeroize")]
    let response_packet = zeroize::Zeroizing::new(response_packet);

    let decrypted = decrypt_response(session, &response_packet, &client_nonce_z[..])?;
    #[cfg(feature = "zeroize")]
    let decrypted = zeroize::Zeroizing::new(decrypted);

    let mut unpadded = unpad_response(&decrypted).map_err(Error::Protocol)?;

    if unpadded.len() >= 3 && (unpadded[2] & 0x02) != 0 {
        // Response is truncated! Fallback to TCP.
        let tcp_response_packet =
            crate::net::send_dns_query_async(session.resolver_ip, &request_packet, true)
                .await
                .map_err(Error::Network)?;
        #[cfg(feature = "zeroize")]
        let tcp_response_packet = zeroize::Zeroizing::new(tcp_response_packet);

        let decrypted_tcp = decrypt_response(session, &tcp_response_packet, &client_nonce_z[..])?;
        #[cfg(feature = "zeroize")]
        let decrypted_tcp = zeroize::Zeroizing::new(decrypted_tcp);

        unpadded = unpad_response(&decrypted_tcp).map_err(Error::Protocol)?;
    }
    #[cfg(feature = "zeroize")]
    let unpadded = zeroize::Zeroizing::new(unpadded);

    if unpadded.len() < 2 || unpadded[0..2] != txid {
        return Err(Error::Protocol("inner txid mismatch".to_string()));
    }

    let ips = parse_dns_response(&unpadded);
    if ips.is_empty() {
        return Err(Error::NoRecord);
    }
    Ok(ips)
}

/// Resolve a domain name to an [`IpAddr`] using `DNSCrypt` asynchronously.
///
/// Only available when the `tokio` feature is enabled.
///
/// # Errors
///
/// Returns [`Error`] if no hardcoded resolver can be reached, certificate
/// validation fails for all resolvers, or the DNS response contains no
/// address record.
#[cfg(feature = "tokio")]
pub async fn resolve_async(
    resolvers: &'static [HardcodedResolver],
    domain: &str,
) -> Result<Vec<IpAddr>, Error> {
    let sessions = establish_dnscrypt_session_async(resolvers).await?;

    #[cfg(feature = "max-security")]
    {
        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.to_string();
            tasks.push(tokio::spawn(async move {
                resolve_domain_via_dnscrypt_session_async(&session_clone, &domain_clone).await
            }));
        }

        for task in tasks {
            if let Ok(Ok(ips)) = task.await {
                for ip in ips {
                    *results.entry(ip).or_insert(0) += 1;
                }
            }
        }

        if results.is_empty() {
            return Err(Error::Protocol(
                "max-security: all providers failed to resolve domain".into(),
            ));
        }

        Ok(majority_ips(results))
    }

    #[cfg(not(feature = "max-security"))]
    {
        resolve_with_cache_async(&sessions, domain).await
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::panic)]
mod tests {
    use super::*;
    use std::io::{Read, Write};
    use std::net::{IpAddr, Ipv4Addr, SocketAddr, TcpListener, UdpSocket};
    use std::thread;

    #[test]
    fn test_establish_session_all_fail() {
        static BAD_RESOLVERS: &[HardcodedResolver] = &[HardcodedResolver {
            // Use a local port that is closed, causing UDP/TCP to fail or timeout
            ip: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 65535),
            provider_name: "2.dnscrypt-cert.bad.net",
            provider_pk: [0u8; 32],
        }];

        let result = establish_dnscrypt_session(BAD_RESOLVERS);
        assert!(result.is_err());
        match result {
            Err(Error::CertificateError(msg)) => {
                assert!(msg.contains("all resolvers failed:"));
            }
            _ => panic!("Expected CertificateError"),
        }
    }

    #[test]
    fn test_resolve_domain_via_session_no_record_when_unreachable() {
        // Both the A and AAAA queries fail to reach anything (closed port),
        // so `resolve_domain_via_dnscrypt_session` must fall through to its
        // own NoRecord error rather than propagating the per-query network
        // error or panicking.
        let session = DnscryptClientState {
            resolver_ip: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 65535),
            client_pk_bytes: [0u8; 32],
            client_magic: [0u8; 8],
            shared_key: [0u8; 32],
            ts_end: u32::MAX,
        };
        let result = resolve_domain_via_dnscrypt_session(&session, "example.com");
        assert!(matches!(result, Err(Error::NoRecord)));
    }

    #[test]
    fn test_query_via_session_inner_txid_mismatch() {
        let udp_socket = UdpSocket::bind("127.0.0.1:0").unwrap();
        let port = udp_socket.local_addr().unwrap().port();
        let resolver_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), port);
        let shared_key = [11u8; 32];

        let shared_key_srv = shared_key;
        thread::spawn(move || {
            let mut buf = [0u8; 2048];
            let (_, src) = udp_socket.recv_from(&mut buf).unwrap();
            let client_nonce = &buf[40..52];

            // DNS header transaction ID [0x99, 0x99] deliberately does not
            // match the client's query txid of [0, 0].
            let mut dns_resp = vec![0u8; 12];
            dns_resp[0] = 0x99;
            dns_resp[1] = 0x99;
            let padded_resp = crate::packet::pad_query(&dns_resp, 64);

            let mut resp_nonce = [0u8; 24];
            resp_nonce[..12].copy_from_slice(client_nonce);
            resp_nonce[12..].copy_from_slice(&[1u8; 12]);
            let encrypted =
                crate::crypto::xchacha20_djb_poly1305_encrypt(&shared_key_srv, &resp_nonce, &padded_resp);

            let mut resp_packet = Vec::new();
            resp_packet.extend_from_slice(RESOLVER_MAGIC.as_slice());
            resp_packet.extend_from_slice(&resp_nonce);
            resp_packet.extend_from_slice(&encrypted);
            udp_socket.send_to(&resp_packet, src).unwrap();
        });

        let session = DnscryptClientState {
            resolver_ip: resolver_addr,
            client_pk_bytes: [0u8; 32],
            client_magic: [0u8; 8],
            shared_key,
            ts_end: u32::MAX,
        };
        let query = crate::packet::build_a_record_query("example.com", [0, 0]);
        let err = query_via_dnscrypt_session(&session, &query, [0, 0]).unwrap_err();
        assert!(matches!(err, Error::Protocol(msg) if msg.contains("inner txid mismatch")));
    }

    #[test]
    fn test_check_and_record_serial_rejects_rollback() {
        // A unique provider key so this test doesn't collide with the
        // process-global serial cache used by other tests/resolvers.
        let provider_pk = [42u8; 32];
        let cert_high = DnscryptCert {
            resolver_pk: [0u8; 32],
            client_magic: [0u8; 8],
            serial: 10,
            ts_start: 0,
            ts_end: 0,
        };
        assert!(check_and_record_serial(provider_pk, &cert_high).is_ok());

        let cert_low = DnscryptCert {
            serial: 5,
            ..cert_high
        };
        let err = check_and_record_serial(provider_pk, &cert_low).unwrap_err();
        assert!(err.contains("cert serial 5 < cached 10"));

        // A newer-or-equal serial is accepted and updates the cache.
        let cert_higher = DnscryptCert {
            serial: 11,
            ..cert_high
        };
        assert!(check_and_record_serial(provider_pk, &cert_higher).is_ok());
    }

    fn dummy_session(shared_key: [u8; 32]) -> DnscryptClientState {
        DnscryptClientState {
            resolver_ip: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0),
            client_pk_bytes: [0u8; 32],
            client_magic: [0u8; 8],
            shared_key,
            ts_end: u32::MAX,
        }
    }

    #[test]
    fn test_decrypt_response_rejects_short_packet() {
        let session = dummy_session([1u8; 32]);
        let err = decrypt_response(&session, &[0u8; 31], &[0u8; 12]).unwrap_err();
        assert!(matches!(err, Error::Protocol(msg) if msg.contains("too short")));
    }

    #[test]
    fn test_decrypt_response_rejects_bad_magic() {
        let session = dummy_session([1u8; 32]);
        let packet = [0u8; 40]; // all zero, not RESOLVER_MAGIC
        let err = decrypt_response(&session, &packet, &[0u8; 12]).unwrap_err();
        assert!(matches!(err, Error::Protocol(msg) if msg.contains("invalid resolver magic")));
    }

    #[test]
    fn test_decrypt_response_rejects_nonce_mismatch() {
        let session = dummy_session([1u8; 32]);
        let mut packet = vec![0u8; 40];
        packet[..8].copy_from_slice(&RESOLVER_MAGIC);
        // Nonce prefix (bytes 8..20) left as zero, but we claim a different
        // client nonce, so the check must fail before decryption is attempted.
        let err = decrypt_response(&session, &packet, &[9u8; 12]).unwrap_err();
        assert!(matches!(err, Error::Protocol(msg) if msg.contains("nonce mismatch")));
    }

    #[test]
    fn test_decrypted_tc_bit_fallback() {
        let udp_socket = UdpSocket::bind("127.0.0.1:0").unwrap();
        let port = udp_socket.local_addr().unwrap().port();
        let tcp_listener = TcpListener::bind(("127.0.0.1", port)).unwrap();
        let resolver_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), port);

        let shared_key = [7u8; 32];

        // Spawn UDP server thread
        let shared_key_udp = shared_key;
        thread::spawn(move || {
            let mut buf = [0u8; 2048];
            let (_, src) = udp_socket.recv_from(&mut buf).unwrap();
            
            // Query packet structure:
            // [client_magic (8)] [client_pk (32)] [client_nonce (12)] [encrypted_query]
            let client_nonce = &buf[40..52];

            // Build truncated encrypted response
            // DNS response header with TC bit set:
            // Transaction ID (2 bytes) matched to the query's inner transaction ID (we'll just use [0, 0])
            // Flags (2 bytes) = [0x02, 0x00] (TC bit set)
            let mut dns_resp = vec![0u8; 12];
            dns_resp[2] = 0x02; // TC bit set

            let padded_resp = crate::packet::pad_query(&dns_resp, 64);

            let mut resp_nonce = [0u8; 24];
            resp_nonce[..12].copy_from_slice(client_nonce);
            resp_nonce[12..].copy_from_slice(&[9u8; 12]);

            let encrypted = crate::crypto::xchacha20_djb_poly1305_encrypt(&shared_key_udp, &resp_nonce, &padded_resp);

            let mut resp_packet = Vec::new();
            resp_packet.extend_from_slice(RESOLVER_MAGIC.as_slice());
            resp_packet.extend_from_slice(&resp_nonce);
            resp_packet.extend_from_slice(&encrypted);

            udp_socket.send_to(&resp_packet, src).unwrap();
        });

        // Spawn TCP server thread
        let shared_key_tcp = shared_key;
        thread::spawn(move || {
            let (mut stream, _) = tcp_listener.accept().unwrap();

            // Read query length (2 bytes)
            let mut len_buf = [0u8; 2];
            stream.read_exact(&mut len_buf).unwrap();
            let query_len = u16::from_be_bytes(len_buf) as usize;
            
            let mut query_buf = vec![0u8; query_len];
            stream.read_exact(&mut query_buf).unwrap();

            let client_nonce = &query_buf[40..52];

            // Build valid non-truncated response with an IP address (140.82.121.4)
            // Header: Transaction ID = [0, 0], Flags = [0x81, 0x80] (standard response), QDCOUNT = 1, ANCOUNT = 1
            let mut response = Vec::new();
            response.extend_from_slice(&[
                0x00, 0x00, 0x81, 0x80, 0x00, 0x01, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00,
            ]);
            // QNAME: "github.com"
            response.extend_from_slice(&[6]);
            response.extend_from_slice(b"github");
            response.extend_from_slice(&[3]);
            response.extend_from_slice(b"com");
            response.extend_from_slice(&[0]);
            // QTYPE = A, QCLASS = IN
            response.extend_from_slice(&[0x00, 0x01, 0x00, 0x01]);
            // Answer: Name offset 12, Type = A, Class = IN, TTL = 60, RDLength = 4, RData = 140.82.121.4
            response.extend_from_slice(&[0xc0, 0x0c]);
            response.extend_from_slice(&[0x00, 0x01, 0x00, 0x01]);
            response.extend_from_slice(&[0x00, 0x00, 0x00, 0x3c]);
            response.extend_from_slice(&[0x00, 0x04]);
            response.extend_from_slice(&[140, 82, 121, 4]);

            let padded_resp = crate::packet::pad_query(&response, 64);

            let mut resp_nonce = [0u8; 24];
            resp_nonce[..12].copy_from_slice(client_nonce);
            resp_nonce[12..].copy_from_slice(&[10u8; 12]);

            let encrypted = crate::crypto::xchacha20_djb_poly1305_encrypt(&shared_key_tcp, &resp_nonce, &padded_resp);

            let mut resp_packet = Vec::new();
            resp_packet.extend_from_slice(RESOLVER_MAGIC.as_slice());
            resp_packet.extend_from_slice(&resp_nonce);
            resp_packet.extend_from_slice(&encrypted);

            // Write length-prefixed response
            let resp_len = u16::try_from(resp_packet.len()).unwrap().to_be_bytes();
            stream.write_all(&resp_len).unwrap();
            stream.write_all(&resp_packet).unwrap();
        });

        // Run client query
        let session = DnscryptClientState {
            resolver_ip: resolver_addr,
            client_pk_bytes: [0u8; 32],
            client_magic: [0u8; 8],
            shared_key,
            ts_end: u32::MAX,
        };

        // Query A record for github.com, using transaction ID [0, 0]
        let query = crate::packet::build_a_record_query("github.com", [0, 0]);
        let ips = query_via_dnscrypt_session(&session, &query, [0, 0]).unwrap();
        assert_eq!(ips.len(), 1);
        assert_eq!(ips[0], IpAddr::V4(Ipv4Addr::new(140, 82, 121, 4)));
    }

    #[cfg(feature = "tokio")]
    mod async_tests {
        use super::*;

        #[tokio::test]
        async fn test_establish_session_all_fail_async() {
            static BAD_RESOLVERS: &[HardcodedResolver] = &[HardcodedResolver {
                ip: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 65534),
                provider_name: "2.dnscrypt-cert.bad-async.net",
                provider_pk: [0u8; 32],
            }];

            let result = establish_dnscrypt_session_async(BAD_RESOLVERS).await;
            assert!(result.is_err());
            match result {
                Err(Error::CertificateError(msg)) => {
                    assert!(msg.contains("all resolvers failed:"));
                }
                _ => panic!("Expected CertificateError"),
            }
        }

        #[tokio::test]
        async fn test_resolve_domain_via_session_no_record_when_unreachable_async() {
            let session = DnscryptClientState {
                resolver_ip: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 65534),
                client_pk_bytes: [0u8; 32],
                client_magic: [0u8; 8],
                shared_key: [0u8; 32],
                ts_end: u32::MAX,
            };
            let result = resolve_domain_via_dnscrypt_session_async(&session, "example.com").await;
            assert!(matches!(result, Err(Error::NoRecord)));
        }

        #[tokio::test]
        async fn test_query_via_session_inner_txid_mismatch_async() {
            let udp_socket = UdpSocket::bind("127.0.0.1:0").unwrap();
            let port = udp_socket.local_addr().unwrap().port();
            let resolver_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), port);
            let shared_key = [13u8; 32];

            let shared_key_srv = shared_key;
            thread::spawn(move || {
                let mut buf = [0u8; 2048];
                let (_, src) = udp_socket.recv_from(&mut buf).unwrap();
                let client_nonce = &buf[40..52];

                let mut dns_resp = vec![0u8; 12];
                dns_resp[0] = 0x77;
                dns_resp[1] = 0x77;
                let padded_resp = crate::packet::pad_query(&dns_resp, 64);

                let mut resp_nonce = [0u8; 24];
                resp_nonce[..12].copy_from_slice(client_nonce);
                resp_nonce[12..].copy_from_slice(&[2u8; 12]);
                let encrypted = crate::crypto::xchacha20_djb_poly1305_encrypt(
                    &shared_key_srv,
                    &resp_nonce,
                    &padded_resp,
                );

                let mut resp_packet = Vec::new();
                resp_packet.extend_from_slice(RESOLVER_MAGIC.as_slice());
                resp_packet.extend_from_slice(&resp_nonce);
                resp_packet.extend_from_slice(&encrypted);
                udp_socket.send_to(&resp_packet, src).unwrap();
            });

            let session = DnscryptClientState {
                resolver_ip: resolver_addr,
                client_pk_bytes: [0u8; 32],
                client_magic: [0u8; 8],
                shared_key,
                ts_end: u32::MAX,
            };
            let query = crate::packet::build_a_record_query("example.com", [0, 0]);
            let err = query_via_dnscrypt_session_async(&session, &query, [0, 0])
                .await
                .unwrap_err();
            assert!(matches!(err, Error::Protocol(msg) if msg.contains("inner txid mismatch")));
        }

        #[tokio::test]
        async fn test_decrypted_tc_bit_fallback_async() {
            let udp_socket = UdpSocket::bind("127.0.0.1:0").unwrap();
            let port = udp_socket.local_addr().unwrap().port();
            let tcp_listener = TcpListener::bind(("127.0.0.1", port)).unwrap();
            let resolver_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), port);

            let shared_key = [8u8; 32];

            let shared_key_udp = shared_key;
            thread::spawn(move || {
                let mut buf = [0u8; 2048];
                let (_, src) = udp_socket.recv_from(&mut buf).unwrap();
                let client_nonce = &buf[40..52];

                let mut dns_resp = vec![0u8; 12];
                dns_resp[2] = 0x02; // TC bit set
                let padded_resp = crate::packet::pad_query(&dns_resp, 64);

                let mut resp_nonce = [0u8; 24];
                resp_nonce[..12].copy_from_slice(client_nonce);
                resp_nonce[12..].copy_from_slice(&[9u8; 12]);
                let encrypted = crate::crypto::xchacha20_djb_poly1305_encrypt(
                    &shared_key_udp,
                    &resp_nonce,
                    &padded_resp,
                );

                let mut resp_packet = Vec::new();
                resp_packet.extend_from_slice(RESOLVER_MAGIC.as_slice());
                resp_packet.extend_from_slice(&resp_nonce);
                resp_packet.extend_from_slice(&encrypted);
                udp_socket.send_to(&resp_packet, src).unwrap();
            });

            let shared_key_tcp = shared_key;
            thread::spawn(move || {
                let (mut stream, _) = tcp_listener.accept().unwrap();
                let mut len_buf = [0u8; 2];
                stream.read_exact(&mut len_buf).unwrap();
                let query_len = usize::from(u16::from_be_bytes(len_buf));
                let mut query_buf = vec![0u8; query_len];
                stream.read_exact(&mut query_buf).unwrap();
                let client_nonce = &query_buf[40..52];

                let mut response = Vec::new();
                response.extend_from_slice(&[
                    0x00, 0x00, 0x81, 0x80, 0x00, 0x01, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00,
                ]);
                response.extend_from_slice(&[6]);
                response.extend_from_slice(b"github");
                response.extend_from_slice(&[3]);
                response.extend_from_slice(b"com");
                response.extend_from_slice(&[0]);
                response.extend_from_slice(&[0x00, 0x01, 0x00, 0x01]);
                response.extend_from_slice(&[0xc0, 0x0c]);
                response.extend_from_slice(&[0x00, 0x01, 0x00, 0x01]);
                response.extend_from_slice(&[0x00, 0x00, 0x00, 0x3c]);
                response.extend_from_slice(&[0x00, 0x04]);
                response.extend_from_slice(&[140, 82, 121, 4]);
                let padded_resp = crate::packet::pad_query(&response, 64);

                let mut resp_nonce = [0u8; 24];
                resp_nonce[..12].copy_from_slice(client_nonce);
                resp_nonce[12..].copy_from_slice(&[10u8; 12]);
                let encrypted = crate::crypto::xchacha20_djb_poly1305_encrypt(
                    &shared_key_tcp,
                    &resp_nonce,
                    &padded_resp,
                );

                let mut resp_packet = Vec::new();
                resp_packet.extend_from_slice(RESOLVER_MAGIC.as_slice());
                resp_packet.extend_from_slice(&resp_nonce);
                resp_packet.extend_from_slice(&encrypted);

                let resp_len = u16::try_from(resp_packet.len()).unwrap().to_be_bytes();
                stream.write_all(&resp_len).unwrap();
                stream.write_all(&resp_packet).unwrap();
            });

            let session = DnscryptClientState {
                resolver_ip: resolver_addr,
                client_pk_bytes: [0u8; 32],
                client_magic: [0u8; 8],
                shared_key,
                ts_end: u32::MAX,
            };

            let query = crate::packet::build_a_record_query("github.com", [0, 0]);
            let ips = query_via_dnscrypt_session_async(&session, &query, [0, 0])
                .await
                .unwrap();
            assert_eq!(ips.len(), 1);
            assert_eq!(ips[0], IpAddr::V4(Ipv4Addr::new(140, 82, 121, 4)));
        }
    }
}