From 0065b0459cec601d7d3fb8ef8ce8b47d0fb2c5ae Mon Sep 17 00:00:00 2001
From: Chris Smith <1979423+chris13524@users.noreply.github.com>
Date: Tue, 12 Dec 2023 15:56:34 -0500
Subject: [PATCH] fix: add x request (#161)

* fix: add timing traces

* fix: add x-request-id

* chore: don't span with state
---
 Cargo.lock                       |  2 ++
 Cargo.toml                       |  2 +-
 justfile                         | 24 +++++++++++++++++-------
 src/handlers/identity/resolve.rs |  2 +-
 src/lib.rs                       | 29 +++++++++++++++++++----------
 5 files changed, 40 insertions(+), 19 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 110bda6..701289d 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3399,9 +3399,11 @@ dependencies = [
  "pin-project-lite",
  "tokio",
  "tokio-util",
+ "tower",
  "tower-layer",
  "tower-service",
  "tracing",
+ "uuid",
 ]
 
 [[package]]
diff --git a/Cargo.toml b/Cargo.toml
index 448768c..1d406e7 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -20,7 +20,7 @@ axum = "0.6.1"
 axum-macros = "0.3.0"
 tokio = { version = "1.0", features = ["full"] }
 tower = { version = "0.4", features = ["util", "timeout", "load-shed", "limit"] }
-tower-http = { version = "0.3.0", features = ["add-extension", "auth", "compression-full", "trace", "cors"] }
+tower-http = { version = "0.3.0", features = ["add-extension", "auth", "compression-full", "trace", "cors", "request-id", "util"] }
 hyper = "0.14"
 http = "0.2.8"
 
diff --git a/justfile b/justfile
index c883ed4..fcc5605 100644
--- a/justfile
+++ b/justfile
@@ -37,10 +37,16 @@ clean:
   @echo '==> Cleaning project target/*'
   cargo clean
 
-# Reformat code
-fmt:
-  @echo '==> Reformatting code'
-  cargo +nightly fmt
+fmt-imports:
+  #!/bin/bash
+  set -euo pipefail
+
+  if command -v cargo-fmt >/dev/null; then
+    echo '==> Running rustfmt'
+    cargo +nightly fmt -- --config group_imports=StdExternalCrate,imports_granularity=One
+  else
+    echo '==> rustfmt not found in PATH, skipping'
+  fi
 
 fmt-imports:
   #!/bin/bash
@@ -96,7 +102,11 @@ restart-keyserver-docker:
   docker-compose -f ./ops/docker-compose.keyserver.yml -f ./ops/docker-compose.storage.yml up -d --build --force-recreate --no-deps keyserver
 
 # Lint the project for any quality issues
-lint: rs-check-fmt check clippy commit-check
+lint: clippy fmt commit-check
+
+unit: lint test test-all test-doc tf-lint
+
+devloop: unit fmt-imports
 
 # Run project linter
 clippy:
@@ -111,13 +121,13 @@ clippy:
   fi
 
 # Run code formatting check
-rs-check-fmt:
+fmt:
   #!/bin/bash
   set -euo pipefail
 
   if command -v cargo-fmt >/dev/null; then
     echo '==> Running rustfmt'
-    cargo fmt -- --check
+    cargo fmt
   else
     echo '==> rustfmt not found in PATH, skipping'
   fi
diff --git a/src/handlers/identity/resolve.rs b/src/handlers/identity/resolve.rs
index b0c24f8..4e245d9 100644
--- a/src/handlers/identity/resolve.rs
+++ b/src/handlers/identity/resolve.rs
@@ -35,7 +35,7 @@ impl From<ResolveIdentityResponse> for Value {
     }
 }
 
-#[instrument(name = "resolve_handler", skip(state))]
+#[instrument(name = "resolve_handler", skip_all)]
 pub async fn handler(
     State(state): State<Arc<AppState>>,
     Query(params): Query<ResolveIdentityPayload>,
diff --git a/src/lib.rs b/src/lib.rs
index 7069656..aacc414 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -14,7 +14,9 @@ use {
     tower::ServiceBuilder,
     tower_http::{
         cors::CorsLayer,
+        request_id::MakeRequestUuid,
         trace::{DefaultMakeSpan, DefaultOnRequest, DefaultOnResponse, TraceLayer},
+        ServiceBuilderExt,
     },
     tracing::Level,
     wc::geoip::{
@@ -63,16 +65,23 @@ pub async fn bootstrap(
 
     let state_arc = Arc::new(state);
 
-    let global_middleware = ServiceBuilder::new().layer(
-        TraceLayer::new_for_http()
-            .make_span_with(DefaultMakeSpan::new().include_headers(true))
-            .on_request(DefaultOnRequest::new().level(Level::INFO))
-            .on_response(
-                DefaultOnResponse::new()
-                    .level(Level::INFO)
-                    .include_headers(true),
-            ),
-    );
+    let global_middleware = ServiceBuilder::new()
+        .set_x_request_id(MakeRequestUuid)
+        .layer(
+            TraceLayer::new_for_http()
+                .make_span_with(
+                    DefaultMakeSpan::new()
+                        .level(Level::INFO)
+                        .include_headers(true),
+                )
+                .on_request(DefaultOnRequest::new().level(Level::INFO))
+                .on_response(
+                    DefaultOnResponse::new()
+                        .level(Level::INFO)
+                        .include_headers(true),
+                ),
+        )
+        .propagate_x_request_id();
 
     let cors_layer = CorsLayer::new()
         .allow_headers([http::header::CONTENT_TYPE])