rust implementation

src/proxy.rs

@@ -0,0 +1,150 @@
+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))
+}