feat: 支持通过 udp 转发。

src/main.rs

@@ -7,6 +7,7 @@ use tokio_tungstenite::connect_async;
 
 mod clash_conn_msg;
 mod statistics;
+mod udp_server;
 
 #[derive(Parser)]
 #[clap(
@@ -27,6 +28,8 @@ struct Args {
 
 #[tokio::main]
 async fn main() {
+    env_logger::init();
+
     println!("Hello, world!");
 
     let args = Args::parse();
@@ -61,22 +64,26 @@ async fn main() {
 
             let state = statistics_manager.get_state().await;
 
-            stdout()
-                .write_all(
-                    format!(
-                        "len: {},speed_upload: {}, speed_download: {}, update_direct_speed: {}, download_direct_speed: {}, update_proxy_speed: {}, download_proxy_speed: {}\n",
-                        state.connections,
-                        state.speed_upload,
-                        state.speed_download,
-                        state.direct_upload_speed,
-                        state.direct_download_speed,
-                        state.proxy_upload_speed,
-                        state.proxy_download_speed,
-                    )
-                    .as_bytes(),
-                )
-                .await
-                .unwrap();
+            // stdout()
+            //     .write_all(
+            //         format!(
+            //             "len: {},speed_upload: {}, speed_download: {}, update_direct_speed: {}, download_direct_speed: {}, update_proxy_speed: {}, download_proxy_speed: {}\n",
+            //             state.connections,
+            //             state.speed_upload,
+            //             state.speed_download,
+            //             state.direct_upload_speed,
+            //             state.direct_download_speed,
+            //             state.proxy_upload_speed,
+            //             state.proxy_download_speed,
+            //         )
+            //         .as_bytes(),
+            //     )
+            //     .await
+            //     .unwrap();
+
+            let udp_server = udp_server::UdpServer::global().await;
+            let buf = serde_json::to_vec(&state).unwrap();
+            udp_server.publish(&buf).await;
         })
     };
 

src/statistics.rs

@@ -1,10 +1,11 @@
 use std::{collections::HashMap, sync::Arc};
 
+use serde::Serialize;
 use tokio::sync::{Mutex, OnceCell};
 
 use crate::clash_conn_msg::{ClashConnectionMessage, ClashConnectionsWrapper};
 
-#[derive(Clone)]
+#[derive(Clone, Serialize)]
 pub struct StatisticsState {
     pub connections: usize,
     pub download_total: u64,
@@ -15,6 +16,7 @@ pub struct StatisticsState {
     pub direct_download_speed: u64,
     pub proxy_upload_speed: u64,
     pub proxy_download_speed: u64,
+    #[serde(skip)]
     connection_map: Arc<Mutex<HashMap<String, ClashConnectionMessage>>>,
 }
 

src/udp_server.rs

@@ -0,0 +1,119 @@
+use std::{
+    collections::HashMap,
+    net::{Ipv4Addr, SocketAddr},
+    sync::Arc,
+};
+
+use log::{error, info};
+use tokio::{
+    net::UdpSocket,
+    sync::{Mutex, OnceCell},
+    time::Instant,
+};
+use clap::Parser;
+
+use crate::Args;
+
+struct Client {
+    last_seen: Instant,
+    socket: UdpSocket,
+}
+
+#[derive(Clone)]
+pub struct UdpServer {
+    listen_addr: Ipv4Addr,
+    listen_port: u16,
+    clients: Arc<Mutex<HashMap<SocketAddr, Client>>>,
+}
+
+impl UdpServer {
+    pub async fn global() -> &'static Self {
+        static UDP_SERVER: OnceCell<UdpServer> = OnceCell::const_new();
+
+        UDP_SERVER
+            .get_or_init(|| async {
+                let args = Args::parse();
+                let listen_addr = Ipv4Addr::new(0, 0, 0, 0);
+                let listen_port = args.listen_port;
+
+                let udp_server = UdpServer::new(listen_addr, listen_port);
+                let server = udp_server.clone();
+                tokio::spawn(async move {
+                    if let Err(err) = server.run().await {
+                        error!("Failed to run UDP server: {}", err);
+                    }
+                });
+
+                udp_server
+            })
+            .await
+    }
+
+    pub fn new(listen_addr: Ipv4Addr, listen_port: u16) -> Self {
+        Self {
+            listen_addr,
+            listen_port,
+            clients: Arc::new(Mutex::new(HashMap::new())),
+        }
+    }
+
+    pub async fn run(&self) -> anyhow::Result<()> {
+        let mut buf = [0; 1024];
+        let socket = UdpSocket::bind((self.listen_addr, self.listen_port)).await?;
+
+        loop {
+            info!("Waiting for data");
+            let (len, addr) = socket.recv_from(&mut buf).await?;
+
+            info!("Received data({}) from {}", len, addr);
+
+            let mut clients = self.clients.lock().await;
+
+            let client = clients.get_mut(&addr);
+
+            if client.is_none() {
+                let socket = UdpSocket::bind((self.listen_addr, 0)).await;
+                if let Err(err) = socket {
+                    error!("Failed to bind socket: {}", err);
+                    continue;
+                }
+                let socket = socket.unwrap();
+                if let Err(err) = socket.connect(addr).await {
+                    error!("Failed to connect socket: {}", err);
+                    continue;
+                }
+
+                let client = Client {
+                    socket,
+                    last_seen: Instant::now(),
+                };
+
+                clients.insert(addr, client);
+            } else {
+                let client = client.unwrap();
+                client.last_seen = Instant::now();
+            }
+        }
+    }
+
+    pub async fn publish(&self, buf: &[u8]) {
+        let mut to_remove = Vec::new();
+        let mut clients = self.clients.lock().await;
+
+        for (addr, client) in clients.iter() {
+            if client.last_seen.elapsed().as_secs() > 10 {
+                to_remove.push(addr.clone());
+            } else {
+                if let Err(err) = client.socket.send(buf).await {
+                    error!("Failed to send data to {}: {}", addr, err);
+                } else {
+                    info!("Sent data to {}", addr);
+                }
+            }
+        }
+
+        for addr in to_remove {
+            clients.remove(&addr);
+        }
+    }
+}