Refactoring how some things work and add doc tests for middleware updates

masongup-mdsol · masongup-mdsol · commit 89a0735b988b · 2025-01-03T14:50:27.000-05:00
diff --git a/Cargo.toml b/Cargo.toml
@@ -26,17 +26,18 @@ dirs = "5"
 chrono = "0.4"
 tokio = { version = "1", features = ["fs"] }
 tower = { version = "0.4", optional = true }
-axum = { version = ">= 0.7.2", optional = true }
+axum = { version = ">= 0.8", optional = true }
 futures-core = { version = "0.3", optional = true }
 http = "1"
 bytes = { version = "1", optional = true }
 thiserror = "1"
 mauth-core = "0.6"
+tracing = { version = "0.1", optional = true }
 
 [dev-dependencies]
 tokio = { version = "1", features = ["rt-multi-thread", "macros"] }
 
 [features]
-axum-service = ["tower", "futures-core", "axum", "bytes"]
+axum-service = ["tower", "futures-core", "axum", "bytes", "tracing"]
 tracing-otel-26 = ["reqwest-tracing/opentelemetry_0_26"]
 tracing-otel-27 = ["reqwest-tracing/opentelemetry_0_27"]
diff --git a/README.md b/README.md
@@ -7,6 +7,8 @@ the MAuth protocol, and verify the responses. Usage example:
 release any code to Production or deploy in a Client-accessible environment without getting
 approval for the full stack used through the Architecture and Security groups.
 
+## Outgoing Requests
+
 ```no_run
 use mauth_client::MAuthInfo;
 use reqwest::Client;
@@ -49,6 +51,8 @@ match client.get("https://www.example.com/").send().await {
 # }
 ```
 
+## Incoming Requests
+
 The optional `axum-service` feature provides for a Tower Layer and Service that will
 authenticate incoming requests via MAuth V2 or V1 and provide to the lower layers a
 validated app_uuid from the request via the `ValidatedRequestDetails` struct. Note that
@@ -66,6 +70,104 @@ a different response, or respond to the lack of the extension in another way, yo
 use a more manual mechanism to check for the extension and decide how to proceed if it
 is not present.
 
+### Examples for `RequiredMAuthValidationLayer`
+
+```no_run
+# async fn run_server() {
+use mauth_client::{
+    axum_service::RequiredMAuthValidationLayer,
+    validate_incoming::ValidatedRequestDetails,
+};
+use axum::{http::StatusCode, Router, routing::get, serve};
+use tokio::net::TcpListener;
+
+// If there is not a valid mauth signature, this function will never run at all, and
+// the request will return an empty 401 Unauthorized
+async fn foo() -> StatusCode {
+    StatusCode::OK
+}
+
+// In addition to returning a 401 Unauthorized without running if there is not a valid
+// MAuth signature, this also makes the validated requesting app UUID available to
+// the function
+async fn bar(details: ValidatedRequestDetails) -> StatusCode {
+    println!("Got a request from app with UUID: {}", details.app_uuid);
+    StatusCode::OK
+}
+
+// This function will run regardless of whether or not there is a mauth signature
+async fn baz() -> StatusCode {
+    StatusCode::OK
+}
+
+// Attaching the baz route handler after the layer means the layer is not run for
+// requests to that path, so no mauth checking will be performed for that route and
+// any other routes attached after the layer
+let router = Router::new()
+    .route("/foo", get(foo))
+    .route("/bar", get(bar))
+    .layer(RequiredMAuthValidationLayer::from_default_file().unwrap())
+    .route("/baz", get(baz));
+let listener = TcpListener::bind("0.0.0.0:3000").await.unwrap();
+serve(listener, router).await.unwrap();
+# }
+```
+
+### Examples for `OptionalMAuthValidationLayer`
+
+```no_run
+# async fn run_server() {
+use mauth_client::{
+    axum_service::OptionalMAuthValidationLayer,
+    validate_incoming::ValidatedRequestDetails,
+};
+use axum::{http::StatusCode, Router, routing::get, serve};
+use tokio::net::TcpListener;
+
+// This request will run no matter what the authorization status is
+async fn foo() -> StatusCode {
+    StatusCode::OK
+}
+
+// If there is not a valid mauth signature, this function will never run at all, and
+// the request will return an empty 401 Unauthorized
+async fn bar(_: ValidatedRequestDetails) -> StatusCode {
+    StatusCode::OK
+}
+
+// In addition to returning a 401 Unauthorized without running if there is not a valid
+// MAuth signature, this also makes the validated requesting app UUID available to
+// the function
+async fn baz(details: ValidatedRequestDetails) -> StatusCode {
+    println!("Got a request from app with UUID: {}", details.app_uuid);
+    StatusCode::OK
+}
+
+// This request will run whether or not there is a valid mauth signature, but the Option
+// provided can be used to tell you whether there was a valid signature, so you can
+// implement things like multiple possible types of authentication or behavior other than
+// a 401 return if there is no authentication
+async fn bam(optional_details: Option<ValidatedRequestDetails>) -> StatusCode {
+    match optional_details {
+        Some(details) => println!("Got a request from app with UUID: {}", details.app_uuid),
+        None => println!("Got a request without a valid mauth signature"),
+    }
+    StatusCode::OK
+}
+
+let router = Router::new()
+    .route("/foo", get(foo))
+    .route("/bar", get(bar))
+    .route("/baz", get(baz))
+    .route("/bam", get(bam))
+    .layer(OptionalMAuthValidationLayer::from_default_file().unwrap());
+let listener = TcpListener::bind("0.0.0.0:3000").await.unwrap();
+serve(listener, router).await.unwrap();
+# }
+```
+
+### OpenTelemetry Integration
+
 There are also optional features `tracing-otel-26` and `tracing-otel-27` that pair with
 the `axum-service` feature to ensure that any outgoing requests for credentials that take
 place in the context of an incoming web request also include the proper OpenTelemetry span
