psocks/src/main.rs

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::{State, serde::json::Json};
use stats::{Snapshot, Total};
use std::{future::Future, sync::Arc};
use structopt::StructOpt;
use tokio::{net::TcpListener, task};

#[rocket::get("/")]
async fn index(totals: &State<Arc<Total>>) -> Json<Snapshot> {
    Json(totals.inner().snapshot())
}

#[rocket::get("/allow/<rule>")]
async fn allow(rule: &str, list: &State<Arc<List>>, totals: &State<Arc<Total>>) -> Json<bool> {
    let result = list.allow(rule).await;
    totals.set_entries(list.entries().await);
    info!("Delete `{rule}` from the in-memory rules (operation status: {result:?})");
    Json(result.is_ok_and(|v| v))
}

#[rocket::get("/block/<rule>")]
async fn block(rule: &str, list: &State<Arc<List>>, totals: &State<Arc<Total>>) -> Json<bool> {
    let result = list.block(rule).await;
    totals.set_entries(list.entries().await);
    info!("Add `{rule}` to the in-memory rules (operation status: {result:?})");
    Json(result.is_ok_and(|v| v))
}

#[rocket::launch]
async fn rocket() -> _ {
    env_logger::init();

    let opt: &'static Opt = Box::leak(Box::new(Opt::from_args()));

    let list = Arc::new(
        List::from_opt(&opt.allow_list, opt.cache.clone())
            .await
            .unwrap(),
    );

    let totals = Arc::new(Total::with_entries(list.entries().await));

    tokio::spawn({
        let socks_list = list.clone();
        let socks_totals = totals.clone();
        async move {
            if let Err(err) = spawn_socks_server(opt, socks_list, socks_totals).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(list)
        .manage(totals)
        .mount("/", rocket::routes![index, allow, block])
}

async fn spawn_socks_server(opt: &'static Opt, list: Arc<List>, totals: Arc<Total>) -> Result<()> {
    if opt.allow_udp && opt.public_addr.is_none() {
        return Err(SocksError::ArgumentInputError(
            "Can't allow UDP if public-addr is not set",
        ));
    }
    if opt.skip_auth && opt.auth != AuthMode::NoAuth {
        return Err(SocksError::ArgumentInputError(
            "Can't use skip-auth flag and authentication altogether.",
        ));
    }

    let listener = TcpListener::bind(&opt.listen_addr).await?;

    info!("Listen for socks connections @ {}", &opt.listen_addr);

    loop {
        match listener.accept().await {
            Ok((socket, _client_addr)) => {
                spawn_and_log_error(serve_socks5(opt, socket, list.clone(), totals.clone()));
            }
            Err(err) => error!("accept error = {:?}", err),
        }
    }
}

async fn serve_socks5(
    opt: &Opt,
    socket: tokio::net::TcpStream,
    list: Arc<List>,
    totals: Arc<Total>,
) -> Result<(), SocksError> {
    totals.increase_request();

    let request = match &opt.auth {
        AuthMode::NoAuth if opt.skip_auth => {
            Socks5ServerProtocol::skip_auth_this_is_not_rfc_compliant(socket)
        }
        AuthMode::NoAuth => Socks5ServerProtocol::accept_no_auth(socket).await?,
        AuthMode::Password { username, password } => {
            Socks5ServerProtocol::accept_password_auth(socket, |user, pass| {
                user == *username && pass == *password
            })
            .await?
            .0
        }
    }
    .read_command()
    .await?;

    let (host, _) = request.2.clone().into_string_and_port(); // @TODO ref
    let (proto, cmd, addr) = request.resolve_dns().await?;

    if !list.any(&[&host, &addr.to_string()]).await {
        totals.increase_blocked();
        info!("Blocked connection attempt to: {host}");
        proto.reply_error(&ReplyError::ConnectionNotAllowed).await?;
        return Err(ReplyError::ConnectionNotAllowed.into());
    }

    match cmd {
        Socks5Command::TCPConnect => {
            run_tcp_proxy(proto, &addr, opt.request_timeout, false).await?;
        }
        Socks5Command::UDPAssociate if opt.allow_udp => {
            run_udp_proxy(
                proto,
                &addr,
                None,
                opt.public_addr.context("invalid reply ip")?,
                None,
            )
            .await?;
        }
        _ => {
            proto.reply_error(&ReplyError::CommandNotSupported).await?;
            return Err(ReplyError::CommandNotSupported.into());
        }
    };
    Ok(())
}

fn spawn_and_log_error<F>(fut: F) -> task::JoinHandle<()>
where
    F: Future<Output = Result<()>> + Send + 'static,
{
    task::spawn(async move {
        match fut.await {
            Ok(()) => {}
            Err(err) => match err {
                SocksError::ReplyError(reply_error) => {
                    if !matches!(reply_error, ReplyError::ConnectionNotAllowed) {
                        error!("{reply_error:#}")
                    }
                }
                _ => error!("{err:#}"),
            },
        }
    })
}