Add WAF CountryCode matching, migrate middlewares, rework tests

Cargo.toml

@@ -17,13 +17,13 @@ categories = []
 ahash = "0.8.11"
 anyhow = "1.0.93"
 arc-swap = "1.7.1"
-arrayvec = "0.7"
+arrayvec = { version = "0.7", features = ["serde"] }
 async-channel = "2.3.1"
 async-trait = "0.1.83"
 axum = "0.8.1"
 axum-extra = { version = "0.12.6", features = ["middleware"] }
 axum-server = "0.8.0"
-base64 = "0.22.1"
+base64 = "0.23.1"
 bytes = "1.10.0"
 candid = "0.10.10"
 candid_parser = "0.4"
@@ -69,24 +69,27 @@ ic-cdk-timers = "1.0"
 ic-custom-domains-canister-api = { version = "0.2", path = "./custom-domains/api" }
 ic-http-types = "0.1.0"
 ic-stable-structures = "0.7.0"
+ic-transport-types = "0.49"
 ic-verify-bls-signature = "0.6"
+ipnet = { version = "2", features = ["serde"] }
 instant-acme = { version = "0.8.5", default-features = false, features = [
     "aws-lc-rs",
     "hyper-rustls",
 ] }
 itertools = "0.15.0"
-mail-auth = "0.11.0"
+mail-auth = "0.12.0"
 mail-parser = { version = "0.11.3", features = ["full_encoding"] }
 mail-send = { version = "0.6.0", default-features = false, features = [
     "builder",
 ] }
-mockall = "0.13.0"
+maxminddb = "0.29"
+mockall = "0.15"
 mock-io = { version = "0.3.2", features = ["full"] }
 moka = { version = "0.12.15", features = ["sync"] }
 nix = { version = "0.31.0", features = ["signal"] }
 ppp = "2.3.0"
 parse-size = { version = "1.1.0", features = ["std"] }
-pem = "3.0.5"
+pem = "4"
 pocket-ic = "=12.0.0"
 prometheus = "0.14.0"
 prost = { version = "0.14.1" }
@@ -155,13 +158,12 @@ tokio-io-timeout = "1.2.0"
 tokio-test = "0.4"
 tokio-util = { version = "0.7.12", features = ["full"] }
 tower = { version = "0.5.1", features = ["util"] }
-tower_governor = { version = "0.8" }
 tower-service = "0.3.3"
 tower-http = { version = "0.7", features = ["trace"] }
 tracing = { version = "0.1.40", features = ["attributes"] }
 url = "2.5.3"
 utoipa = { version = "5.4.0", features = ["axum_extras"] }
-utoipa-swagger-ui = { version = "9.0.2", features = ["axum"] }
+utoipa-swagger-ui = { version = "9.0.2", features = ["axum", "vendored"] }
 # DO NOT upgrade, this breaks monorepo compatibility
 # Read https://github.com/uuid-rs/uuid/releases/tag/1.13.0
 uuid = { version = "=1.12.1", features = ["v7", "serde"] }

ic-bn-lib/Cargo.toml

@@ -27,7 +27,11 @@ acme-dns = ["acme"]
 acme-alpn = ["acme", "dep:rustls-acme"]
 clients-hyper = ["dep:hyper-rustls"]
 custom-domains = ["acme-dns"]
-custom-domains-openapi = ["custom-domains", "dep:utoipa", "dep:utoipa-swagger-ui"]
+custom-domains-openapi = [
+    "custom-domains",
+    "dep:utoipa",
+    "dep:utoipa-swagger-ui",
+]
 cert-providers = ["custom-domains"]
 lb = []
 pubsub = []
@@ -68,11 +72,14 @@ hyper-rustls = { workspace = true, optional = true }
 hyper-util = { workspace = true }
 ic-agent = { workspace = true }
 ic-custom-domains-canister-api = { workspace = true }
+ic-transport-types = { workspace = true }
 indoc = { workspace = true }
 instant-acme = { workspace = true, optional = true }
+ipnet = { workspace = true }
 itertools = { workspace = true }
 mail-auth = { workspace = true, optional = true }
 mail-parser = { workspace = true, optional = true }
+maxminddb = { workspace = true }
 moka = { workspace = true }
 nix = { workspace = true }
 parse-size = { workspace = true }
@@ -112,7 +119,6 @@ tokio-util = { workspace = true }
 tokio-rustls = { workspace = true }
 tokio-io-timeout = { workspace = true }
 tower = { workspace = true }
