dnscrypt 0.1.0

//! 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,
};

// ── 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;

// ── 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,
        ],
    },
    // DNSCry.pt Moscow
    HardcodedResolver {
        ip: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(93, 183, 106, 222)), 443),
        provider_name: "2.dnscrypt-cert.dnscry.pt",
        provider_pk: [
            0x14, 0x3a, 0xb8, 0x27, 0x1e, 0x45, 0x43, 0x55, 0x24, 0x61, 0x41, 0xd6, 0x41, 0xa5,
            0xb6, 0xe5, 0xbf, 0x26, 0xbd, 0x8d, 0xb2, 0x2c, 0xef, 0x73, 0x4e, 0x81, 0x7c, 0xc2,
            0x37, 0xd2, 0x22, 0x01,
        ],
    },
    // adguard-dns-unfiltered
    HardcodedResolver {
        ip: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(94, 140, 14, 141)), 5443),
        provider_name: "2.dnscrypt.unfiltered.ns1.adguard.com",
        provider_pk: [
            0xb5, 0xe8, 0x44, 0xd6, 0xb8, 0x3a, 0x3e, 0x3e, 0x12, 0x68, 0xeb, 0x68, 0x1f, 0xbe,
            0xa7, 0x0d, 0xe3, 0xc5, 0xfa, 0x68, 0xdb, 0xd6, 0x88, 0x14, 0x1d, 0xbf, 0x2e, 0xc0,
            0x66, 0xb2, 0xda, 0xfa,
        ],
    },
    // 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,
}

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

