args

Author: develon2015

src/main.rs

@@ -6,32 +6,50 @@ 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
             query[0] = (received / 0xff) as u8;
             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]);
+                }
             }
         });
     }