mirror-server/src/main.rs

use axum::{
    extract::{Host, OriginalUri, Json},
    http::{header::HeaderMap, Method},
    Router,
};
use clap::Parser;
use serde::Serialize;
use serde_json::Value;
use std::collections::BTreeMap;
use tower_http::trace;
use tracing::{Level, info, warn};

#[derive(Parser)]
#[command(author, version, about, long_about = None)]
struct CliArgs {
    /// Port to run on.
    #[arg(short, long, default_value_t = 8080)]
    port: u16,

    /// Listen to IP mask
    #[arg(short, long, default_value = "0.0.0.0")]
    ips: String,
}

#[derive(Serialize)]
struct EchoResponse {
    method: String,
    path: String,
    host: String,
    headers: BTreeMap<String, String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    body: Option<Value>,
}

async fn echo_request(
    method: Method,
    original_uri: OriginalUri,
    host: Host,
    header_map: HeaderMap,
    body: Option<Json<Value>>
) -> Json<EchoResponse> {
    let method = method.to_string();
    let host = host.0;
    let path = original_uri.path().to_string();
    let headers = header_map
        .iter()
        .map(|(name, value)| (name.to_string(), value.to_str().unwrap_or("").to_string()))
        .collect();

    let json_body = match body {
        None => {
            warn!("Received a non-JSON body.");
            None
        },
        Some(Json(value)) => {
            info!("JSON request: {}", value.to_string());
            Some(value)},
    };
    let response = EchoResponse {
        method,
        host,
        path,
        headers,
        body: json_body,
    };
    Json(response)
}

#[tokio::main]
async fn main() {
    let cli_args = CliArgs::parse();
    let listen_on = (cli_args.ips, cli_args.port);
    let listen_on = format!("{ip}:{port}", ip = listen_on.0, port = listen_on.1);

    // From https://stackoverflow.com/questions/75009289/how-to-enable-logging-tracing-with-axum
    tracing_subscriber::fmt()
        .with_max_level(Level::INFO)
        .init();
    let app = Router::new().fallback(echo_request).layer(
        trace::TraceLayer::new_for_http()
            .make_span_with(trace::DefaultMakeSpan::new().level(Level::INFO))
            .on_response(trace::DefaultOnResponse::new().level(Level::INFO)),
    );

    info!("Starting the mirror-server to listen to {}", listen_on);
    let listener = tokio::net::TcpListener::bind(&listen_on)
        .await
        .unwrap_or_else(|_| panic!("Failed to binding to {}", listen_on));
    axum::serve(listener, app)
        .await
        .expect("Server should start");
}