feat: middleware (WIP)

src/middleware.rs

@@ -0,0 +1,178 @@
+use std::str;
+use std::sync::Arc;
+
+use axum::{
+    extract::{OriginalUri, State},
+    http::{Request, StatusCode},
+    middleware::Next,
+    response::{IntoResponse, Response},
+};
+use http_body_util::BodyExt;
+use serde_json::Value;
+use tracing::{info, warn};
+
+use crate::{AppState, crypto};
+
+enum ResponseBody {
+    Original(String),
+    Modified(Value),
+}
+
+pub async fn log_middleware(
+    State(state): State<Arc<AppState>>,
+    OriginalUri(uri): OriginalUri,
+    req: Request<axum::body::Body>,
+    next: Next,
+) -> impl IntoResponse {
+    let (parts, body) = req.into_parts();
+    let body_bytes = match body.collect().await {
+        Ok(body) => body.to_bytes(),
+        Err(e) => {
+            warn!("Failed to read request body: {}", e);
+            return (
+                StatusCode::INTERNAL_SERVER_ERROR,
+                "Failed to read request body".to_string(),
+            )
+                .into_response();
+        }
+    };
+
+    let path = uri.path();
+
+    let (decrypted_request_log, method) = match str::from_utf8(&body_bytes)
+        .map_err(anyhow::Error::from)
+        .and_then(|body| process_and_log_request(body, &state.key, &state.iv))
+    {
+        Ok(request_data) => {
+            let method = request_data
+                .get("method")
+                .and_then(Value::as_str)
+                .map(str::to_string);
+            (request_data, method)
+        }
+        Err(e) => {
+            warn!("Failed to process request for logging: {}", e);
+            let val = Value::String("Could not decrypt request".to_string());
+            (val, None)
+        }
+    };
+
+    let req = Request::from_parts(parts, axum::body::Body::from(body_bytes));
+    let res = next.run(req).await;
+
+    let (resp_parts, body_bytes) = {
+        let (parts, body) = res.into_parts();
+        let bytes = match body.collect().await {
+            Ok(b) => b.to_bytes(),
+            Err(e) => {
+                warn!("Failed to read response body: {}", e);
+                return (
+                    StatusCode::INTERNAL_SERVER_ERROR,
+                    "Failed to read response body".to_string(),
+                )
+                    .into_response();
+            }
+        };
+        (parts, bytes)
+    };
+    let resp_body_text = String::from_utf8(body_bytes.clone().to_vec()).unwrap_or_default();
+
+    let response_body_to_log = if Some("com.linspirer.device.getcommand") == method.as_deref() {
+        match handle_getcommand_response(&resp_body_text, &state).await {
+            Ok(new_body) => ResponseBody::Modified(new_body),
+            Err(e) => {
+                warn!(
+                    "Failed to handle getcommand response: {}. Responding with empty command list.",
+                    e
+                );
+                let mut empty_response =
+                    serde_json::from_str::<Value>(&resp_body_text).unwrap_or(Value::Null);
+                if let Some(obj) = empty_response.as_object_mut() {
+                    obj.insert("result".to_string(), Value::Array(vec![]));
+                }
+                ResponseBody::Modified(empty_response)
+            }
+        }
+    } else {
+        ResponseBody::Original(resp_body_text)
+    };
+
+    let (decrypted_response_for_log, final_response_body) = match response_body_to_log {
+        ResponseBody::Original(body_text) => {
+            let decrypted = decrypt_and_format(&body_text, &state.key, &state.iv)
+                .unwrap_or_else(|_| "Could not decrypt or format response".to_string());
+            (decrypted, body_text)
+        }
+        ResponseBody::Modified(body_value) => {
+            let pretty_printed = serde_json::to_string_pretty(&body_value).unwrap_or_default();
+            let encrypted = crypto::encrypt(&pretty_printed, &state.key, &state.iv)
+                .unwrap_or_else(|_| "Failed to encrypt modified response".to_string());
+            (pretty_printed, encrypted)
+        }
+    };
+
+    info!(
+        "{}\nRequest:\n{}\nResponse:\n{}\n{}",
+        path,
+        serde_json::to_string_pretty(&decrypted_request_log).unwrap_or_default(),
+        decrypted_response_for_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(final_response_body))
+        .unwrap()
+}
+
+fn process_and_log_request(body: &str, key: &str, iv: &str) -> anyhow::Result<Value> {
+    let mut request_data: Value = serde_json::from_str(body)?;
+
+    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 handle_getcommand_response(body_text: &str, state: &Arc<AppState>) -> anyhow::Result<Value> {
+    let decrypted = crypto::decrypt(body_text, &state.key, &state.iv)?;
+    let mut response_json: Value = serde_json::from_str(&decrypted)?;
+
+    if let Some(result) = response_json.get("result")
+        && let Some(commands) = result.as_array()
+        && !commands.is_empty()
+    {
+        let mut queue = state.command_queue.write().await;
+        for cmd in commands {
+            queue.push(cmd.clone());
+            info!("Added command to the queue: {:?}", cmd);
+        }
+    }
+
+    if let Some(obj) = response_json.as_object_mut() {
+        obj.insert("result".to_string(), Value::Array(vec![]));
+    }
+
+    Ok(response_json)
+}
+
+fn decrypt_and_format(body_text: &str, key: &str, iv: &str) -> anyhow::Result<String> {
+    let decrypted = crypto::decrypt(body_text, key, iv)?;
+    let formatted: Value = serde_json::from_str(&decrypted)?;
+    Ok(serde_json::to_string_pretty(&formatted)?)
+}