Add auth middleware

Author: tcerqueira

lambda/benches/mix_node.rs

@@ -78,7 +78,7 @@ fn bench_requests(c: &mut Criterion) {
     let enc_key = EncryptionKey::from(&Scalar::random(&mut rng) * &GENERATOR_TABLE);
 
     let rt = tokio::runtime::Runtime::new().unwrap();
-    let test_app = rt.block_on(async { testing::create_app().await });
+    let test_app = rt.block_on(async { testing::create_app(None).await });
     let client = Arc::new(reqwest::Client::new());
 
     let payload = Arc::new(EncryptedCodes {

lambda/src/bin/mix_node.rs

@@ -1,4 +1,4 @@
-use lambda::mix_node;
+use lambda::mix_node::{self, AppState};
 
 use mimalloc::MiMalloc as GlobalAllocator;
 
@@ -17,5 +17,6 @@ async fn main() -> Result<(), lambda_http::Error> {
 
     lambda_http::tracing::init_default_subscriber();
 
-    lambda_http::run(mix_node::app()).await
+    let state = AppState::new(std::env::var("AUTH_TOKEN").ok());
+    lambda_http::run(mix_node::app(state)).await
 }

lambda/src/bin/simple_server.rs

@@ -1,4 +1,4 @@
-use lambda::mix_node;
+use lambda::mix_node::{self, AppState};
 
 use mimalloc::MiMalloc as GlobalAllocator;
 
@@ -11,6 +11,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let port = listener.local_addr().unwrap().port();
 
     println!("Listening on http://localhost:{port}...");
-    axum::serve(listener, mix_node::app()).await.unwrap();
+    let state = AppState::new(std::env::var("AUTH_TOKEN").ok());
+    axum::serve(listener, mix_node::app(state)).await.unwrap();
     Ok(())
 }

lambda/src/lib.rs

@@ -6,20 +6,38 @@ pub const N_BITS: usize = 25600;
 pub mod mix_node {
     use super::*;
 
-    use axum::extract::DefaultBodyLimit;
+    use axum::extract::{DefaultBodyLimit, Request, State};
+    use axum::middleware::{self, Next};
+    use axum::response::{IntoResponse, Response};
     use axum::{response::Json, routing::post, Router};
 
-    use axum::http::StatusCode;
+    use axum::http::{HeaderMap, StatusCode};
     use rand::{rngs::StdRng, SeedableRng};
     use rust_elgamal::{Ciphertext, EncryptionKey, RistrettoPoint};
     use serde::{Deserialize, Deserializer, Serialize};
     use std::sync::OnceLock;
 
-    pub fn app() -> Router {
+    #[derive(Debug, Clone)]
+    pub struct AppState {
+        auth_token: Option<String>,
+    }
+
+    impl AppState {
+        pub fn new(auth_token: Option<String>) -> Self {
+            Self { auth_token }
+        }
+    }
+
+    pub fn app(state: AppState) -> Router {
         Router::new()
             .route("/remix", post(remix_handler))
+            .layer(middleware::from_fn_with_state(
+                state.clone(),
+                auth_middleware,
+            ))
             // TODO: for security reasons set max instead of disabling (measured payload was ~11MB)
             .layer(DefaultBodyLimit::disable())
+            .with_state(state)
     }
 
     #[derive(Debug, Serialize, Deserialize)]
@@ -51,6 +69,25 @@ pub mod mix_node {
         Ok(Json(codes))
     }
 
+    async fn auth_middleware(
+        State(AppState { auth_token }): State<AppState>,
+        headers: HeaderMap,
+        request: Request,
+        next: Next,
+    ) -> Response {
+        let next_run = async { next.run(request).await };
+
+        match (auth_token, headers.get("X-AUTH-TOKEN")) {
+            // AUTH_TOKEN is set on the server and in the request header so we check
+            (Some(auth_token), Some(header_auth_token)) if auth_token == *header_auth_token => {
+                next_run.await
+            }
+            // AUTH_TOKEN is not set on the server so we disable auth
+            (None, _) => next_run.await,
+            _ => (StatusCode::UNAUTHORIZED, "Unauthorized").into_response(),
+        }
+    }
+
     fn enc_key() -> &'static EncryptionKey {
         // TODO: remove hardcoded encryption key from a fixed seed
         static ENC_KEY: OnceLock<EncryptionKey> = OnceLock::new();

lambda/src/testing.rs

@@ -1,18 +1,19 @@
 use tokio::task::JoinHandle;
 
-use crate::mix_node;
+use crate::mix_node::{self, AppState};
 
 pub struct TestApp {
     pub port: u16,
     pub join_handle: JoinHandle<()>,
 }
 
-pub async fn create_app() -> TestApp {
+pub async fn create_app(auth_token: Option<String>) -> TestApp {
     let listener = tokio::net::TcpListener::bind("0.0.0.0:0").await.unwrap();
     let port = listener.local_addr().unwrap().port();
 
     let join_handle = tokio::spawn(async move {
-        axum::serve(listener, mix_node::app()).await.unwrap();
+        let state = AppState::new(auth_token);
+        axum::serve(listener, mix_node::app(state)).await.unwrap();
     });
 
     TestApp { port, join_handle }

lambda/tests/mix_node_integration.rs

@@ -43,7 +43,7 @@ fn set_up_payload() -> (EncryptedCodes, DecryptionKey) {
 
 #[tokio::test]
 async fn test_mix_node() -> Result<(), Box<dyn Error>> {
-    let TestApp { port, .. } = testing::create_app().await;
+    let TestApp { port, .. } = testing::create_app(None).await;
 
     let (codes, dec_key) = set_up_payload();
 
@@ -78,7 +78,7 @@ async fn test_mix_node() -> Result<(), Box<dyn Error>> {
 
 #[tokio::test]
 async fn test_mix_node_bad_request() -> Result<(), Box<dyn Error>> {
-    let TestApp { port, .. } = testing::create_app().await;
+    let TestApp { port, .. } = testing::create_app(None).await;
 
     let (mut codes, _dec_key) = set_up_payload();
     // Remove elements to cause a size mismatch
@@ -98,6 +98,44 @@ async fn test_mix_node_bad_request() -> Result<(), Box<dyn Error>> {
     Ok(())
 }
 
+#[tokio::test]
+async fn test_mix_node_unauthorized() -> Result<(), Box<dyn Error>> {
+    let TestApp { port, .. } =
+        testing::create_app(Some("test_mix_node_unauthorized".to_string())).await;
+
+    // Bad request + Serialization
+    let client = reqwest::Client::new();
+    let response = client
+        .post(format!("http://localhost:{port}/remix"))
+        .send()
+        .await?;
+
+    // Assert
+    assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
+    Ok(())
+}
+
+#[tokio::test]
+async fn test_mix_node_authorized() -> Result<(), Box<dyn Error>> {
+    let auth_token = "test_mix_node_authorized";
+    let TestApp { port, .. } = testing::create_app(Some(auth_token.to_string())).await;
+
+    let (codes, _dec_key) = set_up_payload();
+
+    // Bad request + Serialization
+    let client = reqwest::Client::new();
+    let response = client
+        .post(format!("http://localhost:{port}/remix"))
+        .header("X-AUTH-TOKEN", auth_token)
+        .json(&codes)
+        .send()
+        .await?;
+
+    // Assert
+    assert_eq!(response.status(), StatusCode::OK);
+    Ok(())
+}
+
 #[test]
 fn test_encode_bits() {
     let bits = BitVec::<u8, Msb0>::from_slice(&[0b11100100]);