From 6a7869a80cd0633ebc512711cda29ab097fc48e1 Mon Sep 17 00:00:00 2001
From: Dorian Pula <dorian.pula@points.com>
Date: Wed, 7 Feb 2024 10:12:40 -0500
Subject: [PATCH] Break out into using axum extractor. Add support for JSON
 extraction.

---
 Cargo.toml  |  1 +
 README.md   |  4 +++-
 src/main.rs | 57 ++++++++++++++++++++++++++++++++---------------------
 3 files changed, 39 insertions(+), 23 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 49d8d59..f80cc59 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -19,4 +19,5 @@ section = "web"
 axum = "0.7"
 clap = { version = "4.4", features = ["derive"] }
 serde = { version = "1.0", features = ["derive"] }
+serde_json = "1.0.113"
 tokio = { version = "1.36.0", features = ["macros", "rt-multi-thread"] }
diff --git a/README.md b/README.md
index 55522d1..8f0c175 100644
--- a/README.md
+++ b/README.md
@@ -11,6 +11,7 @@ A simple server for mirroring HTTP requests for testing.
 * Create a DEB package:
   * Install cargo-deb: `cargo install cargo-deb`
   * Create the DEB package: `cargo deb`
+* Faster builds using [cargo-watch](https://watchexec.github.io/#cargo-watch): `cargo watch -x run`
 
 ## Install
 
@@ -25,7 +26,8 @@ sudo dpkg -i "mirror-server-${VERSION}_amd64.deb"
 
 ## TODO
 
-* [ ] Migrate actix to axum for easier maintainability.
+* [x] Migrate actix to axum for easier maintainability.
+* [x] Add mirroring of JSON request.
 * [ ] Add logging to server.
 * [ ] Add convenience path to access logging for a certain date / records.
 * [ ] Dockerize mirror-server for easier distribution.
diff --git a/src/main.rs b/src/main.rs
index 77e5d9c..292a60c 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,9 +1,11 @@
 use axum::{
-    extract::{Json, Request},
+    extract::{Host, OriginalUri, Json},
+    http::{header::HeaderMap, Method},
     Router,
 };
 use clap::Parser;
 use serde::Serialize;
+use serde_json::{json, Value};
 use std::collections::BTreeMap;
 
 #[derive(Parser)]
@@ -30,29 +32,40 @@ struct EchoResponse {
     path: String,
     host: String,
     headers: BTreeMap<String, String>,
+    body: Value,
 }
 
-impl EchoResponse {
-    fn new(req: &Request) -> Self {
-        let req_uri = req.uri();
-        let mut headers = BTreeMap::new();
-        for (header_name, header_value) in req.headers().iter() {
-            headers.insert(
-                header_name.to_string(),
-                header_value.to_str().unwrap_or("ERROR").to_string(),
-            );
-        }
-        EchoResponse {
-            method: req.method().to_string(),
-            path: req_uri.path().to_string(),
-            host: req_uri.host().unwrap_or("").to_string(),
-            headers,
-        }
-    }
-}
+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();
+    // TODO: Logging the method and path.
 
-async fn echo_request(request: Request) -> Json<EchoResponse> {
-    Json(EchoResponse::new(&request))
+    let json_body = match body {
+        None => {
+            // TODO: Adding logging.
+            json!({"error": "Non-JSON request sent."})
+        },
+        Some(Json(value)) => value,
+    };
+    let response = EchoResponse {
+        method,
+        host,
+        path,
+        headers,
+        body: json_body,
+    };
+    Json(response)
 }
 
 #[tokio::main]
@@ -64,7 +77,7 @@ async fn main() {
     let app = Router::new().fallback(echo_request);
     let listener = tokio::net::TcpListener::bind(&listen_on)
         .await
-        .expect(&format!("Attempted binding to {}", listen_on));
+        .unwrap_or_else(|_| panic!("Attempted binding to {}", listen_on));
     axum::serve(listener, app)
         .await
         .expect("Server should start");