implement web-api server with stats

Cargo.toml

@@ -1,10 +1,10 @@
 [package]
 name = "psocks"
-version = "0.1.0"
+version = "0.2.0"
 edition = "2024"
 license = "MIT"
 readme = "README.md"
-description = "Experimental async SOCKS5 (TCP/UDP) proxy server based on fast-socks5, featuring allowlist-based access control (drop everything but allowed by user)"
+description = "Experimental async SOCKS5 (TCP/UDP) proxy server based on fast-socks5, featuring allowlist-based access control (drop everything but allowed by user) with JSON/API based on Rocket framework"
 keywords = ["proxy", "whitelist", "privacy"]
 categories = ["network-programming"]
 repository = "https://codeberg.org/postscriptum/psocks"
@@ -15,5 +15,7 @@ env_logger = "0.11.9"
 fast-socks5 = "1.0.0"
 log = "0.4.29"
 reqwest = "0.13.2"
+rocket = { version = "0.5.1", features = ["json"] }
+serde = { version = "1.0.228", features = ["derive"] }
 structopt = "0.3.26"
 tokio = { version = "1", features = ["full"] }

README.md

@@ -1,22 +1,24 @@
 # psocks
 
-Experimental async SOCKS5 (TCP/UDP) proxy server based on [fast-socks5](https://github.com/dizda/fast-socks5/blob/master/examples/server.rs), featuring allowlist-based access control (drop everything but allowed by user)
+Experimental async SOCKS5 (TCP/UDP) proxy server based on [fast-socks5](https://github.com/dizda/fast-socks5/blob/master/examples/server.rs), featuring allowlist-based access control (drop everything but allowed by user) with JSON/API based on [Rocket](https://rocket.rs) framework.
 
 ## Roadmap
 
 * [ ] Range support
 * [ ] Local Web-API
-  * [ ] Block stats
+  * [x] Block stats
   * [ ] In-memory list update (without server restart)
   * [ ] Performance optimization
 
 ## Usage
 
 ``` bash
-RUST_LOG=trace cargo run -- -a=/path/to/allow1.txt \
+RUST_LOG=psocks=trace cargo run -- -a=/path/to/allow1.txt \
                                    -a=http://localhost/allow2.txt \
                                    no-auth
 ```
+* set `socks5://127.0.0.1:1080` proxy in your application
+* open http://127.0.0.1:8010 in browser for stats & control API
 
 ### Allow list example
 

src/list.rs

@@ -48,4 +48,9 @@ impl List {
             }),
         }
     }
+    pub fn entries(&self) -> usize {
+        match self {
+            Self::Allow(list) => list.len(),
+        }
+    }
 }

src/main.rs

@@ -1,28 +1,56 @@
 mod list;
 mod opt;
+mod stats;
 
 use anyhow::Context;
 use fast_socks5::{
     ReplyError, Result, Socks5Command, SocksError,
     server::{DnsResolveHelper as _, Socks5ServerProtocol, run_tcp_proxy, run_udp_proxy},
 };
+use list::List;
 use log::*;
 use opt::{AuthMode, Opt};
+use rocket::serde::json::Json;
+use stats::Stats;
 use std::{future::Future, sync::Arc};
 use structopt::StructOpt;
-use tokio::net::TcpListener;
-use tokio::task;
+use tokio::{net::TcpListener, sync::RwLock, task};
 
-use crate::list::List;
+type SharedStats = Arc<RwLock<Stats>>;
+type SharedList = Arc<List>;
 
