From 15066af5723c172e580247cb355a7b3b8c21756b Mon Sep 17 00:00:00 2001
From: kasumi <302615249@qq.com>
Date: Sat, 24 Apr 2021 01:49:26 +0800
Subject: [PATCH] args

---
 src/main.rs | 37 +++++++++++++++++++++++++++----------
 1 file changed, 27 insertions(+), 10 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 231a5a8..89f6943 100755
--- a/src/main.rs
+++ b/src/main.rs
@@ -6,24 +6,37 @@ const LISTEN: &str = "0.0.0.0:53";
 const SERVER: &str = "1.1.1.1:53";
 
 async fn udp_serv() -> std::io::Result<()> {
-    let listener = tokio::net::UdpSocket::bind(LISTEN).await;
+    let listen = if let Some(listen) = std::env::args().nth(1) {
+        listen
+    } else {
+        LISTEN.to_string()
+    };
+    let log = if let Some(_) = std::env::args().nth(3) {
+        true
+    } else {
+        false
+    };
+    let listener = tokio::net::UdpSocket::bind(listen).await;
     if let Err(e) = listener {
         println!("无法启用监听服务: {}", e);
         return Err(e);
     }
     let listener: UdpSocket = listener.ok().unwrap();
     let listener = std::sync::Arc::new(listener);
+    println!("DNS代理服务已启动");
 
+    let mut query = [0u8; 1024]; // DNS query request data
     loop {
         let listener = listener.clone();
-        let mut query = [0u8; 10240*2]; // DNS query request data
         let recv_result = listener.recv_from(&mut query[2..]).await;
         if let Err(_) = recv_result {
             continue;
         }
         let (received, client) = recv_result.ok().unwrap();
-        println!("DNS查询...");
-        println!("data: {:?}", &query[..received + 2]);
+        if log {
+            println!("DNS查询...");
+            println!("data: {:?}", &query[..received + 2]);
+        }
 
         tokio::spawn(async move {
             // TCP data body length
@@ -31,7 +44,12 @@ async fn udp_serv() -> std::io::Result<()> {
             query[1] = (received % 0xff) as u8;
 
             // Connect server
-            let tcp = tokio::net::TcpStream::connect(SERVER).await;
+            let server = if let Some(server) = std::env::args().nth(2) {
+                server
+            } else {
+                SERVER.to_string()
+            };
+            let tcp = tokio::net::TcpStream::connect(server).await;
             if let Err(e) = tcp {
                 println!("无法连接服务器{}：{}", SERVER, e);
                 return;
@@ -46,11 +64,10 @@ async fn udp_serv() -> std::io::Result<()> {
             }
 
             if let Ok(le) = tcp.read(&mut query).await {
-                listener
-                    .send_to(&query[2..le], client)
-                    .await
-                    .unwrap();
-                println!("resp: {:?}", &query[..le]);
+                listener.send_to(&query[2..le], client).await.unwrap();
+                if log {
+                    println!("resp: {:?}", &query[..le]);
+                }
             }
         });
     }