add custom key extractor to use a fly-specific header, disable locally

Author: b4sen

fly-app/src/key_extractor.rs

@@ -0,0 +1,20 @@
+use actix_governor::{KeyExtractor, SimpleKeyExtractionError};
+
+#[derive(Clone, Default)]
+pub struct Extractor {}
+
+impl KeyExtractor for Extractor {
+    type Key = String;
+    type KeyExtractionError = SimpleKeyExtractionError<&'static str>;
+
+    fn extract(
+        &self,
+        req: &actix_web::dev::ServiceRequest,
+    ) -> Result<Self::Key, Self::KeyExtractionError> {
+        let head = req.head();
+        match head.headers().get("Fly-Client-IP") {
+            Some(data) => return Ok(data.to_str().unwrap().to_string()),
+            None => return Err(SimpleKeyExtractionError::new("can not find any token")),
+        }
+    }
+}

fly-app/src/main.rs

@@ -1,5 +1,7 @@
+use std::env;
+
 use actix_governor::{Governor, GovernorConfigBuilder};
-use actix_web::{web, App, HttpResponse, HttpServer};
+use actix_web::{middleware::Condition, web, App, HttpResponse, HttpServer};
 use serde::{Deserialize, Serialize};
 use serde_qs::actix::QsQuery;
 use sqlx::postgres::PgPoolOptions;
@@ -10,6 +12,7 @@ use tracing_actix_web::{DefaultRootSpanBuilder, RequestId, TracingLogger};
 use crate::error::{ErrorResponse, InternalErrorResponse};
 use crate::tracing::init_tracing;
 mod error;
+mod key_extractor;
 mod tracing;
 mod util;
 
@@ -844,7 +847,9 @@ async fn main() -> std::io::Result<()> {
         .parse()
         .expect("PORT must be a valid number");
 
+    let is_fly = env::var("FLY_APP_NAME").is_ok();
     let governor_conf = GovernorConfigBuilder::default()
+        .key_extractor(key_extractor::Extractor::default())
         .seconds_per_request(2)
         .burst_size(5)
         .finish()
@@ -867,13 +872,11 @@ async fn main() -> std::io::Result<()> {
             .route("/", web::get().to(index))
             .service(
                 web::scope("/v1")
-                    .wrap(Governor::new(&governor_conf))
+                    // add rate-limiter only when deployed to fly.io as it fetches a fly-specific header
+                    .wrap(Condition::new(is_fly, Governor::new(&governor_conf)))
                     .route("", web::get().to(index_v1))
                     .route("/wasms", web::get().to(get_wasms))
-                    .route(
-                        "/wasms/{wasm_name}",
-                        web::get().to(get_wasm_root_channel),
-                    )
+                    .route("/wasms/{wasm_name}", web::get().to(get_wasm_root_channel))
                     .route(
                         "/wasms/{channel}/{wasm_name}",
                         web::get().to(get_wasm_latest),