-#[tokio::main]
-async fn main() -> Result<()> {
-    env_logger::init();
-    spawn_socks_server().await
+#[rocket::get("/")]
+async fn index(stats: &rocket::State<SharedStats>) -> Json<Stats> {
+    Json(*stats.read().await)
 }
 
-async fn spawn_socks_server() -> Result<()> {
+#[rocket::launch]
+async fn rocket() -> _ {
+    env_logger::init();
+
     let opt: &'static Opt = Box::leak(Box::new(Opt::from_args()));
+    let stats = Arc::new(RwLock::new(Stats::default()));
+
+    tokio::spawn({
+        let socks_stats = stats.clone();
+        async move {
+            if let Err(err) = spawn_socks_server(opt, socks_stats).await {
+                error!("SOCKS server failed: `{err}`");
+            }
+        }
+    });
+
+    rocket::build()
+        .configure(rocket::Config {
+            port: opt.api_addr.port(),
+            address: opt.api_addr.ip(),
+            ..rocket::Config::release_default()
+        })
+        .manage(stats)
+        .mount("/", rocket::routes![index])
+}
+
+async fn spawn_socks_server(opt: &'static Opt, stats: SharedStats) -> Result<()> {
     if opt.allow_udp && opt.public_addr.is_none() {
         return Err(SocksError::ArgumentInputError(
             "Can't allow UDP if public-addr is not set",
@@ -38,6 +66,10 @@ async fn spawn_socks_server() -> Result<()> {
         error!("Can't parse whitelist: `{err}`");
         SocksError::ArgumentInputError("Can't parse whitelist")
     })?);
+    {
+        let mut s = stats.write().await;
+        s.entries = list.entries();
+    }
 
     let listener = TcpListener::bind(&opt.listen_addr).await?;
 
@@ -46,7 +78,7 @@ async fn spawn_socks_server() -> Result<()> {
     loop {
         match listener.accept().await {
             Ok((socket, _client_addr)) => {
-                spawn_and_log_error(serve_socks5(opt, socket, list.clone()));
+                spawn_and_log_error(serve_socks5(opt, socket, list.clone(), stats.clone()));
             }
             Err(err) => error!("accept error = {:?}", err),
         }
@@ -56,8 +88,12 @@ async fn spawn_socks_server() -> Result<()> {
 async fn serve_socks5(
     opt: &Opt,
     socket: tokio::net::TcpStream,
-    list: Arc<List>,
+    list: SharedList,
+    stats: SharedStats,
 ) -> Result<(), SocksError> {
+    let mut s = stats.write().await;
+    s.request += 1;
+
     let request = match &opt.auth {
         AuthMode::NoAuth if opt.skip_auth => {
             Socks5ServerProtocol::skip_auth_this_is_not_rfc_compliant(socket)
@@ -78,6 +114,7 @@ async fn serve_socks5(
     let (proto, cmd, addr) = request.resolve_dns().await?;
 
     if !list.has(&host) && !list.has(&addr.to_string()) {
+        s.blocked += 1;
         info!("Blocked connection attempt to: {host}");
         proto.reply_error(&ReplyError::ConnectionNotAllowed).await?;
         return Err(ReplyError::ConnectionNotAllowed.into());

src/opt.rs

@@ -1,4 +1,4 @@
-use std::{num::ParseFloatError, time::Duration};
+use std::{net::SocketAddr, num::ParseFloatError, time::Duration};
 use structopt::StructOpt;
 
 /// # How to use it:
@@ -18,8 +18,12 @@ use structopt::StructOpt;
 )]
 pub struct Opt {
     /// Bind on address address. eg. `127.0.0.1:1080`
-    #[structopt(short, long, default_value = "127.0.0.1:1080")]
-    pub listen_addr: String,
+    #[structopt(short = "l", long, default_value = "127.0.0.1:1080")]
+    pub listen_addr: SocketAddr,
+
+    /// Bind on address address. eg. `127.0.0.1:8010`
+    #[structopt(short = "A", long, default_value = "127.0.0.1:8010")]
+    pub api_addr: SocketAddr,
 
     /// Our external IP address to be sent in reply packets (required for UDP)
     #[structopt(long)]

src/stats.rs

@@ -0,0 +1,9 @@
+use rocket::serde::Serialize;
+
+#[derive(Debug, Default, Clone, Copy, Serialize)]
+#[serde(crate = "rocket::serde")]
+pub struct Stats {
+    pub entries: usize,
+    pub blocked: usize,
+    pub request: usize,
+}