feat: adding OpenWRT WAN port speeds.

src/main.rs

@@ -12,8 +12,11 @@ use tokio_tungstenite::connect_async;
 mod clash_conn_msg;
 mod statistics;
 mod udp_server;
+mod wan;
 
-#[derive(Parser)]
+use wan::poll_wan_traffic;
+
+#[derive(Parser, Debug)]
 #[clap(
     version = "0.1.0",
     author = "Ivan Li",
@@ -22,22 +25,40 @@ mod udp_server;
 struct Args {
     #[clap(
         short,
+        env,
         long,
-        default_value = "ws://192.168.31.1:9090/connections?token=123456"
+        default_value = "ws://192.168.1.1:9090/connections?token=123456"
     )]
     clash_url: String,
-    #[clap(short, long, default_value = "17890")]
+    #[clap(short = 'p', long, env, default_value = "17890")]
     listen_port: u16,
+    #[clap(short, long, env, default_value = "http://192.168.1.1/cgi-bin/luci")]
+    luci_url: String,
+    #[clap(short = 'u', long, env, default_value = "root")]
+    luci_username: String,
+    #[clap(short = 'P', long, env, default_value = "123456")]
+    luci_password: String,
 }
 
 #[tokio::main]
 async fn main() {
+    dotenvy::dotenv().ok();
     env_logger::init();
 
     println!("Hello, world!");
 
     let args = Args::parse();
 
+    tokio::spawn(async move {
+        poll_wan_traffic(
+            args.luci_url.as_str(),
+            args.luci_username.as_str(),
+            args.luci_password.as_str(),
+        )
+        .await
+        .unwrap();
+    });
+
     let connect_addr = args.clash_url;
 
     loop {
@@ -99,7 +120,7 @@ async fn pipe(connect_addr: String) {
             buf[8..16].copy_from_slice(&state.direct_download_speed.to_le_bytes());
             buf[16..24].copy_from_slice(&state.proxy_upload_speed.to_le_bytes());
             buf[24..32].copy_from_slice(&state.proxy_download_speed.to_le_bytes());
-            udp_server.publish(&buf).await;
+            udp_server.publish_clash(&buf).await;
         })
     };
 

src/udp_server.rs

@@ -4,13 +4,13 @@ use std::{
     sync::Arc,
 };
 
+use clap::Parser;
 use log::{error, info};
 use tokio::{
     net::UdpSocket,
     sync::{Mutex, OnceCell},
     time::Instant,
 };
-use clap::Parser;
 
 use crate::Args;
 
@@ -24,6 +24,7 @@ pub struct UdpServer {
     listen_addr: Ipv4Addr,
     listen_port: u16,
     clients: Arc<Mutex<HashMap<SocketAddr, Client>>>,
+    tx_buffer: Arc<Mutex<[u8; 48]>>,
 }
 
 impl UdpServer {
@@ -54,6 +55,7 @@ impl UdpServer {
             listen_addr,
             listen_port,
             clients: Arc::new(Mutex::new(HashMap::new())),
+            tx_buffer: Arc::new(Mutex::new([0u8; 48])),
         }
     }
 
@@ -95,7 +97,25 @@ impl UdpServer {
         }
     }
 
