sdlan-lib-rs/src/utils/dns.rs

use std::sync::atomic::{AtomicU8, Ordering};
use std::sync::{Arc, atomic::AtomicU16};
use std::{net::{Ipv4Addr, SocketAddr}, time::Duration};

use arc_swap::ArcSwap;
use sdlan_sn_rs::utils::Result;
use simple_dns::{Name, QCLASS, QTYPE, Question, rdata::RData};
use tokio::signal;
use tokio::sync::mpsc::{channel, Sender, Receiver};
use tokio::{net::UdpSocket, time};

use dashmap::DashMap;
use sdlan_sn_rs::utils::get_current_timestamp;

pub struct DNSCacheInfo {
    pub src_ip: u32,
    pub src_port: u16,
    pub origin_transaction_id: u16,
    pub added: u64,
}

pub struct DNSMatcher {
    current_transaction_id: AtomicU16,
    // match transaction_id to (client ip, client port)
    matcher: DashMap<u16, DNSCacheInfo>,
}

impl DNSMatcher {
    pub fn new() -> Self {
        Self {
            current_transaction_id: AtomicU16::new(0),
            matcher: DashMap::new(),
        }
    }

    pub fn retain(&self) {
        let now = get_current_timestamp();
        self.matcher.retain(|_, v| (now - v.added) < 10);
    }

    pub fn generate_transaction_id(&self, client_ip: u32, client_port: u16, origin_transaction_id: u16) -> u16 {
        let transaction_id = self.current_transaction_id.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        self.matcher.insert(transaction_id, DNSCacheInfo{
            src_ip: client_ip, 
            src_port: client_port,
            origin_transaction_id,
            added: get_current_timestamp(),
        });
        transaction_id
    }

    pub fn get_client_info(&self, transaction_id: u16) -> Option<(u32, u16, u16)> {
        let res = self.matcher.remove(&transaction_id)?;
        Some((res.1.src_ip, res.1.src_port, res.1.origin_transaction_id))
    }
}



/// usage
/// ```rust
/// let (tx, rx) = tokio::sync::mpsc::channel(10);
/// let client = DynamicDNSClient::new("223.5.5.5:53".parse().unwrap(), tx);
/// 
/// ```
pub struct DynamicDNSClient {
    current_socket: Arc<ArcSwap<UdpSocket>>,
    rotate_send_channel: Sender<bool>,
    query_count: AtomicU8,
    response_tx: Sender<(Vec<u8>, SocketAddr)>,
    dns_server: SocketAddr,
}

impl DynamicDNSClient {
    pub async fn new(dns_server: SocketAddr, response_tx: Sender<(Vec<u8>, SocketAddr)>) -> Result<Self> {
        let initial_socket = Arc::new(UdpSocket::bind("0.0.0.0:0").await?);
        let current_socket = Arc::new(ArcSwap::from(initial_socket.clone()));

        Self::spawn_receiver(initial_socket, response_tx.clone());

        let (tx, rx) = channel(5);

        let result = Self {
            current_socket,
            dns_server,
            response_tx,
            rotate_send_channel: tx,
            query_count: AtomicU8::new(0),
        };
        result.start_refresher(rx);
        Ok(result)
    }

    pub async fn send_query(&self, packet: &[u8]) -> Result<usize> {
        let sock_guard = self.current_socket.load();
        let size = sock_guard.send_to(packet, self.dns_server).await?;
        let count = self.query_count.fetch_add(1, Ordering::Relaxed);
        if count >= 100 {
            self.query_count.store(0, Ordering::Release);
            let _ = self.rotate_send_channel.send(true).await;
        }
        Ok(size)
    }

    fn start_refresher(&self, mut rx: Receiver<bool>) {
        let socket_swap = self.current_socket.clone();
        let response_tx = self.response_tx.clone();

        tokio::spawn(async move {
            loop {
                tokio::select! {
                    _ = tokio::time::sleep(Duration::from_secs(60)) => {

                    }
                    data = rx.recv() => {
                        if data.is_none() {
                            panic!("global dns rx None");
                        }
                    }
                };
                // tokio::time::sleep(Duration::from_secs(60)).await;
                match UdpSocket::bind("0.0.0.0:0").await {
                    Ok(new_socket) => {
                        let new_socket = Arc::new(new_socket);
                        // let port = new_socket.local_addr().unwrap().port();
                        Self::spawn_receiver(new_socket.clone(), response_tx.clone());
                        socket_swap.store(new_socket);
                    }
                    Err(e) => {
                        eprintln!("failed to refresh: {}", e);
                    }
                }
            }
        });
    }

    fn spawn_receiver(socket: Arc<UdpSocket>, tx: Sender<(Vec<u8>, SocketAddr)>) {
        tokio::spawn(async move {
            let mut buf = vec![0u8; 1024];
            let port = socket.local_addr().unwrap().port();

            loop {
                match time::timeout(Duration::from_secs(65), socket.recv_from(&mut buf)).await {
                    Ok(Ok((len, from))) => {
                        let data = buf[..len].to_vec();
                        if tx.send((data, from)).await.is_err() {
                            break;
                        }
                    }
                    Ok(Err(_)) => break,
                    Err(_) => {
                        // timeout occured
                        break;
                    }
                }
            }
            eprintln!("port {} has been closed", socket.local_addr().unwrap().port());
        });
    }
}

/*
async fn test_simple_dns() {
    let mut id = 0;
    let (tx, mut rx) = tokio::sync::mpsc::channel(10);
    let client = DynamicDNSClient::new("223.5.5.5:53".parse().unwrap(), tx).await.unwrap();

    tokio::spawn(async move {
        loop {
            let data = rx.recv().await;
            if let Some(data) = data {
                let packet = simple_dns::Packet::parse(&data.0).unwrap();
                println!("got response: id = {}", packet.id());
                for answer in packet.answers {
                    println!("  domain: {}", answer.name);
                    match answer.rdata {
                        RData::A(a_record) => {
                            println!("    {}", Ipv4Addr::from_bits(a_record.address).to_string())
                        }
                        _other => {
                            println!("    other response type")
                        }
                    }
                }
            }
        }
    });

    loop {
        let mut packet = simple_dns::Packet::new_query(id);
        id += 1;
        let name = Name::new("www.baidu.com").unwrap();
        let question = Question {
            qname: name,
            qtype: QTYPE::TYPE(simple_dns::TYPE::A),
            qclass: QCLASS::CLASS(simple_dns::CLASS::IN),
            unicast_response: false,
        };
        packet.questions.push(question);
        let question = packet.build_bytes_vec().unwrap();
        client.send_query(&question).await;
        tokio::time::sleep(Duration::from_secs(5)).await;
    }
}
*/