diff --git a/src/axum_service.rs b/src/axum_service.rs
@@ -1,11 +1,17 @@
 //! Structs and impls related to providing a Tower Service and Layer to verify incoming requests
 
-use axum::extract::{FromRequestParts, Request};
+use axum::{
+    body::Body,
+    extract::{FromRequestParts, OptionalFromRequestParts, Request},
+    response::IntoResponse,
+};
 use futures_core::future::BoxFuture;
-use http::{request::Parts, StatusCode};
+use http::{request::Parts, Response, StatusCode};
+use std::convert::Infallible;
 use std::error::Error;
 use std::task::{Context, Poll};
 use tower::{Layer, Service};
+use tracing::error;
 
 use crate::validate_incoming::ValidatedRequestDetails;
 use crate::{
@@ -27,24 +33,31 @@ where
     S: Service<Request> + Send + Clone + 'static,
     S::Future: Send + 'static,
     S::Error: Into<Box<dyn Error + Sync + Send>>,
+    S::Response: Into<Response<Body>>,
 {
-    type Response = S::Response;
-    type Error = Box<dyn Error + Sync + Send>;
+    type Response = Response<Body>;
+    type Error = S::Error;
     type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
 
     fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
-        self.service.poll_ready(cx).map_err(|e| e.into())
+        self.service.poll_ready(cx)
     }
 
     fn call(&mut self, request: Request) -> Self::Future {
         let mut cloned = self.clone();
         Box::pin(async move {
             match cloned.mauth_info.validate_request(request).await {
                 Ok(valid_request) => match cloned.service.call(valid_request).await {
-                    Ok(response) => Ok(response),
-                    Err(err) => Err(err.into()),
+                    Ok(response) => Ok(response.into()),
+                    Err(err) => Err(err),
                 },
-                Err(err) => Err(Box::new(err) as Box<dyn Error + Send + Sync>),
+                Err(err) => {
+                    error!(
+                        error = ?err,
+                        "Failed to validate MAuth signature, rejecting request"
+                    );
+                    Ok(StatusCode::UNAUTHORIZED.into_response())
+                }
             }
         })
     }
@@ -121,23 +134,18 @@ where
     S::Error: Into<Box<dyn Error + Sync + Send>>,
 {
     type Response = S::Response;
-    type Error = Box<dyn Error + Sync + Send>;
+    type Error = S::Error;
     type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
 
     fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
-        self.service.poll_ready(cx).map_err(|e| e.into())
+        self.service.poll_ready(cx)
     }
 
     fn call(&mut self, request: Request) -> Self::Future {
         let mut cloned = self.clone();
         Box::pin(async move {
-            match cloned.mauth_info.validate_request_optionally(request).await {
-                Ok(valid_request) => match cloned.service.call(valid_request).await {
-                    Ok(response) => Ok(response),
-                    Err(err) => Err(err.into()),
-                },
-                Err(err) => Err(Box::new(err) as Box<dyn Error + Send + Sync>),
-            }
+            let processed_request = cloned.mauth_info.validate_request_optionally(request).await;
+            cloned.service.call(processed_request).await
         })
     }
 }