-    pub async fn publish(&self, buf: &[u8]) {
+    pub async fn publish_clash(&self, buf: &[u8]) {
+        let mut tx_buffer = self.tx_buffer.lock().await;
+        tx_buffer[..32].copy_from_slice(buf);
+        let buf = tx_buffer.clone();
+        drop(tx_buffer);
+
+        self.publish(&buf).await;
+    }
+
+    pub async fn publish_wan(&self, buf: &[u8]) {
+        let mut tx_buffer = self.tx_buffer.lock().await;
+        tx_buffer[32..].copy_from_slice(buf);
+        let buf = tx_buffer.clone();
+        drop(tx_buffer);
+
+        self.publish(&buf).await;
+    }
+
+    async fn publish(&self, buf: &[u8]) {
         let mut to_remove = Vec::new();
         let mut clients = self.clients.lock().await;
 
@@ -103,7 +123,7 @@ impl UdpServer {
             if client.last_seen.elapsed().as_secs() > 10 {
                 to_remove.push(addr.clone());
             } else {
-                if let Err(err) = client.socket.send(buf).await {
+                if let Err(err) = client.socket.send(&buf).await {
                     error!("Failed to send data to {}: {}", addr, err);
                 } else {
                     info!("Sent data to {}", addr);

src/wan.rs

@@ -0,0 +1,194 @@
+use std::time::Duration;
+
+use anyhow::anyhow;
+use anyhow::bail;
+use anyhow::Result;
+use http_body_util::BodyExt;
+use http_body_util::Empty;
+use hyper::StatusCode;
+use hyper::{
+    body::{Buf, Bytes},
+    Request,
+};
+use hyper_util::rt::TokioIo;
+use log::debug;
+use serde::{Deserialize, Serialize};
+use tokio::net::TcpStream;
+use tokio::time::sleep;
+
+use crate::udp_server;
+
+pub type OpenWRTIFaces = Vec<OpenWRTIFace>;
+
+pub async fn poll_wan_traffic(luci_url: &str, username: &str, password: &str) -> Result<()> {
+    // Parse our URL...
+    let url = format!("{}/admin/network/iface_status/wan", luci_url).parse::<hyper::Uri>()?;
+
+    // Get the host and the port
+    let host = url.host().expect("uri has no host");
+    let port = url.port_u16().unwrap_or(80);
+
+    let address = format!("{}:{}", host, port);
+
+    let mut cookies_str = String::new();
+    // Open a TCP connection to the remote host
+    loop {
+        let stream = TcpStream::connect(address.clone()).await?;
+
+        // Use an adapter to access something implementing `tokio::io` traits as if they implement
+        // `hyper::rt` IO traits.
+        let io = TokioIo::new(stream);
+
+        // Create the Hyper client
+        let (mut sender, conn) = hyper::client::conn::http1::handshake(io).await?;
+
+        // Spawn a task to poll the connection, driving the HTTP state
+        tokio::task::spawn(async move {
+            if let Err(err) = conn.await {
+                println!("Connection failed: {:?}", err);
+            }
+        });
+
+        let mut prev_up = 0;
+        let mut prev_down = 0;
+
+        loop {
+            if sender.is_closed() {
+                break;
+            }
+
+            if !sender.is_ready() {
+                sleep(Duration::from_millis(100)).await;
+                continue;
+            }
+
+            let poll_req = generate_poll_wan_req(luci_url, &cookies_str)?;
+
+            // Await the response...
+
+            let res = sender.send_request(poll_req.clone()).await?;
+
+            debug!("Response status: {}", res.status());
+
+            match res.status() {
+                StatusCode::FORBIDDEN | StatusCode::UNAUTHORIZED => {
+                    cookies_str = login(luci_url, username, password).await?;
+                    continue;
+                }
+                _ => {}
+            }
+
+            // asynchronously aggregate the chunks of the body
+            let body = res.collect().await?.aggregate();
+
+            // try to parse as json with serde_json
+            let interfaces: OpenWRTIFaces = serde_json::from_reader(body.reader())?;
+
+            let wan = interfaces.first().unwrap();
+
+            if prev_up > wan.tx_bytes || prev_down > wan.rx_bytes {
+                prev_up = wan.tx_bytes;
+                prev_down = wan.rx_bytes;
+            }
+
+            let up_speed = (wan.tx_bytes - prev_up) as u64;
+            let down_speed = (wan.rx_bytes - prev_down) as u64;
+            prev_up = wan.tx_bytes;
+            prev_down = wan.rx_bytes;
+
+            let udp_server = udp_server::UdpServer::global().await;
+            let mut buf = [0; 16];
+            buf[0..8].copy_from_slice(&up_speed.to_le_bytes());
+            buf[8..16].copy_from_slice(&down_speed.to_le_bytes());
+            udp_server.publish_wan(&buf).await;
+
+            // println!("speed: ↑ {}, ↓ {}", up_speed, down_speed);
+
+            sleep(Duration::from_secs(1)).await;
+        }
+    }
+}
+
+async fn login(luci_url: &str, username: &str, password: &str) -> Result<String> {
+    let url = format!("{}/admin/network/iface_status/wan", luci_url).parse::<hyper::Uri>()?;
+    let host = url.host().expect("uri has no host");
+    let port = url.port_u16().unwrap_or(80);
+
+    let address = format!("{}:{}", host, port);
+    let stream = TcpStream::connect(address).await?;
+    let io = TokioIo::new(stream);
+    let (mut sender, conn) = hyper::client::conn::http1::handshake(io).await?;
+    tokio::task::spawn(async move {
+        if let Err(err) = conn.await {
+            println!("Connection failed: {:?}", err);
+        }
+    });
+
+    let login_req = generate_login_req(luci_url, username, password)?;
+    let res = sender.send_request(login_req.clone()).await?;
+
+    if res.status() == StatusCode::FORBIDDEN {
+        bail!("Login failed, got status: {}", res.status());
+    }
+
+    let cookies = res.headers().get_all(hyper::header::SET_COOKIE);
+
+    let cookies = cookies
+        .iter()
+        .map(|cookie| cookie.to_str().unwrap().split(';').next());
+    let cookies = cookies.filter_map(|cookie| cookie);
+    let cookies_str = cookies.collect::<Vec<_>>().join("; ");
+
+    Ok(cookies_str)
+}
+
+fn generate_poll_wan_req(luci_url: &str, cookie: &str) -> Result<Request<Empty<Bytes>>> {
+    let url = format!("{}/admin/network/iface_status/wan", luci_url).parse::<hyper::Uri>()?;
+    let authority = url.authority().unwrap().clone();
+    let target = url.path_and_query().unwrap().clone();
+
+    debug!("Polling WAN traffic. Cookie: {:?}", cookie);
+    return Request::builder()
+        .uri(target)
+        .header(hyper::header::HOST, authority.as_str())
+        .header(hyper::header::CONNECTION, "Keep-Alive")
+        .header(hyper::header::COOKIE, cookie)
+        .body(Empty::<Bytes>::new())
+        .map_err(|e| anyhow!(e));
+}
+
+fn generate_login_req(luci_url: &str, username: &str, password: &str) -> Result<Request<String>> {
+    let url = format!("{}/admin/network/iface_status/wan", luci_url).parse::<hyper::Uri>()?;
+    let authority = url.authority().unwrap().clone();
+    let target = url.path_and_query().unwrap().clone();
+
+    let login_params = format!("luci_username={}&luci_password={}", username, password);
+    let body = login_params;
+
+    return Request::builder()
+        .uri(target)
+        .method("POST")
+        .header(hyper::header::HOST, authority.as_str())
+        .header(
+            hyper::header::CONTENT_TYPE,
+            "application/x-www-form-urlencoded",
+        )
+        .header(hyper::header::CONNECTION, "Keep-Alive")
+        .body(body)
+        .map_err(|e| anyhow!(e));
+}
+
+#[derive(Serialize, Deserialize)]
+pub struct OpenWRTIFace {
+    rx_bytes: u64,
+    tx_bytes: u64,
+}
+
+impl Default for OpenWRTIFace {
+    fn default() -> Self {
+        Self {
+            rx_bytes: 0,
+            tx_bytes: 0,
+        }
+    }
+}