/// Establish a DNSCrypt session against the first available resolver in the list.
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 {
        if seen_pks.contains(&resolver.provider_pk) {
            continue;
        }

        let mut txid = [0u8; 2];
        getrandom::fill(&mut txid).expect("getrandom failed");
        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 best_cert: Option<DnscryptCert> = parse_txt_response(&resp_data)
            .into_iter()
            .filter_map(|raw| verify_cert(&raw, &resolver.provider_pk))
            .reduce(|a, b| if b.serial > a.serial { b } else { a });

        let cert = match best_cert {
            Some(c) => c,
            None => {
                errors.push(format!("{}: no valid certificate", resolver.ip));
                continue;
            }
        };

        {
            let mut cache = get_serial_cache().lock().unwrap();
            if let Some(&highest) = cache.get(&resolver.provider_pk) {
                if cert.serial < highest {
                    errors.push(format!(
                        "{}: cert serial {} < cached {}",
                        resolver.ip, cert.serial, highest
                    ));
                    continue;
                }
            }
            cache.insert(resolver.provider_pk, cert.serial);
        }

        let mut secret = [0u8; 32];
        getrandom::fill(&mut secret).expect("getrandom failed");
        let client_sk = StaticSecret::from(secret);
        let client_pk = X25519PublicKey::from(&client_sk);

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

        let peer_pk = X25519PublicKey::from(cert.resolver_pk);
        let shared_secret = client_sk.diffie_hellman(&peer_pk);

        let shared_key = run_hchacha20(shared_secret.as_bytes(), &[0u8; 16]);

        active_sessions.push(DnscryptClientState {
            resolver_ip: resolver.ip,
            client_pk_bytes,
            client_magic: cert.client_magic,
            shared_key,
            ts_end: cert.ts_end,
        });

        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`.
pub fn resolve_domain_via_dnscrypt_session(
    session: &DnscryptClientState,
    domain: &str,
) -> Result<Vec<IpAddr>, Error> {
    let mut all_ips = Vec::new();

    // A Query
    let mut txid_a = [0u8; 2];
    getrandom::fill(&mut txid_a).expect("getrandom failed");
    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 mut txid_aaaa = [0u8; 2];
    getrandom::fill(&mut txid_aaaa).expect("getrandom failed");
    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: Vec<u8>,
    txid: [u8; 2],
) -> Result<Vec<IpAddr>, Error> {
    #[cfg(feature = "zeroize")]
    let client_query = zeroize::Zeroizing::new(client_query);
    let padded_query = pad_query(&client_query, QUERY_MIN_PAD_LEN);
    #[cfg(feature = "zeroize")]
    let padded_query = zeroize::Zeroizing::new(padded_query);

    let mut client_nonce = [0u8; 12];
    getrandom::fill(&mut client_nonce).expect("getrandom failed");
    #[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)
}

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

/// Establish a DNSCrypt session against the first available resolver in the list asynchronously.
///
/// Only available when the `tokio` feature is enabled.
#[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 {
        if seen_pks.contains(&resolver.provider_pk) {
            continue;
        }

        let mut txid = [0u8; 2];
        getrandom::fill(&mut txid).expect("getrandom failed");
        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 best_cert: Option<DnscryptCert> = parse_txt_response(&resp_data)
            .into_iter()
            .filter_map(|raw| verify_cert(&raw, &resolver.provider_pk))
            .reduce(|a, b| if b.serial > a.serial { b } else { a });

        let cert = match best_cert {
            Some(c) => c,
            None => {
                errors.push(format!("{}: no valid certificate", resolver.ip));
                continue;
            }
        };

        {
            let mut cache = get_serial_cache().lock().unwrap();
            if let Some(&highest) = cache.get(&resolver.provider_pk) {
                if cert.serial < highest {
                    errors.push(format!(
                        "{}: cert serial {} < cached {}",
                        resolver.ip, cert.serial, highest
                    ));
                    continue;
                }
            }
            cache.insert(resolver.provider_pk, cert.serial);
        }

        let mut secret = [0u8; 32];
        getrandom::fill(&mut secret).expect("getrandom failed");
        let client_sk = StaticSecret::from(secret);
        let client_pk = X25519PublicKey::from(&client_sk);

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

        let peer_pk = X25519PublicKey::from(cert.resolver_pk);
        let shared_secret = client_sk.diffie_hellman(&peer_pk);

        let shared_key = run_hchacha20(shared_secret.as_bytes(), &[0u8; 16]);

        active_sessions.push(DnscryptClientState {
            resolver_ip: resolver.ip,
            client_pk_bytes,
            client_magic: cert.client_magic,
            shared_key,
            ts_end: cert.ts_end,
        });

        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.
#[cfg(feature = "tokio")]
pub async fn resolve_domain_via_dnscrypt_session_async(
    session: &DnscryptClientState,
    domain: &str,
) -> Result<Vec<IpAddr>, Error> {
    let mut txid_a = [0u8; 2];
    getrandom::fill(&mut txid_a).expect("getrandom failed");
    let query_a = build_a_record_query(domain, txid_a);

    let mut txid_aaaa = [0u8; 2];
    getrandom::fill(&mut txid_aaaa).expect("getrandom failed");
    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)
}

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

    let mut client_nonce = [0u8; 12];
    getrandom::fill(&mut client_nonce).expect("getrandom failed");
    #[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.
#[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(),
            ));
        }

        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())
    }

    #[cfg(not(feature = "max-security"))]
    {
        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;
        for i in 0..sessions.len() {
            let idx = (start_idx + i) % sessions.len();
            let session = &sessions[idx];
            match resolve_domain_via_dnscrypt_session_async(session, domain).await {
                Ok(ip) => return Ok(ip),
                Err(e) => last_err = Some(e),
            }
        }
        Err(last_err.unwrap_or_else(|| Error::Protocol("no sessions available".into())))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::net::{IpAddr, Ipv4Addr, SocketAddr};

    #[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::new(127, 0, 0, 1)), 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_decrypted_tc_bit_fallback() {
        use std::net::{TcpListener, UdpSocket};
        use std::thread;

        let udp_socket = UdpSocket::bind("127.0.0.1:0").unwrap();
        let port = udp_socket.local_addr().unwrap().port();
        let tcp_listener = TcpListener::bind(format!("127.0.0.1:{}", port)).unwrap();
        let resolver_addr = SocketAddr::new(
            std::net::IpAddr::V4(std::net::Ipv4Addr::new(127, 0, 0, 1)),
            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)
            use std::io::Read;
            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
            use std::io::Write;
            let resp_len = (resp_packet.len() as u16).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)));
    }
}