<commit_before>use std::{
    io,
    net::{IpAddr, SocketAddr},
    sync::Arc,
    time::Duration,
};

use tokio::sync::mpsc;

#[cfg(not(target_os = "android"))]
use std::net::{Ipv4Addr, SocketAddrV4};
#[cfg(not(target_os = "android"))]
use tokio::net::UdpSocket;

use log::{debug, error, info, warn};
use trust_dns_proto::{
    op::{header::MessageType, response_code::ResponseCode, Message, Query},
    rr::RData,
};

use crate::{
    config::ConfigType,
    context::SharedContext,
    relay::{
        loadbalancing::server::{PlainPingBalancer, ServerType},
        sys::create_udp_socket,
        utils::try_timeout,
    },
};

mod upstream;

async fn acl_lookup<Local, Remote>(
    context: SharedContext,
    local: Arc<Local>,
    remote: Arc<Remote>,
    query: &Query
) -> io::Result<(Message, bool)>
    where
        Local: upstream::Upstream,
        Remote: upstream::Upstream,
{
    // Start querying name servers
    debug!(
        "attempting lookup of {:?} {} with ns {:?} and {:?}",
        query.query_type(), query.name(), local, remote
    );

    let qname_in_proxy_list = context.check_query_in_proxy_list(query);

    let remote_response_fut = async {
        match qname_in_proxy_list {
            Some(false) => None,
            _ => {
                let timeout = Some(Duration::new(3, 0));
                try_timeout(remote.lookup(query), timeout)
                    .await
                    .ok()
            }
        }
    };

    let local_response = match qname_in_proxy_list {
        Some(true) => None,
        _ => {
            let timeout = Some(Duration::new(3, 0));
            try_timeout(local.lookup(query), timeout).await.ok()
        }
    }
    .unwrap_or_else(Message::new);

    match qname_in_proxy_list {
        Some(true) => {
            let remote_response = remote_response_fut.await.unwrap_or_else(Message::new);
            debug!("pick remote response (qname): {:?}", remote_response);
            return Ok((remote_response, true));
        }
        Some(false) => {
            debug!("pick local response (qname): {:?}", local_response);
            return Ok((local_response, false));
        }
        None => (),
    }

    if local_response.answer_count() == 0 {
        let remote_response = remote_response_fut.await.unwrap_or_else(Message::new);
        return Ok((remote_response, true));
    }

    for rec in local_response.answers() {
        if rec.record_type() != query.query_type() {
            warn!("local DNS response has inconsistent answer type {} for query {}", rec.record_type(), query);
            break
        }
        let forward = match rec.rdata() {
            RData::A(ref ip) => context.check_ip_in_proxy_list(&IpAddr::from(*ip)),
            RData::AAAA(ref ip) => context.check_ip_in_proxy_list(&IpAddr::from(*ip)),
            RData::PTR(_) => panic!("PTR records should not reach here"),
            _ => context.is_default_in_proxy_list(),
        };
        if !forward {
            debug!("pick local response (response): {:?}", local_response);
            return Ok((local_response, false));
        }
    }

    let remote_response = remote_response_fut.await.unwrap_or_else(Message::new);
    debug!("pick remote response (response): {:?}", remote_response);
    Ok((remote_response, true))
}

/// Start a DNS relay local server
pub async fn run(context: SharedContext) -> io::Result<()> {
    let local_addr = match context.config().config_type {
        ConfigType::DnsLocal => {
            // Standalone server
            context.config().local_addr.as_ref().expect("local config")<commit_msg>Refine the futures<commit_after>use std::{
    io,
    net::{IpAddr, SocketAddr},
    sync::Arc,
    time::Duration,
};

use tokio::sync::mpsc;

#[cfg(not(target_os = "android"))]
use std::net::{Ipv4Addr, SocketAddrV4};
#[cfg(not(target_os = "android"))]
use tokio::net::UdpSocket;

use log::{debug, error, info, warn};
use trust_dns_proto::{
    op::{header::MessageType, response_code::ResponseCode, Message, Query},
    rr::RData,
};

use crate::{
    config::ConfigType,
    context::SharedContext,
    relay::{
        loadbalancing::server::{PlainPingBalancer, ServerType},
        sys::create_udp_socket,
        utils::try_timeout,
    },
};

mod upstream;

async fn acl_lookup<Local, Remote>(
    context: SharedContext,
    local: Arc<Local>,
    remote: Arc<Remote>,
    query: &Query
) -> io::Result<(Message, bool)>
    where
        Local: upstream::Upstream,
        Remote: upstream::Upstream,
{
    // Start querying name servers
    debug!(
        "attempting lookup of {:?} {} with ns {:?} and {:?}",
        query.query_type(), query.name(), local, remote
    );

    let qname_in_proxy_list = context.check_query_in_proxy_list(query);
    let remote_response_fut = try_timeout(remote.lookup(query), Some(Duration::new(3, 0)));
    let local_response_fut = try_timeout(local.lookup(query), Some(Duration::new(3, 0)));

    match qname_in_proxy_list {
        Some(true) => {
            let remote_response = remote_response_fut.await.unwrap_or_else(|_| Message::new());
            debug!("pick remote response (qname): {:?}", remote_response);
            return Ok((remote_response, true));
        }
        Some(false) => {
            let local_response = local_response_fut.await.unwrap_or_else(|_| Message::new());
            debug!("pick local response (qname): {:?}", local_response);
            return Ok((local_response, false));
        }
        None => (),
    }

    let local_response = local_response_fut.await.unwrap_or_else(|_| Message::new());
    for rec in local_response.answers() {
        if rec.record_type() != query.query_type() {
            warn!("local DNS response has inconsistent answer type {} for query {}", rec.record_type(), query);
            break
        }
        let forward = match rec.rdata() {
            RData::A(ref ip) => context.check_ip_in_proxy_list(&IpAddr::from(*ip)),
            RData::AAAA(ref ip) => context.check_ip_in_proxy_list(&IpAddr::from(*ip)),
            RData::PTR(_) => panic!("PTR records should not reach here"),
            _ => context.is_default_in_proxy_list(),
        };
        if !forward {
            debug!("pick local response (response): {:?}", local_response);
            return Ok((local_response, false));
        }
    }

    let remote_response = remote_response_fut.await.unwrap_or_else(|_| Message::new());
    debug!("pick remote response (response): {:?}", remote_response);
    Ok((remote_response, true))
}

/// Start a DNS relay local server
pub async fn run(context: SharedContext) -> io::Result<()> {
    let local_addr = match context.config().config_type {
        ConfigType::DnsLocal => {
            // Standalone server
            context.config().local_addr.as_ref().expect("local config")
{'fraction_non_alphanumeric': 0.1497294926158795, 'fraction_numerical': 0.0019008626992250328, 'mean_word_length': 2.836231071228267, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 37, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3276605', 'n_tokens_mistral': 2107, 'n_tokens_neox': 2102, 'n_words': 418}