xuproxy/src/main.rs

use hmac_sha256::HMAC;
use std::net::SocketAddr;
use std::sync::Arc;
use std::{env, fs, thread};

use crate::database::{DbExecutor, ExecutorConnection};
use log::{debug, error, info};
use serde::Deserialize;
use tokio::sync::oneshot;
use warp::http::StatusCode;
use warp::hyper::body::Bytes;
use warp::path::FullPath;
use warp::{any, body, header, path, query, Filter, Reply};

mod database;
mod discord;

#[cfg(unix)]
async fn terminate_signal() {
    use tokio::signal::unix::{signal, SignalKind};
    let mut sigterm = signal(SignalKind::terminate()).unwrap();
    let mut sigint = signal(SignalKind::interrupt()).unwrap();
    debug!("Installed ctrl+c handler");
    select! {
        _ = sigterm.recv() => {},
        _ = sigint.recv() => {}
    }
}
#[cfg(windows)]
async fn terminate_signal() {
    use tokio::signal::windows::ctrl_c;
    let mut ctrlc = ctrl_c().unwrap();
    debug!("Installed ctrl+c handler");
    let _ = ctrlc.recv().await;
}

#[derive(Debug, Deserialize)]
struct Config {
    tls: Option<TlsConfig>,
    address: SocketAddr,
    webhook: String,
    secret: String,
    dbpath: String,
}

#[derive(Debug, Deserialize)]
struct TlsConfig {
    cert: String,
    key: String,
}

#[tokio::main(flavor = "current_thread")]
async fn main() -> anyhow::Result<()> {
    pretty_env_logger::init_custom_env("XUPROXY_LOG");
    let config_var = env::var("XUPROXY_CONFIG");
    let config_path = config_var.as_deref().unwrap_or("xuproxy.toml");
    info!("Loading config '{}'", config_path);
    let config_str = fs::read_to_string(config_path)?;
    let config: Config = toml::from_str(&config_str)?;
    info!("Initializing database...");
    let (mut db_exec, db_conn) = DbExecutor::create(&config.dbpath)?;
    let executor_thread = thread::spawn(move || {
        db_exec.run();
        log::info!("Database executor shutting down");
    });

    let (ctx, crx) = oneshot::channel();
    let server_task = tokio::spawn(run_server(config, db_conn, crx));
    terminate_signal().await;
    info!("Shutdown signal received, powering down");
    let _ = ctx.send(());
    server_task.await.unwrap_or_else(|e| error!("Couldn't await the server task: {}", e));
    executor_thread.join().unwrap();
    Ok(())
}

#[derive(Deserialize)]
struct PutQueryString {
    v: String,
}

async fn handle_put(
    filename: FullPath,
    length: u64,
    query: PutQueryString,
    body: Bytes,
    config: Arc<Config>,
    db: ExecutorConnection
) -> impl Reply {
    let filename_str = &filename.as_str()[1..];
    debug!(
        "Received PUT request, name({}) length({}) token({})",
        filename_str, length, query.v
    );
    let mut supplied_token = [0_u8; 32];
    if let Err(e) = hex::decode_to_slice(&query.v, &mut supplied_token) {
        debug!("Failed to parse hex string '{}': {}", query.v, e);
        return StatusCode::FORBIDDEN;
    }
    let hmac_input = format!("{} {}", filename_str, length);
    let calculated_token = HMAC::mac(hmac_input.as_bytes(), config.secret.as_bytes());
    if supplied_token != calculated_token {
        debug!("Token '{}' doesn't match HMAC secret", query.v);
        return StatusCode::FORBIDDEN;
    }
    match discord::upload_webhook(&config.webhook, body, filename_str).await {
        Err(e) => {
            debug!("Could not upload file to Discord: {}", e);
            StatusCode::FORBIDDEN
        }
        Ok(url) => {
            if !db.add_file(filename_str.to_string(), url).await {
                StatusCode::FORBIDDEN
            } else {
                StatusCode::CREATED
            }
        }
    }
}

async fn run_server(
    conf: Config,
    db: ExecutorConnection,
    cancel: oneshot::Receiver<()>
) {
    let conf = Arc::new(conf);

    let put_route = warp::put()
        .and(path::full())
        .and(header::<u64>("content-length"))
        .and(query::<PutQueryString>())
        .and(body::bytes())
        .and(any().map({
            let conf = conf.clone();
            move || conf.clone()
        }))
        .and(any().map({
            let db = db.clone();
            move || db.clone()
        }))
        .then(handle_put);

    let routes = put_route;

    if let Some(tls) = &conf.tls {
        warp::serve(routes)
            .tls()
            .cert_path(&tls.cert)
            .key_path(&tls.key)
            .bind_with_graceful_shutdown(conf.address, async {
                let _ = cancel.await;
            }).1.await;
    } else {
        warp::serve(routes).bind_with_graceful_shutdown(conf.address, async {
            let _ = cancel.await;
        }).1.await;
    };
    info!("Webserver shutting down");
}