xuproxy/src/main.rs

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

use log::{debug, info};
use serde::Deserialize;
use tokio::select;
use warp::http::StatusCode;
use warp::{body, header, path, query, Buf, Filter, Reply, any};

#[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() => return,
        _ = sigint.recv() => return
    }
}
#[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,
}

#[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)?;

    select! {
        r = run_server(config) => r?,
        _ = terminate_signal() => {}
    }
    info!("Shutting down...");
    Ok(())
}

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

async fn handle_put(
    id: String,
    name: String,
    length: u64,
    query: PutQueryString,
    body: impl Buf,
    config: Arc<Config>
) -> impl Reply {
    debug!(
        "Received PUT request, id({}) name({}) length({}) token({})",
        id, name, length, query.v
    );
    let mut supplied_token = [0_u8; 32];
    if let Err(_) = hex::decode_to_slice(&query.v, &mut supplied_token) {
        debug!("Failed to parse hex string '{}'", query.v);
        return StatusCode::FORBIDDEN
    }
    let hmac_input = format!("{}/{} {}", id, name, 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
    }
    let mut vec = Vec::new();
    body.reader().read_to_end(&mut vec).unwrap();
    StatusCode::CREATED
}

async fn run_server(conf: Config) -> anyhow::Result<()> {
    let conf = Arc::new(conf);

    let put_route = warp::put()
        .and(path::param::<String>())
        .and(path::param::<String>())
        .and(header::<u64>("content-length"))
        .and(query::<PutQueryString>())
        .and(body::aggregate())
        .and(any().map({
            let conf = conf.clone();
            move || conf.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)
            .run(conf.address)
            .await
    } else {
        warp::serve(routes).run(conf.address).await
    }
    Ok(())
}