implement web-api server with stats

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());