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, sync::RwLock, task};

type SharedStats = Arc<RwLock<Stats>>;
type SharedList = Arc<List>;

#[rocket::get("/")]
async fn index(stats: &rocket::State<SharedStats>) -> Json<Stats> {
    Json(*stats.read().await)
}

#[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",
        ));
    }
    if opt.skip_auth && opt.auth != AuthMode::NoAuth {
        return Err(SocksError::ArgumentInputError(
            "Can't use skip-auth flag and authentication altogether.",
        ));
    }

    let list = Arc::new(List::from_opt(&opt.allow_list).await.map_err(|err| {
        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?;

    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(), stats.clone()));
            }
            Err(err) => error!("accept error = {:?}", err),
        }
    }
}

async fn serve_socks5(
    opt: &Opt,
    socket: tokio::net::TcpStream,
    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)
        }
        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.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());
    }

    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:#}"),
            },
        }
    })
}