middleware/block_traffic: Implement block_by_ip() middleware fn (#7435)

Author: Turbo87

src/config/server.rs

@@ -37,6 +37,7 @@ pub struct Server {
     pub rate_limiter: HashMap<LimitedAction, RateLimiterConfig>,
     pub new_version_rate_limit: Option<u32>,
     pub blocked_traffic: Vec<(String, Vec<String>)>,
+    pub blocked_ips: HashSet<IpAddr>,
     pub max_allowed_page_offset: u32,
     pub page_offset_ua_blocklist: Vec<String>,
     pub page_offset_cidr_blocklist: Vec<IpNetwork>,
@@ -113,6 +114,15 @@ impl Server {
 
         let port = env_optional("PORT").unwrap_or(8888);
 
+        let blocked_ips = match env_optional::<String>("BLOCKED_IPS") {
+            None => HashSet::new(),
+            Some(s) if s.is_empty() => HashSet::new(),
+            Some(s) => s
+                .split(',')
+                .map(|s| s.trim().parse())
+                .collect::<Result<_, _>>()?,
+        };
+
         let allowed_origins = AllowedOrigins::from_default_env()?;
         let page_offset_ua_blocklist = match env_optional::<String>("WEB_PAGE_OFFSET_UA_BLOCKLIST")
         {
@@ -190,6 +200,7 @@ impl Server {
             rate_limiter,
             new_version_rate_limit: env_optional("MAX_NEW_VERSIONS_DAILY"),
             blocked_traffic: blocked_traffic(),
+            blocked_ips,
             max_allowed_page_offset: env_optional("WEB_MAX_ALLOWED_PAGE_OFFSET").unwrap_or(200),
             page_offset_ua_blocklist,
             page_offset_cidr_blocklist,

src/middleware.rs

@@ -62,6 +62,10 @@ pub fn apply_axum_middleware(state: AppState, router: Router<(), TimeoutBody<Bod
             state.clone(),
             require_user_agent::require_user_agent,
         ))
+        .layer(from_fn_with_state(
+            state.clone(),
+            block_traffic::block_by_ip,
+        ))
         .layer(from_fn_with_state(
             state.clone(),
             block_traffic::block_by_header,

src/middleware/block_traffic.rs

@@ -1,21 +1,20 @@
-//! Middleware that blocks requests if a header matches the given list
-//!
-//! To use, set the `BLOCKED_TRAFFIC` environment variable to a comma-separated list of pairs
-//! containing a header name, an equals sign, and the name of another environment variable that
-//! contains the values of that header that should be blocked. For example, set `BLOCKED_TRAFFIC`
-//! to `User-Agent=BLOCKED_UAS,X-Real-Ip=BLOCKED_IPS`, `BLOCKED_UAS` to `curl/7.54.0,cargo 1.36.0
-//! (c4fcfb725 2019-05-15)`, and `BLOCKED_IPS` to `192.168.0.1,127.0.0.1` to block requests from
-//! the versions of curl or Cargo specified or from either of the IPs (values are nonsensical
-//! examples). Values of the headers must match exactly.
-
 use crate::app::AppState;
 use crate::middleware::log_request::RequestLogExt;
+use crate::middleware::real_ip::RealIp;
 use crate::util::errors::RouteBlocked;
-use axum::extract::MatchedPath;
+use axum::extract::{Extension, MatchedPath};
 use axum::middleware::Next;
 use axum::response::IntoResponse;
 use http::StatusCode;
 
+/// Middleware that blocks requests if a header matches the given list
+///
+/// To use, set the `BLOCKED_TRAFFIC` environment variable to a comma-separated list of pairs
+/// containing a header name, an equals sign, and the name of another environment variable that
+/// contains the values of that header that should be blocked. For example, set `BLOCKED_TRAFFIC`
+/// to `User-Agent=BLOCKED_UAS` and `BLOCKED_UAS` to `curl/7.54.0,cargo 1.36.0 (c4fcfb725 2019-05-15)`
+/// to block requests from the versions of curl or Cargo specified (values are nonsensical examples).
+/// Values of the headers must match exactly.
 pub async fn block_by_header<B>(
     state: AppState,
     req: http::Request<B>,
@@ -40,6 +39,19 @@ pub async fn block_by_header<B>(
     next.run(req).await
 }
 
+pub async fn block_by_ip<B>(
+    Extension(real_ip): Extension<RealIp>,
+    state: AppState,
+    req: http::Request<B>,
+    next: Next<B>,
+) -> axum::response::Response {
+    if state.config.blocked_ips.contains(&real_ip) {
+        rejection_response_from(&state, &req).into_response()
+    } else {
+        next.run(req).await
+    }
+}
+
 fn rejection_response_from<B>(state: &AppState, req: &http::Request<B>) -> impl IntoResponse {
     let domain_name = &state.config.domain_name;
 

src/tests/server.rs

@@ -1,5 +1,6 @@
 use crate::builders::*;
 use crate::util::*;
+use std::collections::HashSet;
 
 use ::insta::assert_display_snapshot;
 use http::{header, Method, StatusCode};
@@ -76,3 +77,16 @@ fn block_traffic_via_arbitrary_header_and_value() {
     let resp = anon.run::<()>(req);
     assert_eq!(resp.status(), StatusCode::FOUND);
 }
+
+#[test]
+fn block_traffic_via_ip() {
+    let (_app, anon) = TestApp::init()
+        .with_config(|config| {
+            config.blocked_ips = HashSet::from(["127.0.0.1".parse().unwrap()]);
+        })
+        .empty();
+
+    let resp = anon.get::<()>("/api/v1/crates");
+    assert_eq!(resp.status(), StatusCode::FORBIDDEN);
+    assert_display_snapshot!(resp.into_text());
+}

src/tests/snapshots/all__server__block_traffic_via_ip.snap

@@ -0,0 +1,5 @@
+---
+source: src/tests/server.rs
+expression: resp.into_text()
+---
+We are unable to process your request at this time. This usually means that you are in violation of our crawler policy (https://crates.io/policies#crawlers). Please open an issue at https://github.com/rust-lang/crates.io or email [email protected] and provide the request id 

src/tests/util/test_app.rs

@@ -411,6 +411,7 @@ fn simple_config() -> config::Server {
         rate_limiter: Default::default(),
         new_version_rate_limit: Some(10),
         blocked_traffic: Default::default(),
+        blocked_ips: Default::default(),
         max_allowed_page_offset: 200,
         page_offset_ua_blocklist: vec![],
         page_offset_cidr_blocklist: vec![],