From feecd42df46d4ed3ff7f7983418d45bf5045fa44 Mon Sep 17 00:00:00 2001
From: Xavier Basty <xavier@hexbee.net>
Date: Mon, 2 Oct 2023 10:18:28 +0200
Subject: [PATCH] fix: apply geoblock at beginning of chain (#143)

---
 src/lib.rs | 66 +++++++++++++++++++++---------------------------------
 1 file changed, 26 insertions(+), 40 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 35b0ad6..a9c5d9d 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -3,7 +3,6 @@ use {
     aws_config::meta::region::RegionProviderChain,
     aws_sdk_s3::{config::Region, Client as S3Client},
     axum::{
-        body::HttpBody,
         routing::{get, post},
         Router,
     },
@@ -79,26 +78,32 @@ pub async fn bootstrap(
         .allow_origin("*".parse::<HeaderValue>().unwrap())
         .allow_methods([Method::GET, Method::POST]);
 
-    let app = new_geoblocking_router(
-        geoip_resolver.clone(),
-        state_arc.config.blocked_countries.clone(),
-    )
-    .route("/health", get(handlers::health::handler))
-    .route(
-        "/identity",
-        get(handlers::identity::resolve::handler)
-            .post(handlers::identity::register::handler)
-            .delete(handlers::identity::unregister::handler),
-    )
-    .route(
-        "/invite",
-        post(handlers::invite::register::handler)
-            .delete(handlers::invite::unregister::handler)
-            .get(handlers::invite::resolve::handler),
-    )
-    .layer(global_middleware)
-    .layer(cors_layer)
-    .with_state(state_arc.clone());
+    let app = Router::new()
+        .route("/health", get(handlers::health::handler))
+        .route(
+            "/identity",
+            get(handlers::identity::resolve::handler)
+                .post(handlers::identity::register::handler)
+                .delete(handlers::identity::unregister::handler),
+        )
+        .route(
+            "/invite",
+            post(handlers::invite::register::handler)
+                .delete(handlers::invite::unregister::handler)
+                .get(handlers::invite::resolve::handler),
+        )
+        .layer(global_middleware)
+        .layer(cors_layer);
+    let app = if let Some(resolver) = geoip_resolver {
+        app.layer(GeoBlockLayer::new(
+            resolver.clone(),
+            state_arc.config.blocked_countries.clone(),
+            BlockingPolicy::AllowAll,
+        ))
+    } else {
+        app
+    };
+    let app = app.with_state(state_arc.clone());
 
     let private_app = Router::new()
         .route("/metrics", get(handlers::metrics::handler))
@@ -116,25 +121,6 @@ pub async fn bootstrap(
     Ok(())
 }
 
-fn new_geoblocking_router<S, B>(
-    geoip_resolver: Option<Arc<MaxMindResolver>>,
-    blocked_countries: Vec<String>,
-) -> Router<S, B>
-where
-    S: Clone + Send + Sync + 'static,
-    B: HttpBody + Send + 'static,
-{
-    if let Some(resolver) = geoip_resolver {
-        Router::new().layer(GeoBlockLayer::new(
-            resolver.clone(),
-            blocked_countries.clone(),
-            BlockingPolicy::AllowAll,
-        ))
-    } else {
-        Router::new()
-    }
-}
-
 async fn get_s3_client(config: &Configuration) -> S3Client {
     let region_provider = RegionProviderChain::first_try(Region::new("eu-central-1"));
     let shared_config = aws_config::from_env().region(region_provider).load().await;