-tower_governor = { workspace = true }
 tower-service = { workspace = true }
 tower-http = { workspace = true }
 tracing = { workspace = true }

ic-bn-lib/src/custom_domains/backend/router.rs

@@ -120,7 +120,7 @@ pub fn create_router(
 
 #[cfg(test)]
 mod tests {
-    use std::{str::FromStr, sync::Arc};
+    use std::{net::IpAddr, str::FromStr, sync::Arc};
 
     use axum::{
         body::{Body, to_bytes},
@@ -132,18 +132,21 @@ mod tests {
     use serde_json::Value;
     use tower::{Service, util::ServiceExt};
 
-    use crate::custom_domains::{
-        backend::router::{RateLimitConfig, create_router},
-        base::{
-            traits::{
-                repository::{MockRepository, RepositoryError},
-                validation::{MockValidatesDomains, ValidationError},
-            },
-            types::{
-                domain::{DomainStatus, RegistrationStatus},
-                task::{InputTask, TaskKind},
+    use crate::{
+        custom_domains::{
+            backend::router::{RateLimitConfig, create_router},
+            base::{
+                traits::{
+                    repository::{MockRepository, RepositoryError},
+                    validation::{MockValidatesDomains, ValidationError},
+                },
+                types::{
+                    domain::{DomainStatus, RegistrationStatus},
+                    task::{InputTask, TaskKind},
+                },
             },
         },
+        http::middleware::RemoteAddr,
     };
 
     const BODY_LIMIT: usize = 5000;
@@ -1028,30 +1031,30 @@ mod tests {
 
         let router =
             create_test_router_with_rate_limiter(mock_repository, mock_validator, rate_limits);
-        let test_ip = "192.168.1.100";
+        let test_ip = IpAddr::from_str("192.168.1.100").unwrap();
 
         // Create a service that maintains state between calls (needed for rate limiting)
         let mut service = router.into_service();
 
         let request1 = Request::builder()
             .method("POST")
             .uri("/v1/example1.org")
-            .header("x-real-ip", test_ip)
+            .extension(RemoteAddr(test_ip))
             .body(Body::empty())
             .unwrap();
 
         let request2 = Request::builder()
             .method("POST")
             .uri("/v1/example2.org")
-            .header("x-real-ip", test_ip)
+            .extension(RemoteAddr(test_ip))
             .body(Body::empty())
             .unwrap();
 
         // Should be rate-limited, as it's 3nd request from the same IP
         let request3 = Request::builder()
             .method("POST")
             .uri("/v1/example3.org")
-            .header("x-real-ip", test_ip)
+            .extension(RemoteAddr(test_ip))
             .body(Body::empty())
             .unwrap();
 

ic-bn-lib/src/custom_domains/base/types/worker.rs

@@ -28,6 +28,7 @@ use tracing::{Span, debug, error, info, instrument, warn};
 use x509_parser::{parse_x509_certificate, prelude::GeneralName};
 
 use crate::{
+    DurationDisplay,
     custom_domains::base::{
         helpers::{format_error_chain, retry_async},
         traits::{repository::Repository, time::UtcTimestamp, validation::ValidatesDomains},
@@ -185,7 +186,7 @@ impl Worker {
                 info!("Certificate revocation starts now");
             } else {
                 info!(
-                    delay_secs = delay.as_secs(),
+                    delay = %delay.display(),
                     "Certificate revocation scheduled"
                 );
             }
@@ -304,7 +305,7 @@ impl Worker {
                 };
 
                 info!(
-                    duration = task_result.duration.as_secs(),
+                    duration = %task_result.duration.display(),
                     not_before = validity.as_ref().map(|x| &x.0),
                     not_after = validity.as_ref().map(|x| &x.1),
                     "Task execution succeeded"
@@ -313,7 +314,7 @@ impl Worker {
 
             TaskOutcome::Failure(err) => {
                 error!(
-                    duration = task_result.duration.as_secs(),
+                    duration = %task_result.duration.display(),
                     error = ?err,
                     "Task execution failed"
                 );
@@ -410,7 +411,7 @@ impl Worker {
             Err(err) => {
                 error!(
                     error = ?err,
-                    duration_secs = self.config.task_fetch_retry_interval.as_secs(),
+                    duration = %self.config.task_fetch_retry_interval.display(),
                     "Failed to fetch pending task, sleeping before retry"
                 );
                 self.shared_metrics
@@ -517,7 +518,7 @@ impl Worker {
                         let attempts = err.attempts.to_string();
 
                         error!(
-                            duration_secs = self.config.task_submit_timeout.as_secs(),
+                            duration = %self.config.task_submit_timeout.display(),
                             "Failed to submit task result after {attempts} attempts: {err:?}",
                         );
 
@@ -544,7 +545,7 @@ impl Worker {
     /// Handles no available tasks, returning whether the worker should continue running
     async fn handle_no_tasks(&self) -> Result<(), WorkerStopped> {
         debug!(
-            duration_secs = self.config.polling_interval_no_tasks.as_secs(),
+            duration = %self.config.polling_interval_no_tasks.display(),
             "No pending tasks found, sleeping"
         );
 

ic-bn-lib/src/custom_domains/client.rs

@@ -32,6 +32,7 @@ use tokio_util::sync::CancellationToken;
 use tracing::{info, instrument, warn};
 
 use crate::{
+    DurationDisplay,
     custom_domains::{
         CustomDomain, ProvidesCustomDomains,
         base::{
@@ -49,8 +50,7 @@ use crate::{
     },
     ic_agent::Agent,
     tasks::Run,
-    tls::Pem,
-    tls::ProvidesCertificates,
+    tls::{Pem, ProvidesCertificates},
 };
 
 #[derive(new)]
@@ -426,9 +426,9 @@ impl Run for CanisterClient {
         interval_refresh.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
 
         warn!(
-            "Started polling every {}s, full refresh every {}s",
-            self.poll_interval.as_secs_f64(),
-            self.refresh_interval.as_secs_f64()
+            "Started polling every {}, full refresh every {}",
+            self.poll_interval.display(),
+            self.refresh_interval.display()
         );
 
         loop {

ic-bn-lib/src/http/cache.rs

@@ -1391,7 +1391,7 @@ mod tests {
         assert_eq!(cache.len(), 4);
     }
 
-    #[tokio::test]
+    #[tokio::test(start_paused = true)]
     async fn test_proxy_cache_lock() {
         let cache = Arc::new(
             CacheBuilder::new(KeyExtractorTest)

ic-bn-lib/src/http/middleware/mod.rs

@@ -1,65 +1,5 @@
 pub mod rate_limiter;
+pub mod request_meta;
 pub mod waf;
 
-use std::{net::IpAddr, str::FromStr, sync::Arc};
-
-use http::Request;
-
-use crate::http::{headers::X_REAL_IP, server::conn::ConnInfo};
-
-/// Extracts IP address from `x-real-ip` header or `ConnInfo` extension
-pub fn extract_ip_from_request<B>(req: &Request<B>) -> Option<IpAddr> {
-    // Try to extract from the header first
-    req.headers()
-        .get(X_REAL_IP)
-        .and_then(|x| x.to_str().ok())
-        .and_then(|x| IpAddr::from_str(x).ok())
-        .or_else(|| {
-            // Then, if that failed, from the ConnInfo extension
-            req.extensions()
-                .get::<Arc<ConnInfo>>()
-                .map(|x| x.remote_addr.ip())
-        })
-}
-
-#[cfg(test)]
-mod test {
-    use std::net::SocketAddr;
-
-    use crate::network::Addr;
-
-    use super::*;
-
-    #[test]
-    fn test_extract_ip_from_request() {
-        let addr1 = IpAddr::from_str("10.0.0.1").unwrap();
-        let addr2 = IpAddr::from_str("192.168.0.1").unwrap();
-
-        let mut ci = ConnInfo::default();
-        ci.remote_addr = Addr::Tcp(SocketAddr::new(addr1, 31337));
-        let ci = Arc::new(ci);
-
-        // Header takes precedence
-        let req = Request::builder()
-            .extension(ci.clone())
-            .header(X_REAL_IP, addr2.to_string())
-            .body("")
-            .unwrap();
-        assert_eq!(extract_ip_from_request(&req), Some(addr2));
-
-        // Only ConnInfo
-        let req = Request::builder().extension(ci).body("").unwrap();
-        assert_eq!(extract_ip_from_request(&req), Some(addr1));
-
-        // Only header
-        let req = Request::builder()
-            .header(X_REAL_IP, addr2.to_string())
-            .body("")
-            .unwrap();
-        assert_eq!(extract_ip_from_request(&req), Some(addr2));
-
-        // Neither
-        let req = Request::builder().body("").unwrap();
-        assert_eq!(extract_ip_from_request(&req), None);
-    }
-}
+pub use request_meta::{RemoteAddr, RequestId};