mylinspirer/src/proxy.rs

use std::sync::Arc;

use axum::{
    body::Bytes,
    extract::{Path, State},
    http::{Request, StatusCode},
    response::{IntoResponse, Response},
};
use http_body_util::BodyExt;
use serde_json::Value;
use tracing::{error, info, warn};

use crate::{crypto, AppState};

pub async fn proxy_handler(
    State(state): State<Arc<AppState>>,
    Path(path): Path<String>,
    req: Request<axum::body::Body>,
) -> impl IntoResponse {
    let (parts, body) = req.into_parts();
    let body_bytes = match body.collect().await {
        Ok(body) => body.to_bytes(),
        Err(e) => {
            error!("Failed to read request body: {}", e);
            return (
                StatusCode::INTERNAL_SERVER_ERROR,
                "Failed to read request body".to_string(),
            )
                .into_response();
        }
    };

    let decrypted_request_log = match process_and_log_request(&body_bytes, &state.key, &state.iv) {
        Ok(log) => log,
        Err(e) => {
            warn!("Failed to process request for logging: {}", e);
            Value::String("Could not decrypt request".to_string())
        }
    };

    let mut target_url = state.target_url.clone();
    target_url.set_path(&path);
    target_url.set_query(parts.uri.query());

    let mut forwarded_req = state.client.request(parts.method, target_url);

    if !parts.headers.is_empty() {
        forwarded_req = forwarded_req.headers(parts.headers);
    }
    forwarded_req = forwarded_req.body(body_bytes);

    let resp = match forwarded_req.send().await {
        Ok(resp) => resp,
        Err(e) => {
            error!("Failed to forward request: {}", e);
            return (
                StatusCode::BAD_GATEWAY,
                format!("Failed to forward request: {}", e),
            )
                .into_response();
        }
    };

    let (resp_parts, resp_body_bytes) = match clone_response(resp).await {
        Ok(tuple) => tuple,
        Err(_) => {
            return (
                StatusCode::INTERNAL_SERVER_ERROR,
                "Failed to clone response".to_string(),
            )
                .into_response();
        }
    };

    let decrypted_response_log =
        match process_and_log_response(&resp_body_bytes, &state.key, &state.iv).await {
            Ok(log) => log,
            Err(e) => {
                warn!("Failed to process response for logging: {}", e);
                "Could not decrypt response".to_string()
            }
        };

    let method = decrypted_request_log
        .get("method")
        .and_then(Value::as_str)
        .unwrap_or("UNKNOWN");
    info!(
        "[{}] {}\nRequest:\n{}\nResponse:\n{}\n{}",
        chrono::Local::now().format("%Y/%m/%d %H:%M:%S"),
        method,
        serde_json::to_string_pretty(&decrypted_request_log).unwrap_or_default(),
        decrypted_response_log,
        "-".repeat(80),
    );

    let mut response_builder = Response::builder().status(resp_parts.status);
    if !resp_parts.headers.is_empty() {
        *response_builder.headers_mut().unwrap() = resp_parts.headers;
    }
    response_builder
        .body(axum::body::Body::from(resp_body_bytes))
        .unwrap()
}

fn process_and_log_request(body_bytes: &Bytes, key: &str, iv: &str) -> anyhow::Result<Value> {
    let mut request_data: Value = serde_json::from_slice(body_bytes)?;

    if let Some(params_value) = request_data.get_mut("params")
        && let Some(params_str) = params_value.as_str() {
            let params_str_owned = params_str.to_string();
            match crypto::decrypt(&params_str_owned, key, iv) {
                Ok(decrypted_str) => {
                    let decrypted_params: Value = serde_json::from_str(&decrypted_str)
                        .unwrap_or(Value::String(decrypted_str));
                    *params_value = decrypted_params;
                }
                Err(e) => {
                    *params_value = Value::String(format!("decrypt failed: {}", e));
                }
            }
        }
    Ok(request_data)
}

async fn process_and_log_response(
    body_bytes: &Bytes,
    key: &str,
    iv: &str,
) -> anyhow::Result<String> {
    let decrypted = crypto::decrypt(std::str::from_utf8(body_bytes)?, key, iv)?;
    let formatted: Value = serde_json::from_str(&decrypted)?;
    Ok(serde_json::to_string_pretty(&formatted)?)
}

async fn clone_response(
    resp: reqwest::Response,
) -> anyhow::Result<(axum::http::response::Parts, Bytes)> {
    let mut parts_builder = axum::http::response::Builder::new()
        .status(resp.status())
        .version(resp.version());

    if !resp.headers().is_empty() {
        *parts_builder.headers_mut().unwrap() = resp.headers().clone();
    }
    let parts = parts_builder.body(())?.into_parts().0;
    let body_bytes = resp.bytes().await?;

    Ok((parts, body_bytes))
}