@@ -192,8 +200,10 @@ impl OptionalMAuthValidationLayer {
     }
 }
 
-#[async_trait::async_trait]
-impl<S> FromRequestParts<S> for ValidatedRequestDetails {
+impl<S> FromRequestParts<S> for ValidatedRequestDetails
+where
+    S: Send + Sync,
+{
     type Rejection = StatusCode;
 
     async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
@@ -204,3 +214,17 @@ impl<S> FromRequestParts<S> for ValidatedRequestDetails {
             .ok_or(StatusCode::UNAUTHORIZED)
     }
 }
+
+impl<S> OptionalFromRequestParts<S> for ValidatedRequestDetails
+where
+    S: Send + Sync,
+{
+    type Rejection = Infallible;
+
+    async fn from_request_parts(
+        parts: &mut Parts,
+        _state: &S,
+    ) -> Result<Option<Self>, Self::Rejection> {
+        Ok(parts.extensions.get::<ValidatedRequestDetails>().cloned())
+    }
+}
diff --git a/src/validate_incoming.rs b/src/validate_incoming.rs
@@ -1,8 +1,10 @@
 use crate::{MAuthInfo, CLIENT, PUBKEY_CACHE};
 use axum::extract::Request;
+use bytes::Bytes;
 use chrono::prelude::*;
 use mauth_core::verifier::Verifier;
 use thiserror::Error;
+use tracing::error;
 use uuid::Uuid;
 
 /// This struct holds the app UUID for a validated request. It is meant to be used with the
@@ -59,15 +61,26 @@ impl MAuthInfo {
         }
     }
 
-    pub(crate) async fn validate_request_optionally(
-        &self,
-        req: Request,
-    ) -> Result<Request, axum::Error> {
+    pub(crate) async fn validate_request_optionally(&self, req: Request) -> Request {
         let (mut parts, body) = req.into_parts();
         if parts.headers.contains_key(MAUTH_V2_SIGNATURE_HEADER)
             || parts.headers.contains_key(MAUTH_V1_SIGNATURE_HEADER)
         {
-            let body_bytes = axum::body::to_bytes(body, usize::MAX).await?;
+            // By my reading of the code for this it should never fail, since we are passing
+            // MAX for the limit. But just to be safe, we will log the error and proceed with
+            // an empty body just in case instead of unwrapping. This would cause the body to
+            // be unavailable to the lower layers, but they would probably also fail to get it
+            // anyways since we just did here.
+            let body_bytes = match axum::body::to_bytes(body, usize::MAX).await {
+                Ok(bytes) => bytes,
+                Err(err) => {
+                    error!(
+                        error = ?err,
+                        "Failed to retrieve request body, continuing with empty body"
+                    );
+                    Bytes::new()
+                }
+            };
 
             match self.validate_request_v2(&parts, &body_bytes).await {
                 Ok(host_app_uuid) => {
@@ -95,9 +108,9 @@ impl MAuthInfo {
 
             let new_body = axum::body::Body::from(body_bytes);
             let new_request = Request::from_parts(parts, new_body);
-            Ok(new_request)
+            new_request
         } else {
-            Ok(Request::from_parts(parts, body))
+            Request::from_parts(parts, body)
         }
     }
 
