From f62e604117d1496b236acee7069a1780b198f789 Mon Sep 17 00:00:00 2001 From: Igor Novgorodov Date: Thu, 27 Aug 2026 12:52:10 +0000 Subject: [PATCH 1/7] Add support for bypass_token & custom canister_id --- Cargo.lock | 38 ++++++++++++- Cargo.toml | 2 +- .../custom_domains/backend/backend_service.rs | 18 +++++-- .../src/custom_domains/backend/handlers.rs | 53 +++++++++++++++++-- .../src/custom_domains/backend/router.rs | 2 +- .../custom_domains/base/traits/validation.rs | 5 ++ .../custom_domains/base/types/validator.rs | 9 ++++ .../src/custom_domains/tests/e2e_test.rs | 4 ++ ic-bn-lib/src/http/middleware/rate_limiter.rs | 18 +++++-- 9 files changed, 135 insertions(+), 14 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 66ea538..a89c7f2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -470,6 +470,7 @@ dependencies = [ "bytes", "futures-core", "futures-util", + "headers", "http", "http-body", "http-body-util", @@ -2251,6 +2252,30 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "headers" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b3314d5adb5d94bcdf56771f2e50dbbc80bb4bdf88967526706205ac9eff24eb" +dependencies = [ + "base64 0.22.1", + "bytes", + "headers-core", + "http", + "httpdate", + "mime", + "sha1 0.10.7", +] + +[[package]] +name = "headers-core" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "54b4a22553d4242c49fddb9ba998a99962b5cc6f22cb5a3482bec22522403ce4" +dependencies = [ + "http", +] + [[package]] name = "heck" version = "0.5.0" @@ -2662,7 +2687,7 @@ dependencies = [ "serde_with", "serde_yaml_ng", "sev", - "sha1", + "sha1 0.11.0", "sha2 0.11.0", "show-option", "smtp-proto", @@ -5578,6 +5603,17 @@ dependencies = [ "x509-cert", ] +[[package]] +name = "sha1" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a978451301f4db1d02937a4ab3ccce137717b81826e79b7d49ffe3244a13c3b8" +dependencies = [ + "cfg-if", + "cpufeatures 0.2.17", + "digest 0.10.7", +] + [[package]] name = "sha1" version = "0.11.0" diff --git a/Cargo.toml b/Cargo.toml index 117061a..fe385bb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -21,7 +21,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-extra = { version = "0.12.6", features = ["middleware", "typed-header"] } axum-server = "0.8.0" base64 = "0.23.1" bytes = "1.10.0" diff --git a/ic-bn-lib/src/custom_domains/backend/backend_service.rs b/ic-bn-lib/src/custom_domains/backend/backend_service.rs index e42e773..e32dccf 100644 --- a/ic-bn-lib/src/custom_domains/backend/backend_service.rs +++ b/ic-bn-lib/src/custom_domains/backend/backend_service.rs @@ -19,9 +19,11 @@ use crate::custom_domains::base::{ #[derive(Clone, new)] pub struct BackendService { /// Repository for storing domain data (e.g. certificates) and tasks - pub repository: Arc, + repository: Arc, /// Domain validator for DNS and canister ownership checks - pub validator: Arc, + validator: Arc, + /// Token that allows to bypass certain validation steps & specify the canister ID directly + pub bypass_token: Option, } impl BackendService { @@ -34,9 +36,19 @@ impl BackendService { domain: &str, task: TaskKind, wildcard: bool, + canister_id: Option, ) -> Result { let fqdn = parse_domain(domain)?; - let canister_id = self.validator.validate(&fqdn).await?; + + // If the canister ID is provided - use limited validation, otherwise full one + // that derives the canister ID from the DNS TXT record + let canister_id = if let Some(canister_id) = canister_id { + self.validator.validate_limited(&fqdn).await?; + canister_id + } else { + self.validator.validate(&fqdn).await? + }; + let task = InputTask::new(task, fqdn, wildcard); match self.repository.try_add_task(task).await { diff --git a/ic-bn-lib/src/custom_domains/backend/handlers.rs b/ic-bn-lib/src/custom_domains/backend/handlers.rs index a02ba8f..409766b 100644 --- a/ic-bn-lib/src/custom_domains/backend/handlers.rs +++ b/ic-bn-lib/src/custom_domains/backend/handlers.rs @@ -2,6 +2,11 @@ use axum::{ extract::{Path, Query, State}, http::StatusCode, }; +use axum_extra::{ + TypedHeader, + headers::{Authorization, authorization::Bearer}, +}; +use candid::Principal; use serde::Deserialize; use tracing::warn; @@ -12,7 +17,7 @@ use super::{ ValidateResponse, error_response, success_response, }, }; -use crate::custom_domains::base::types::task::TaskKind; +use crate::{constant_time_eq, custom_domains::base::types::task::TaskKind}; /// Query parameters for the domain registration endpoint. #[derive(Debug, Default, Deserialize)] @@ -20,6 +25,9 @@ pub struct CreateQuery { /// When `true`, the issued certificate also covers `*.domain`. #[serde(default)] pub wildcard: bool, + /// Canister ID to associate the domain with. + /// Taken into account only if the bypass token is provided in the request. + pub canister_id: Option, } fn log_error(err: &ApiError, domain: &str, operation: &str) { @@ -40,7 +48,8 @@ fn log_error(err: &ApiError, domain: &str, operation: &str) { path = "/v1/{id}", params( ("id" = String, Path, description = "Domain name to register"), - ("wildcard" = Option, Query, description = "Also issue a *.domain wildcard SAN") + ("wildcard" = Option, Query, description = "Also issue a *.domain wildcard SAN"), + ("canister_id" = Option, Query, description = "Canister ID to associate the domain with.") ), responses( (status = 202, description = "Domain registration request accepted", body = super::models::ApiResponse), @@ -54,9 +63,25 @@ pub async fn create_handler( State(backend_service): State, Path(domain): Path, Query(query): Query, + authorization: Option>>, ) -> axum::response::Response { + // Consider the canister_id provided in the query only if the bypass token + // is provided and matches the one configured in the backend service (if any). + let canister_id = authorization + .as_ref() + .map(|x| x.token()) + .zip(backend_service.bypass_token.as_ref()) + .zip(query.canister_id) + .and_then(|((token, bypass_token), canister_id)| { + if constant_time_eq(token.as_bytes(), bypass_token.as_bytes()) { + return Some(canister_id); + } + + None + }); + match backend_service - .submit_task(&domain, TaskKind::Issue, query.wildcard) + .submit_task(&domain, TaskKind::Issue, query.wildcard, canister_id) .await { Ok(canister_id) => success_response( @@ -90,7 +115,8 @@ pub async fn create_handler( patch, path = "/v1/{id}", params( - ("id" = String, Path, description = "Domain name to update") + ("id" = String, Path, description = "Domain name to update"), + ("canister_id" = Option, Query, description = "Canister ID to associate the domain with.") ), responses( (status = 202, description = "Update request accepted", body = super::models::ApiResponse), @@ -104,9 +130,26 @@ pub async fn create_handler( pub async fn update_handler( State(backend_service): State, Path(domain): Path, + Query(query): Query, + authorization: Option>>, ) -> axum::response::Response { + // Consider the canister_id provided in the query only if the bypass token + // is provided and matches the one configured in the backend service (if any). + let canister_id = authorization + .as_ref() + .map(|x| x.token()) + .zip(backend_service.bypass_token.as_ref()) + .zip(query.canister_id) + .and_then(|((token, bypass_token), canister_id)| { + if constant_time_eq(token.as_bytes(), bypass_token.as_bytes()) { + return Some(canister_id); + } + + None + }); + match backend_service - .submit_task(&domain, TaskKind::Update, false) + .submit_task(&domain, TaskKind::Update, false, canister_id) .await { Ok(canister_id) => success_response( diff --git a/ic-bn-lib/src/custom_domains/backend/router.rs b/ic-bn-lib/src/custom_domains/backend/router.rs index e3b3a53..2a0b307 100644 --- a/ic-bn-lib/src/custom_domains/backend/router.rs +++ b/ic-bn-lib/src/custom_domains/backend/router.rs @@ -49,7 +49,7 @@ pub fn create_router( with_metrics_endpoint: bool, bypass_token: Option, ) -> Router { - let backend_service = BackendService::new(repository, validator); + let backend_service = BackendService::new(repository, validator, bypass_token.clone()); let response = (StatusCode::TOO_MANY_REQUESTS, "Too many requests"); // Use ic-bn-lib rate limiting middleware, with key by IP address. diff --git a/ic-bn-lib/src/custom_domains/base/traits/validation.rs b/ic-bn-lib/src/custom_domains/base/traits/validation.rs index 66e166f..5c0ac11 100644 --- a/ic-bn-lib/src/custom_domains/base/traits/validation.rs +++ b/ic-bn-lib/src/custom_domains/base/traits/validation.rs @@ -45,6 +45,11 @@ pub trait ValidatesDomains: Send + Sync { /// canister ownership verification, and ACME challenge setup. async fn validate(&self, domain: &FQDN) -> Result; + /// Validates that a domain can be registered or updated. + /// + /// Skips certain checks compared to validate() + async fn validate_limited(&self, domain: &FQDN) -> Result<(), ValidationError>; + /// Validates that a domain can be safely deleted. /// /// Ensures DNS records are properly cleaned up before certificate revocation. diff --git a/ic-bn-lib/src/custom_domains/base/types/validator.rs b/ic-bn-lib/src/custom_domains/base/types/validator.rs index 91ebf1f..0a5e987 100644 --- a/ic-bn-lib/src/custom_domains/base/types/validator.rs +++ b/ic-bn-lib/src/custom_domains/base/types/validator.rs @@ -52,6 +52,15 @@ impl ValidatesDomains for Validator { Ok(canister_id) } + #[instrument(level = "info", skip_all, fields(domain = %domain))] + async fn validate_limited(&self, domain: &FQDN) -> Result<(), ValidationError> { + info!("Beginning validation"); + self.validate_cname_delegation(domain).await?; + self.validate_no_txt_challenge(domain).await?; + info!("Validation succeeded"); + Ok(()) + } + #[instrument(level = "info", skip_all, fields(domain = %domain))] async fn validate_deletion(&self, domain: &FQDN) -> Result<(), ValidationError> { info!("Beginning deletion validation"); diff --git a/ic-bn-lib/src/custom_domains/tests/e2e_test.rs b/ic-bn-lib/src/custom_domains/tests/e2e_test.rs index f17e659..838e9d1 100644 --- a/ic-bn-lib/src/custom_domains/tests/e2e_test.rs +++ b/ic-bn-lib/src/custom_domains/tests/e2e_test.rs @@ -131,6 +131,10 @@ impl ValidatesDomains for MockValidator { Ok("laqa6-raaaa-aaaam-aehzq-cai".parse().unwrap()) } + async fn validate_limited(&self, _domain: &FQDN) -> Result<(), ValidationError> { + Ok(()) + } + async fn validate_deletion(&self, _domain: &FQDN) -> Result<(), ValidationError> { Ok(()) } diff --git a/ic-bn-lib/src/http/middleware/rate_limiter.rs b/ic-bn-lib/src/http/middleware/rate_limiter.rs index b3fabd8..2a23e73 100644 --- a/ic-bn-lib/src/http/middleware/rate_limiter.rs +++ b/ic-bn-lib/src/http/middleware/rate_limiter.rs @@ -20,7 +20,10 @@ use governor::{ nanos::Nanos, state::keyed::DashMapStateStore, }; -use http::{HeaderName, HeaderValue, StatusCode, header::RETRY_AFTER}; +use http::{ + HeaderName, HeaderValue, StatusCode, + header::{AUTHORIZATION, RETRY_AFTER}, +}; use tower::{Layer, Service}; use crate::{constant_time_eq, hname, http::middleware::RemoteAddr}; @@ -108,12 +111,21 @@ where /// Stale entries cleanup interval - 5 minutes const CLEANUP_INTERVAL: Nanos = Nanos::new(300_000_000_000); - // Check that bypass token is configured, header was sent and it matches + // Check that bypass token is configured, header was sent and it matches. + // Checks both the custom header and the Authorization. let bypass = request .headers() .get(BYPASS_TOKEN_HEADER) + .map(|x| x.as_bytes()) + .or_else(|| { + request + .headers() + .get(AUTHORIZATION) + .and_then(|v| v.to_str().ok()) + .map(|s| s.trim_start_matches("Bearer ").as_bytes()) + }) .zip(self.state.bypass_token.as_ref()) - .is_some_and(|(hdr, token)| constant_time_eq(hdr.as_bytes(), token.as_bytes())); + .is_some_and(|(hdr, token)| constant_time_eq(hdr, token.as_bytes())); // Clean up stale entries from time to time let now = self.state.limiter.clock().now(); From 77e108842196506902b475a494a65531b5455ba5 Mon Sep 17 00:00:00 2001 From: Igor Novgorodov Date: Thu, 27 Aug 2026 13:22:39 +0000 Subject: [PATCH 2/7] Add tests --- .../src/custom_domains/backend/router.rs | 595 +++++++++++++++++- 1 file changed, 585 insertions(+), 10 deletions(-) diff --git a/ic-bn-lib/src/custom_domains/backend/router.rs b/ic-bn-lib/src/custom_domains/backend/router.rs index 2a0b307..17c6b68 100644 --- a/ic-bn-lib/src/custom_domains/backend/router.rs +++ b/ic-bn-lib/src/custom_domains/backend/router.rs @@ -126,7 +126,6 @@ mod tests { body::{Body, to_bytes}, http::{Request, StatusCode}, }; - use candid::Principal; use fqdn::FQDN; use prometheus::Registry; use serde_json::Value; @@ -147,6 +146,7 @@ mod tests { }, }, http::middleware::RemoteAddr, + principal, }; const BODY_LIMIT: usize = 5000; @@ -202,7 +202,7 @@ mod tests { async fn test_post_domain_success_accepted() { // Arrange let mut mock_validator = MockValidatesDomains::new(); - let expected_canister_id = Principal::from_text("rrkah-fqaaa-aaaaa-aaaaq-cai").unwrap(); + let expected_canister_id = principal!("rrkah-fqaaa-aaaaa-aaaaq-cai"); mock_validator .expect_validate() .returning(move |_| Box::pin(async move { Ok(expected_canister_id) })); @@ -264,7 +264,7 @@ mod tests { // Arrange let mut mock_validator = MockValidatesDomains::new(); mock_validator.expect_validate().returning(|_| { - Box::pin(async { Ok(Principal::from_text("rrkah-fqaaa-aaaaa-aaaaq-cai").unwrap()) }) + Box::pin(async { Ok(principal!("rrkah-fqaaa-aaaaa-aaaaq-cai")) }) }); let domain = "example.org"; @@ -300,7 +300,7 @@ mod tests { // Arrange let mut mock_validator = MockValidatesDomains::new(); mock_validator.expect_validate().returning(|_| { - Box::pin(async { Ok(Principal::from_text("rrkah-fqaaa-aaaaa-aaaaq-cai").unwrap()) }) + Box::pin(async { Ok(principal!("rrkah-fqaaa-aaaaa-aaaaq-cai")) }) }); let domain = "example.org"; @@ -383,7 +383,7 @@ mod tests { // Arrange let mut mock_validator = MockValidatesDomains::new(); mock_validator.expect_validate().returning(|_| { - Box::pin(async { Ok(Principal::from_text("rrkah-fqaaa-aaaaa-aaaaq-cai").unwrap()) }) + Box::pin(async { Ok(principal!("rrkah-fqaaa-aaaaa-aaaaq-cai")) }) }); let mut mock_repository = MockRepository::new(); @@ -460,7 +460,7 @@ mod tests { let mut mock_repository = MockRepository::new(); let expected_canister_id = - Principal::from_text("rrkah-fqaaa-aaaaa-aaaaq-cai").unwrap(); + principal!("rrkah-fqaaa-aaaaa-aaaaq-cai"); let domain_status = DomainStatus { domain: FQDN::from_str(domain).unwrap(), canister_id: Some(expected_canister_id), @@ -727,7 +727,7 @@ mod tests { for domain in &domains { // Arrange let mut mock_validator = MockValidatesDomains::new(); - let expected_canister_id = Principal::from_text("rrkah-fqaaa-aaaaa-aaaaq-cai").unwrap(); + let expected_canister_id = principal!("rrkah-fqaaa-aaaaa-aaaaq-cai"); mock_validator .expect_validate() @@ -808,7 +808,7 @@ mod tests { async fn test_post_domain_update_success_accepted() { // Arrange let mut mock_validator = MockValidatesDomains::new(); - let expected_canister_id = Principal::from_text("rrkah-fqaaa-aaaaa-aaaaq-cai").unwrap(); + let expected_canister_id = principal!("rrkah-fqaaa-aaaaa-aaaaq-cai"); mock_validator .expect_validate() .returning(move |_| Box::pin(async move { Ok(expected_canister_id) })); @@ -870,7 +870,7 @@ mod tests { // Arrange let mut mock_validator = MockValidatesDomains::new(); mock_validator.expect_validate().returning(|_| { - Box::pin(async { Ok(Principal::from_text("rrkah-fqaaa-aaaaa-aaaaq-cai").unwrap()) }) + Box::pin(async { Ok(principal!("rrkah-fqaaa-aaaaa-aaaaq-cai")) }) }); let mut mock_repository = MockRepository::new(); @@ -1011,7 +1011,7 @@ mod tests { // Arrange let mut mock_validator = MockValidatesDomains::new(); - let expected_canister_id = Principal::from_text("rrkah-fqaaa-aaaaa-aaaaq-cai").unwrap(); + let expected_canister_id = principal!("rrkah-fqaaa-aaaaa-aaaaq-cai"); mock_validator .expect_validate() .returning(move |_| Box::pin(async move { Ok(expected_canister_id) })) @@ -1073,4 +1073,579 @@ mod tests { let body_bytes = to_bytes(response3.into_body(), BODY_LIMIT).await.unwrap(); assert_eq!(body_bytes.to_vec(), b"Too many requests"); } + + // --- Bypass token & custom canister_id tests --- + + const BYPASS_TOKEN: &str = "s3cr3t-bypass-token"; + const OVERRIDE_CANISTER_ID: &str = "aaaaa-aa"; + + /// Helper function to create a router with a bypass token configured on the backend service. + fn create_test_router_with_bypass_token( + mock_repository: MockRepository, + mock_validator: MockValidatesDomains, + bypass_token: Option, + ) -> axum::Router { + let registry = Registry::new_custom(Some("custom_domains".into()), None).unwrap(); + create_router( + Arc::new(mock_repository), + Arc::new(mock_validator), + registry, + RateLimitConfig::default(), + true, + bypass_token, + ) + } + + /// Helper function to make a request to /v1/{domain} with an optional query string and an + /// optional `Authorization` header, parsing the response body as JSON. + async fn domain_request_with_query_and_auth( + router: axum::Router, + method: &str, + domain: &str, + query: Option<&str>, + auth_header: Option<&str>, + ) -> (StatusCode, Value) { + let uri = query.map_or_else( + || format!("/v1/{domain}"), + |query| format!("/v1/{domain}?{query}"), + ); + let mut builder = Request::builder().method(method).uri(uri); + if let Some(auth) = auth_header { + builder = builder.header("authorization", auth); + } + let request = builder.body(Body::empty()).unwrap(); + + let response = router.oneshot(request).await.unwrap(); + let status = response.status(); + let body_bytes = to_bytes(response.into_body(), BODY_LIMIT).await.unwrap(); + let body_json: Value = serde_json::from_slice(&body_bytes).unwrap(); + (status, body_json) + } + + #[tokio::test] + async fn test_post_domain_bypass_token_and_canister_id_uses_limited_validation() { + // Arrange + let mut mock_validator = MockValidatesDomains::new(); + mock_validator + .expect_validate_limited() + .times(1) + .returning(|_| Box::pin(async { Ok(()) })); + // `validate()` is intentionally left unmocked: it must not be called on the bypass path. + + let domain = "example.org"; + let expected_task = + InputTask::new(TaskKind::Issue, FQDN::from_str(domain).unwrap(), false); + let mut mock_repository = MockRepository::new(); + mock_repository + .expect_try_add_task() + .withf(move |task| *task == expected_task) + .times(1) + .returning(|_| Box::pin(async { Ok(()) })); + + let router = create_test_router_with_bypass_token( + mock_repository, + mock_validator, + Some(BYPASS_TOKEN.to_string()), + ); + + // Act + let (status, response_json) = domain_request_with_query_and_auth( + router, + "POST", + domain, + Some(&format!("canister_id={OVERRIDE_CANISTER_ID}")), + Some(&format!("Bearer {BYPASS_TOKEN}")), + ) + .await; + + // Assert + assert_eq!(status, StatusCode::ACCEPTED); + assert_eq!(response_json["status"], "success"); + assert_eq!(response_json["data"]["domain"], domain); + assert_eq!(response_json["data"]["canister_id"], OVERRIDE_CANISTER_ID); + } + + #[tokio::test] + async fn test_post_domain_wrong_bypass_token_falls_back_to_full_validation() { + // Arrange + let mut mock_validator = MockValidatesDomains::new(); + let derived_canister_id = principal!("rrkah-fqaaa-aaaaa-aaaaq-cai"); + mock_validator + .expect_validate() + .times(1) + .returning(move |_| Box::pin(async move { Ok(derived_canister_id) })); + // `validate_limited()` is intentionally left unmocked: it must not be called when the + // provided token doesn't match the configured bypass token. + + let domain = "example.org"; + let expected_task = + InputTask::new(TaskKind::Issue, FQDN::from_str(domain).unwrap(), false); + let mut mock_repository = MockRepository::new(); + mock_repository + .expect_try_add_task() + .withf(move |task| *task == expected_task) + .times(1) + .returning(|_| Box::pin(async { Ok(()) })); + + let router = create_test_router_with_bypass_token( + mock_repository, + mock_validator, + Some(BYPASS_TOKEN.to_string()), + ); + + // Act + let (status, response_json) = domain_request_with_query_and_auth( + router, + "POST", + domain, + Some(&format!("canister_id={OVERRIDE_CANISTER_ID}")), + Some("Bearer this-is-not-the-right-token"), + ) + .await; + + // Assert + assert_eq!(status, StatusCode::ACCEPTED); + assert_eq!( + response_json["data"]["canister_id"], + "rrkah-fqaaa-aaaaa-aaaaq-cai" + ); + } + + #[tokio::test] + async fn test_post_domain_canister_id_without_auth_header_uses_full_validation() { + // Arrange + let mut mock_validator = MockValidatesDomains::new(); + let derived_canister_id = principal!("rrkah-fqaaa-aaaaa-aaaaq-cai"); + mock_validator + .expect_validate() + .times(1) + .returning(move |_| Box::pin(async move { Ok(derived_canister_id) })); + + let domain = "example.org"; + let mut mock_repository = MockRepository::new(); + mock_repository + .expect_try_add_task() + .times(1) + .returning(|_| Box::pin(async { Ok(()) })); + + let router = create_test_router_with_bypass_token( + mock_repository, + mock_validator, + Some(BYPASS_TOKEN.to_string()), + ); + + // Act: canister_id is provided but there's no Authorization header at all + let (status, response_json) = domain_request_with_query_and_auth( + router, + "POST", + domain, + Some(&format!("canister_id={OVERRIDE_CANISTER_ID}")), + None, + ) + .await; + + // Assert + assert_eq!(status, StatusCode::ACCEPTED); + assert_eq!( + response_json["data"]["canister_id"], + "rrkah-fqaaa-aaaaa-aaaaq-cai" + ); + } + + #[tokio::test] + async fn test_post_domain_valid_bypass_token_without_canister_id_uses_full_validation() { + // Arrange + let mut mock_validator = MockValidatesDomains::new(); + let derived_canister_id = principal!("rrkah-fqaaa-aaaaa-aaaaq-cai"); + mock_validator + .expect_validate() + .times(1) + .returning(move |_| Box::pin(async move { Ok(derived_canister_id) })); + + let domain = "example.org"; + let mut mock_repository = MockRepository::new(); + mock_repository + .expect_try_add_task() + .times(1) + .returning(|_| Box::pin(async { Ok(()) })); + + let router = create_test_router_with_bypass_token( + mock_repository, + mock_validator, + Some(BYPASS_TOKEN.to_string()), + ); + + // Act: correct bypass token, but no canister_id query param + let (status, response_json) = domain_request_with_query_and_auth( + router, + "POST", + domain, + None, + Some(&format!("Bearer {BYPASS_TOKEN}")), + ) + .await; + + // Assert + assert_eq!(status, StatusCode::ACCEPTED); + assert_eq!( + response_json["data"]["canister_id"], + "rrkah-fqaaa-aaaaa-aaaaq-cai" + ); + } + + #[tokio::test] + async fn test_post_domain_no_bypass_token_configured_ignores_header_and_canister_id() { + // Arrange: backend service has no bypass token configured at all + let mut mock_validator = MockValidatesDomains::new(); + let derived_canister_id = principal!("rrkah-fqaaa-aaaaa-aaaaq-cai"); + mock_validator + .expect_validate() + .times(1) + .returning(move |_| Box::pin(async move { Ok(derived_canister_id) })); + + let domain = "example.org"; + let mut mock_repository = MockRepository::new(); + mock_repository + .expect_try_add_task() + .times(1) + .returning(|_| Box::pin(async { Ok(()) })); + + let router = + create_test_router_with_bypass_token(mock_repository, mock_validator, None); + + // Act + let (status, response_json) = domain_request_with_query_and_auth( + router, + "POST", + domain, + Some(&format!("canister_id={OVERRIDE_CANISTER_ID}")), + Some(&format!("Bearer {BYPASS_TOKEN}")), + ) + .await; + + // Assert + assert_eq!(status, StatusCode::ACCEPTED); + assert_eq!( + response_json["data"]["canister_id"], + "rrkah-fqaaa-aaaaa-aaaaq-cai" + ); + } + + #[tokio::test] + async fn test_post_domain_bypass_token_validate_limited_error() { + // Arrange + let mut mock_validator = MockValidatesDomains::new(); + mock_validator + .expect_validate_limited() + .times(1) + .returning(|_| { + Box::pin(async { + Err(ValidationError::MissingDnsCname { + src: "_acme-challenge.example.org.".to_string(), + dst: "_acme-challenge.example.org.icp2.io.".to_string(), + }) + }) + }); + + let mock_repository = MockRepository::new(); + let router = create_test_router_with_bypass_token( + mock_repository, + mock_validator, + Some(BYPASS_TOKEN.to_string()), + ); + + // Act + let (status, response_json) = domain_request_with_query_and_auth( + router, + "POST", + "example.org", + Some(&format!("canister_id={OVERRIDE_CANISTER_ID}")), + Some(&format!("Bearer {BYPASS_TOKEN}")), + ) + .await; + + // Assert + assert_eq!(status, StatusCode::BAD_REQUEST); + assert_eq!(response_json["status"], "error"); + assert_eq!( + response_json["errors"].as_str().unwrap(), + "bad_request: missing DNS CNAME record from _acme-challenge.example.org. to _acme-challenge.example.org.icp2.io." + ); + } + + #[tokio::test] + async fn test_post_domain_malformed_authorization_header_is_rejected() { + // Arrange: a non-Bearer Authorization header should be rejected by the extractor + // itself, before the bypass logic ever runs. + let mock_validator = MockValidatesDomains::new(); + let mock_repository = MockRepository::new(); + let router = create_test_router_with_bypass_token( + mock_repository, + mock_validator, + Some(BYPASS_TOKEN.to_string()), + ); + + // Act + let request = Request::builder() + .method("POST") + .uri(format!("/v1/example.org?canister_id={OVERRIDE_CANISTER_ID}")) + .header("authorization", "Basic dXNlcjpwYXNz") + .body(Body::empty()) + .unwrap(); + let response = router.oneshot(request).await.unwrap(); + + // Assert + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + } + + #[tokio::test] + async fn test_patch_domain_bypass_token_and_canister_id_uses_limited_validation() { + // Arrange + let mut mock_validator = MockValidatesDomains::new(); + mock_validator + .expect_validate_limited() + .times(1) + .returning(|_| Box::pin(async { Ok(()) })); + // `validate()` is intentionally left unmocked: it must not be called on the bypass path. + + let domain = "example.org"; + let expected_task = + InputTask::new(TaskKind::Update, FQDN::from_str(domain).unwrap(), false); + let mut mock_repository = MockRepository::new(); + mock_repository + .expect_try_add_task() + .withf(move |task| *task == expected_task) + .times(1) + .returning(|_| Box::pin(async { Ok(()) })); + + let router = create_test_router_with_bypass_token( + mock_repository, + mock_validator, + Some(BYPASS_TOKEN.to_string()), + ); + + // Act + let (status, response_json) = domain_request_with_query_and_auth( + router, + "PATCH", + domain, + Some(&format!("canister_id={OVERRIDE_CANISTER_ID}")), + Some(&format!("Bearer {BYPASS_TOKEN}")), + ) + .await; + + // Assert + assert_eq!(status, StatusCode::ACCEPTED); + assert_eq!(response_json["status"], "success"); + assert_eq!(response_json["data"]["domain"], domain); + assert_eq!(response_json["data"]["canister_id"], OVERRIDE_CANISTER_ID); + } + + #[tokio::test] + async fn test_patch_domain_wrong_bypass_token_falls_back_to_full_validation() { + // Arrange + let mut mock_validator = MockValidatesDomains::new(); + let derived_canister_id = principal!("rrkah-fqaaa-aaaaa-aaaaq-cai"); + mock_validator + .expect_validate() + .times(1) + .returning(move |_| Box::pin(async move { Ok(derived_canister_id) })); + // `validate_limited()` is intentionally left unmocked: it must not be called when the + // provided token doesn't match the configured bypass token. + + let domain = "example.org"; + let expected_task = + InputTask::new(TaskKind::Update, FQDN::from_str(domain).unwrap(), false); + let mut mock_repository = MockRepository::new(); + mock_repository + .expect_try_add_task() + .withf(move |task| *task == expected_task) + .times(1) + .returning(|_| Box::pin(async { Ok(()) })); + + let router = create_test_router_with_bypass_token( + mock_repository, + mock_validator, + Some(BYPASS_TOKEN.to_string()), + ); + + // Act + let (status, response_json) = domain_request_with_query_and_auth( + router, + "PATCH", + domain, + Some(&format!("canister_id={OVERRIDE_CANISTER_ID}")), + Some("Bearer this-is-not-the-right-token"), + ) + .await; + + // Assert + assert_eq!(status, StatusCode::ACCEPTED); + assert_eq!( + response_json["data"]["canister_id"], + "rrkah-fqaaa-aaaaa-aaaaq-cai" + ); + } + + #[tokio::test] + async fn test_patch_domain_canister_id_without_auth_header_uses_full_validation() { + // Arrange + let mut mock_validator = MockValidatesDomains::new(); + let derived_canister_id = principal!("rrkah-fqaaa-aaaaa-aaaaq-cai"); + mock_validator + .expect_validate() + .times(1) + .returning(move |_| Box::pin(async move { Ok(derived_canister_id) })); + + let domain = "example.org"; + let mut mock_repository = MockRepository::new(); + mock_repository + .expect_try_add_task() + .times(1) + .returning(|_| Box::pin(async { Ok(()) })); + + let router = create_test_router_with_bypass_token( + mock_repository, + mock_validator, + Some(BYPASS_TOKEN.to_string()), + ); + + // Act: canister_id is provided but there's no Authorization header at all + let (status, response_json) = domain_request_with_query_and_auth( + router, + "PATCH", + domain, + Some(&format!("canister_id={OVERRIDE_CANISTER_ID}")), + None, + ) + .await; + + // Assert + assert_eq!(status, StatusCode::ACCEPTED); + assert_eq!( + response_json["data"]["canister_id"], + "rrkah-fqaaa-aaaaa-aaaaq-cai" + ); + } + + #[tokio::test] + async fn test_patch_domain_valid_bypass_token_without_canister_id_uses_full_validation() { + // Arrange + let mut mock_validator = MockValidatesDomains::new(); + let derived_canister_id = principal!("rrkah-fqaaa-aaaaa-aaaaq-cai"); + mock_validator + .expect_validate() + .times(1) + .returning(move |_| Box::pin(async move { Ok(derived_canister_id) })); + + let domain = "example.org"; + let mut mock_repository = MockRepository::new(); + mock_repository + .expect_try_add_task() + .times(1) + .returning(|_| Box::pin(async { Ok(()) })); + + let router = create_test_router_with_bypass_token( + mock_repository, + mock_validator, + Some(BYPASS_TOKEN.to_string()), + ); + + // Act: correct bypass token, but no canister_id query param + let (status, response_json) = domain_request_with_query_and_auth( + router, + "PATCH", + domain, + None, + Some(&format!("Bearer {BYPASS_TOKEN}")), + ) + .await; + + // Assert + assert_eq!(status, StatusCode::ACCEPTED); + assert_eq!( + response_json["data"]["canister_id"], + "rrkah-fqaaa-aaaaa-aaaaq-cai" + ); + } + + #[tokio::test] + async fn test_patch_domain_no_bypass_token_configured_ignores_header_and_canister_id() { + // Arrange: backend service has no bypass token configured at all + let mut mock_validator = MockValidatesDomains::new(); + let derived_canister_id = principal!("rrkah-fqaaa-aaaaa-aaaaq-cai"); + mock_validator + .expect_validate() + .times(1) + .returning(move |_| Box::pin(async move { Ok(derived_canister_id) })); + + let domain = "example.org"; + let mut mock_repository = MockRepository::new(); + mock_repository + .expect_try_add_task() + .times(1) + .returning(|_| Box::pin(async { Ok(()) })); + + let router = + create_test_router_with_bypass_token(mock_repository, mock_validator, None); + + // Act + let (status, response_json) = domain_request_with_query_and_auth( + router, + "PATCH", + domain, + Some(&format!("canister_id={OVERRIDE_CANISTER_ID}")), + Some(&format!("Bearer {BYPASS_TOKEN}")), + ) + .await; + + // Assert + assert_eq!(status, StatusCode::ACCEPTED); + assert_eq!( + response_json["data"]["canister_id"], + "rrkah-fqaaa-aaaaa-aaaaq-cai" + ); + } + + #[tokio::test] + async fn test_patch_domain_bypass_token_validate_limited_error() { + // Arrange + let mut mock_validator = MockValidatesDomains::new(); + mock_validator + .expect_validate_limited() + .times(1) + .returning(|_| { + Box::pin(async { + Err(ValidationError::MissingDnsCname { + src: "_acme-challenge.example.org.".to_string(), + dst: "_acme-challenge.example.org.icp2.io.".to_string(), + }) + }) + }); + + let mock_repository = MockRepository::new(); + let router = create_test_router_with_bypass_token( + mock_repository, + mock_validator, + Some(BYPASS_TOKEN.to_string()), + ); + + // Act + let (status, response_json) = domain_request_with_query_and_auth( + router, + "PATCH", + "example.org", + Some(&format!("canister_id={OVERRIDE_CANISTER_ID}")), + Some(&format!("Bearer {BYPASS_TOKEN}")), + ) + .await; + + // Assert + assert_eq!(status, StatusCode::BAD_REQUEST); + assert_eq!(response_json["status"], "error"); + assert_eq!( + response_json["errors"].as_str().unwrap(), + "bad_request: missing DNS CNAME record from _acme-challenge.example.org. to _acme-challenge.example.org.icp2.io." + ); + } } From d66ac452639fe3c77ae1c1fd95cd80efc0edd1aa Mon Sep 17 00:00:00 2001 From: Igor Novgorodov Date: Thu, 27 Aug 2026 14:14:17 +0000 Subject: [PATCH 3/7] Add canister_id to the tasks --- custom-domains/api/src/lib.rs | 6 +- custom-domains/canister/canister_backend.did | 4 + custom-domains/canister/src/state.rs | 23 ++++++ .../custom_domains/backend/backend_service.rs | 16 ++-- .../src/custom_domains/backend/router.rs | 75 ++++++++++++------- .../src/custom_domains/base/types/task.rs | 5 ++ .../src/custom_domains/base/types/worker.rs | 15 +++- ic-bn-lib/src/custom_domains/client.rs | 1 + ic-bn-lib/src/custom_domains/tests/helpers.rs | 1 + 9 files changed, 105 insertions(+), 41 deletions(-) diff --git a/custom-domains/api/src/lib.rs b/custom-domains/api/src/lib.rs index 11a32e9..ea25e67 100644 --- a/custom-domains/api/src/lib.rs +++ b/custom-domains/api/src/lib.rs @@ -79,8 +79,10 @@ pub enum TaskKind { pub struct InputTask { pub kind: TaskKind, pub domain: String, - // Whether to also include a `*.domain` wildcard SAN in the certificate + /// Whether to also include a `*.domain` wildcard SAN in the certificate pub wildcard: Option, + /// The canister ID associated with the domain (if known at submission time) + pub canister_id: Option, } #[derive(CandidType, Deserialize, Serialize, Debug, Clone, PartialEq, Eq, new)] @@ -91,6 +93,8 @@ pub struct ScheduledTask { pub enc_cert: Option>, // Whether to also include a `*.domain` wildcard SAN in the certificate pub wildcard: Option, + /// The canister ID associated with the domain (if known at submission time) + pub canister_id: Option, } #[derive(CandidType, Deserialize, Serialize, Clone, Debug)] diff --git a/custom-domains/canister/canister_backend.did b/custom-domains/canister/canister_backend.did index 938dcba..dcab9d6 100644 --- a/custom-domains/canister/canister_backend.did +++ b/custom-domains/canister/canister_backend.did @@ -27,6 +27,8 @@ type InputTask = record { domain: text; // Whether to also include a *.domain wildcard SAN in the certificate wildcard: opt bool; + // The canister ID associated with the domain (if known at submission time) + canister_id: opt principal; }; // Reasons why a task might fail @@ -99,6 +101,8 @@ type ScheduledTask = record { enc_cert: opt blob; // Whether to also include a *.domain wildcard SAN in the certificate wildcard: opt bool; + // The canister ID associated with the domain (if known at submission time) + canister_id: opt principal; }; // Status of a domain diff --git a/custom-domains/canister/src/state.rs b/custom-domains/canister/src/state.rs index 0674bd6..864c09d 100644 --- a/custom-domains/canister/src/state.rs +++ b/custom-domains/canister/src/state.rs @@ -299,6 +299,7 @@ impl CanisterState { now, enc_cert, Some(domain_entry.wildcard), + None, ))) } None => Ok(None), @@ -1539,6 +1540,7 @@ mod tests { domain: domain_name.clone(), kind: TaskKind::Issue, wildcard: None, + canister_id: None, }; let result = state.try_add_task(task, now); assert!(result.is_ok()); @@ -1558,6 +1560,7 @@ mod tests { domain: "new.example.com".to_string(), kind: TaskKind::Update, wildcard: None, + canister_id: None, }; let result = state.try_add_task(task, now); assert!(matches!(result, Err(TryAddTaskError::DomainNotFound(_)))); @@ -1577,6 +1580,7 @@ mod tests { domain: domain.clone(), kind: TaskKind::Issue, wildcard: Some(true), + canister_id: None, }, now, ) @@ -1597,6 +1601,7 @@ mod tests { domain: domain.clone(), kind: TaskKind::Update, wildcard: Some(false), + canister_id: None, }, now, ) @@ -1623,6 +1628,7 @@ mod tests { domain: domain.clone(), kind: TaskKind::Update, wildcard: None, + canister_id: None, }; let result = state.try_add_task(task, now); assert!(matches!( @@ -1642,6 +1648,7 @@ mod tests { domain: domain.clone(), kind: task_kind, wildcard: None, + canister_id: None, }; let result = state.try_add_task(task, now); assert!(matches!( @@ -1675,6 +1682,7 @@ mod tests { domain: domain.clone(), kind: TaskKind::Update, wildcard: None, + canister_id: None, }; let result = state.try_add_task(task, now); assert!(result.is_ok()); @@ -1700,6 +1708,7 @@ mod tests { domain: "issue.com".to_string(), kind: TaskKind::Issue, wildcard: None, + canister_id: None, }; let result = state.try_add_task(task, now); assert!(result.is_ok()); @@ -1721,6 +1730,7 @@ mod tests { domain: "update.com".to_string(), kind: TaskKind::Update, wildcard: None, + canister_id: None, }; let result = state.try_add_task(task, now); assert!(result.is_ok()); @@ -1742,6 +1752,7 @@ mod tests { domain: "delete.com".to_string(), kind: TaskKind::Delete, wildcard: None, + canister_id: None, }; let result = state.try_add_task(task, now); assert!(result.is_ok()); @@ -1763,6 +1774,7 @@ mod tests { domain: "renew.com".to_string(), kind: TaskKind::Renew, wildcard: None, + canister_id: None, }; let result = state.try_add_task(task, now); assert!(result.is_ok()); @@ -1790,6 +1802,7 @@ mod tests { domain: "update.com".to_string(), kind: TaskKind::Update, wildcard: None, + canister_id: None, }; let result = state.try_add_task(task, now); // Should fail because certificate is required for Update @@ -1811,6 +1824,7 @@ mod tests { domain: "delete.com".to_string(), kind: TaskKind::Delete, wildcard: None, + canister_id: None, }; let result = state.try_add_task(task, now); // Should succeed as Delete can work with just canister_id @@ -1833,6 +1847,7 @@ mod tests { domain: "issue.com".to_string(), kind: TaskKind::Issue, wildcard: None, + canister_id: None, }; let result = state.try_add_task(task, now); // Should fail because Issue on already registered domain is not allowed @@ -1855,6 +1870,7 @@ mod tests { domain: "renew.com".to_string(), kind: TaskKind::Renew, wildcard: None, + canister_id: None, }; let result = state.try_add_task(task, now); // Should succeed because Renew on registered domain is allowed @@ -1893,11 +1909,13 @@ mod tests { domain: "task1.com".to_string(), kind: TaskKind::Issue, wildcard: None, + canister_id: None, }; let task_2 = InputTask { domain: "task2.com".to_string(), kind: TaskKind::Issue, wildcard: None, + canister_id: None, }; state .try_add_task(task_1, now) @@ -1914,6 +1932,7 @@ mod tests { now, None, Some(false), + None, )); assert_eq!(task, expected_task); let task = state.fetch_next_task(now).unwrap(); @@ -1923,6 +1942,7 @@ mod tests { now, None, Some(false), + None, )); assert_eq!(task, expected_task); let task = state.fetch_next_task(now).unwrap(); @@ -1958,6 +1978,7 @@ mod tests { now - 3000, Some(b"cert_data".to_vec()), Some(false), + None, )); assert_eq!(result, expected_task); } @@ -2467,11 +2488,13 @@ mod tests { domain: "example.com".to_string(), kind: TaskKind::Issue, wildcard: None, + canister_id: None, }; let task2 = InputTask { domain: "example1.com".to_string(), kind: TaskKind::Issue, wildcard: None, + canister_id: None, }; // First task should succeed diff --git a/ic-bn-lib/src/custom_domains/backend/backend_service.rs b/ic-bn-lib/src/custom_domains/backend/backend_service.rs index e32dccf..31ed1f8 100644 --- a/ic-bn-lib/src/custom_domains/backend/backend_service.rs +++ b/ic-bn-lib/src/custom_domains/backend/backend_service.rs @@ -49,24 +49,20 @@ impl BackendService { self.validator.validate(&fqdn).await? }; - let task = InputTask::new(task, fqdn, wildcard); + let task = InputTask::new(task, fqdn, wildcard, Some(canister_id)); - match self.repository.try_add_task(task).await { - Ok(()) => Ok(canister_id), - Err(err) => Err(err.into()), - } + self.repository.try_add_task(task).await?; + Ok(canister_id) } /// Validates domain can be deleted and submits a delete task pub async fn submit_delete_task(&self, domain: &str) -> Result<(), ApiError> { let fqdn = parse_domain(domain)?; self.validator.validate_deletion(&fqdn).await?; - let task = InputTask::new(TaskKind::Delete, fqdn, false); + let task = InputTask::new(TaskKind::Delete, fqdn, false, None); - match self.repository.try_add_task(task).await { - Ok(()) => Ok(()), - Err(err) => Err(err.into()), - } + self.repository.try_add_task(task).await?; + Ok(()) } /// Retrieves the current status of a domain registration diff --git a/ic-bn-lib/src/custom_domains/backend/router.rs b/ic-bn-lib/src/custom_domains/backend/router.rs index 17c6b68..d30dbf2 100644 --- a/ic-bn-lib/src/custom_domains/backend/router.rs +++ b/ic-bn-lib/src/custom_domains/backend/router.rs @@ -215,11 +215,13 @@ mod tests { TaskKind::Issue, FQDN::from_str(domain_normal).unwrap(), false, + Some(expected_canister_id), ); let expected_task_unicode = InputTask::new( TaskKind::Issue, FQDN::from_str(subdomain_unicode).unwrap(), false, + Some(expected_canister_id), ); mock_repository .expect_try_add_task() @@ -263,9 +265,9 @@ mod tests { async fn test_post_domain_conflict_certificate_already_issued() { // Arrange let mut mock_validator = MockValidatesDomains::new(); - mock_validator.expect_validate().returning(|_| { - Box::pin(async { Ok(principal!("rrkah-fqaaa-aaaaa-aaaaq-cai")) }) - }); + mock_validator + .expect_validate() + .returning(|_| Box::pin(async { Ok(principal!("rrkah-fqaaa-aaaaa-aaaaq-cai")) })); let domain = "example.org"; let mut mock_repository = MockRepository::new(); @@ -299,9 +301,9 @@ mod tests { async fn test_post_domain_conflict_another_task_in_progress() { // Arrange let mut mock_validator = MockValidatesDomains::new(); - mock_validator.expect_validate().returning(|_| { - Box::pin(async { Ok(principal!("rrkah-fqaaa-aaaaa-aaaaq-cai")) }) - }); + mock_validator + .expect_validate() + .returning(|_| Box::pin(async { Ok(principal!("rrkah-fqaaa-aaaaa-aaaaq-cai")) })); let domain = "example.org"; let mut mock_repository = MockRepository::new(); @@ -382,9 +384,9 @@ mod tests { async fn test_post_domain_internal_server_error_repository_failure() { // Arrange let mut mock_validator = MockValidatesDomains::new(); - mock_validator.expect_validate().returning(|_| { - Box::pin(async { Ok(principal!("rrkah-fqaaa-aaaaa-aaaaq-cai")) }) - }); + mock_validator + .expect_validate() + .returning(|_| Box::pin(async { Ok(principal!("rrkah-fqaaa-aaaaa-aaaaq-cai")) })); let mut mock_repository = MockRepository::new(); mock_repository.expect_try_add_task().returning(|_| { @@ -459,8 +461,7 @@ mod tests { let mock_validator = MockValidatesDomains::new(); let mut mock_repository = MockRepository::new(); - let expected_canister_id = - principal!("rrkah-fqaaa-aaaaa-aaaaq-cai"); + let expected_canister_id = principal!("rrkah-fqaaa-aaaaa-aaaaq-cai"); let domain_status = DomainStatus { domain: FQDN::from_str(domain).unwrap(), canister_id: Some(expected_canister_id), @@ -821,11 +822,13 @@ mod tests { TaskKind::Update, FQDN::from_str(domain_normal).unwrap(), false, + Some(expected_canister_id), ); let expected_task_unicode = InputTask::new( TaskKind::Update, FQDN::from_str(subdomain_unicode).unwrap(), false, + Some(expected_canister_id), ); mock_repository .expect_try_add_task() @@ -869,9 +872,9 @@ mod tests { async fn test_post_domain_update_bad_request_missing_certificate() { // Arrange let mut mock_validator = MockValidatesDomains::new(); - mock_validator.expect_validate().returning(|_| { - Box::pin(async { Ok(principal!("rrkah-fqaaa-aaaaa-aaaaq-cai")) }) - }); + mock_validator + .expect_validate() + .returning(|_| Box::pin(async { Ok(principal!("rrkah-fqaaa-aaaaa-aaaaq-cai")) })); let mut mock_repository = MockRepository::new(); mock_repository.expect_try_add_task().returning(|_| { @@ -931,11 +934,13 @@ mod tests { TaskKind::Delete, FQDN::from_str(domain_normal).unwrap(), false, + None, ); let expected_task_unicode = InputTask::new( TaskKind::Delete, FQDN::from_str(subdomain_unicode).unwrap(), false, + None, ); mock_repository .expect_try_add_task() @@ -1133,8 +1138,12 @@ mod tests { // `validate()` is intentionally left unmocked: it must not be called on the bypass path. let domain = "example.org"; - let expected_task = - InputTask::new(TaskKind::Issue, FQDN::from_str(domain).unwrap(), false); + let expected_task = InputTask::new( + TaskKind::Issue, + FQDN::from_str(domain).unwrap(), + false, + Some(principal!(OVERRIDE_CANISTER_ID)), + ); let mut mock_repository = MockRepository::new(); mock_repository .expect_try_add_task() @@ -1178,8 +1187,12 @@ mod tests { // provided token doesn't match the configured bypass token. let domain = "example.org"; - let expected_task = - InputTask::new(TaskKind::Issue, FQDN::from_str(domain).unwrap(), false); + let expected_task = InputTask::new( + TaskKind::Issue, + FQDN::from_str(domain).unwrap(), + false, + Some(derived_canister_id), + ); let mut mock_repository = MockRepository::new(); mock_repository .expect_try_add_task() @@ -1310,8 +1323,7 @@ mod tests { .times(1) .returning(|_| Box::pin(async { Ok(()) })); - let router = - create_test_router_with_bypass_token(mock_repository, mock_validator, None); + let router = create_test_router_with_bypass_token(mock_repository, mock_validator, None); // Act let (status, response_json) = domain_request_with_query_and_auth( @@ -1388,7 +1400,9 @@ mod tests { // Act let request = Request::builder() .method("POST") - .uri(format!("/v1/example.org?canister_id={OVERRIDE_CANISTER_ID}")) + .uri(format!( + "/v1/example.org?canister_id={OVERRIDE_CANISTER_ID}" + )) .header("authorization", "Basic dXNlcjpwYXNz") .body(Body::empty()) .unwrap(); @@ -1409,8 +1423,12 @@ mod tests { // `validate()` is intentionally left unmocked: it must not be called on the bypass path. let domain = "example.org"; - let expected_task = - InputTask::new(TaskKind::Update, FQDN::from_str(domain).unwrap(), false); + let expected_task = InputTask::new( + TaskKind::Update, + FQDN::from_str(domain).unwrap(), + false, + Some(principal!(OVERRIDE_CANISTER_ID)), + ); let mut mock_repository = MockRepository::new(); mock_repository .expect_try_add_task() @@ -1454,8 +1472,12 @@ mod tests { // provided token doesn't match the configured bypass token. let domain = "example.org"; - let expected_task = - InputTask::new(TaskKind::Update, FQDN::from_str(domain).unwrap(), false); + let expected_task = InputTask::new( + TaskKind::Update, + FQDN::from_str(domain).unwrap(), + false, + Some(derived_canister_id), + ); let mut mock_repository = MockRepository::new(); mock_repository .expect_try_add_task() @@ -1586,8 +1608,7 @@ mod tests { .times(1) .returning(|_| Box::pin(async { Ok(()) })); - let router = - create_test_router_with_bypass_token(mock_repository, mock_validator, None); + let router = create_test_router_with_bypass_token(mock_repository, mock_validator, None); // Act let (status, response_json) = domain_request_with_query_and_auth( diff --git a/ic-bn-lib/src/custom_domains/base/types/task.rs b/ic-bn-lib/src/custom_domains/base/types/task.rs index 979cc25..b4c8d0d 100644 --- a/ic-bn-lib/src/custom_domains/base/types/task.rs +++ b/ic-bn-lib/src/custom_domains/base/types/task.rs @@ -35,6 +35,8 @@ pub struct InputTask { pub domain: FQDN, /// Whether to also include a `*.domain` wildcard SAN in the certificate pub wildcard: bool, + /// The canister ID associated with the domain (if known at submission time) + pub canister_id: Option, } /// Scheduled task that is ready for execution by a worker. @@ -50,6 +52,8 @@ pub struct ScheduledTask { pub cert: Option>, /// Whether to also include a `*.domain` wildcard SAN in the certificate pub wildcard: bool, + /// The canister ID associated with the domain (if known at submission time) + pub canister_id: Option, } /// Represents the result of a task execution submitted by a worker to the repository. @@ -192,6 +196,7 @@ impl From for ApiInputTask { kind: task.kind.into(), domain: task.domain.to_string(), wildcard: Some(task.wildcard), + canister_id: task.canister_id, } } } diff --git a/ic-bn-lib/src/custom_domains/base/types/worker.rs b/ic-bn-lib/src/custom_domains/base/types/worker.rs index 1aabca8..4c592f5 100644 --- a/ic-bn-lib/src/custom_domains/base/types/worker.rs +++ b/ic-bn-lib/src/custom_domains/base/types/worker.rs @@ -868,6 +868,7 @@ mod tests { use tokio::{spawn, task, time::sleep}; use tokio_util::sync::CancellationToken; + use crate::principal; use crate::tls::acme::Error as AcmeError; use crate::{ custom_domains::base::{ @@ -950,6 +951,7 @@ mod tests { 123, None, false, + Some(principal!("rrkah-fqaaa-aaaaa-aaaaq-cai")), ))) }) }); @@ -960,9 +962,9 @@ mod tests { .returning(|_| Box::pin(async { Ok(()) })); let mut validator = MockValidatesDomains::new(); - validator.expect_validate().returning(|_| { - Box::pin(async { Ok(Principal::from_text("rrkah-fqaaa-aaaaa-aaaaq-cai").unwrap()) }) - }); + validator + .expect_validate() + .returning(|_| Box::pin(async { Ok(principal!("rrkah-fqaaa-aaaaa-aaaaq-cai")) })); // Create worker with short polling interval to speed up the test let config = @@ -1038,6 +1040,7 @@ mod tests { 123, certificate.clone(), false, + Some(principal!("rrkah-fqaaa-aaaaa-aaaaq-cai")), ); // Act @@ -1095,6 +1098,7 @@ mod tests { 123, None, false, + None, ); // Act @@ -1145,6 +1149,7 @@ mod tests { 123, None, false, + Some(principal!("rrkah-fqaaa-aaaaa-aaaaq-cai")), ); // Act @@ -1189,6 +1194,7 @@ mod tests { 123, None, false, + Some(principal!("rrkah-fqaaa-aaaaa-aaaaq-cai")), ); let task_result = TaskResult::success( @@ -1281,6 +1287,7 @@ mod tests { 456, None, false, + None, ); let task_result = TaskResult::failure( @@ -1359,6 +1366,7 @@ mod tests { 123, None, false, + Some(principal!("rrkah-fqaaa-aaaaa-aaaaq-cai")), ); let task_result = TaskResult::success( @@ -1415,6 +1423,7 @@ mod tests { 123, Some(vec![]), false, + None, ); let task_result = TaskResult::success( diff --git a/ic-bn-lib/src/custom_domains/client.rs b/ic-bn-lib/src/custom_domains/client.rs index 6123c1e..cf35454 100644 --- a/ic-bn-lib/src/custom_domains/client.rs +++ b/ic-bn-lib/src/custom_domains/client.rs @@ -297,6 +297,7 @@ impl Repository for CanisterClient { api_task.id, certificate, api_task.wildcard.unwrap_or(false), + api_task.canister_id, ); Ok(Some(task)) } diff --git a/ic-bn-lib/src/custom_domains/tests/helpers.rs b/ic-bn-lib/src/custom_domains/tests/helpers.rs index 30e0612..2798115 100644 --- a/ic-bn-lib/src/custom_domains/tests/helpers.rs +++ b/ic-bn-lib/src/custom_domains/tests/helpers.rs @@ -96,6 +96,7 @@ impl TestEnv { domain, kind, wildcard: None, + canister_id: None, }; let arg = Encode!(&task)?; From e04ddad9778cea9bc3f1c28fea2be7b74c43c2ab Mon Sep 17 00:00:00 2001 From: Igor Novgorodov Date: Fri, 28 Aug 2026 10:43:42 +0000 Subject: [PATCH 4/7] Fix update_task & tests --- .../src/custom_domains/base/types/worker.rs | 158 ++++++++++++------ 1 file changed, 103 insertions(+), 55 deletions(-) diff --git a/ic-bn-lib/src/custom_domains/base/types/worker.rs b/ic-bn-lib/src/custom_domains/base/types/worker.rs index 4c592f5..32c8c2e 100644 --- a/ic-bn-lib/src/custom_domains/base/types/worker.rs +++ b/ic-bn-lib/src/custom_domains/base/types/worker.rs @@ -31,7 +31,11 @@ use crate::{ DurationDisplay, custom_domains::base::{ helpers::{format_error_chain, retry_async}, - traits::{repository::Repository, time::UtcTimestamp, validation::ValidatesDomains}, + traits::{ + repository::Repository, + time::UtcTimestamp, + validation::{ValidatesDomains, ValidationError}, + }, types::task::{ IssueCertificateOutput, ScheduledTask, TaskFailReason, TaskKind, TaskOutcome, TaskOutput, TaskResult, @@ -217,6 +221,7 @@ impl Worker { let task_kind = task.kind; let certificate = task.cert; let wildcard = task.wildcard; + let canister_id = task.canister_id; info!("Task execution started"); @@ -232,6 +237,7 @@ impl Worker { task_id, task_kind, wildcard, + canister_id, ) .await } @@ -247,6 +253,7 @@ impl Worker { task_id, task_kind, wildcard, + canister_id, ) .await; @@ -266,7 +273,14 @@ impl Worker { } TaskKind::Update => { - update_task(domain, self.validator.clone(), task_id, task_kind).await + update_task( + domain, + self.validator.clone(), + task_id, + task_kind, + canister_id, + ) + .await } TaskKind::Delete => { @@ -617,29 +631,43 @@ async fn issue_task( task_id: UtcTimestamp, task_kind: TaskKind, wildcard: bool, + canister_id: Option, ) -> TaskResult { - match validator.validate(&domain).await { - Ok(canister_id) => { - Span::current().record("canister_id", canister_id.to_string()); - - match issue_certificate(&domain, canister_id, wildcard, acme_client).await { - Ok(output) => TaskResult::success(domain.clone(), output, task_id, task_kind), - Err(err) => { - let failure = match err.downcast_ref::() { - Some(err) if err.rate_limited() => TaskFailReason::RateLimited, - _ => TaskFailReason::GenericFailure(format_error_chain(&err)), - }; - TaskResult::failure(domain, failure, task_id, task_kind) - } + let failure = |e: ValidationError| { + TaskResult::failure( + domain.clone(), + TaskFailReason::ValidationFailed(e.to_string()), + task_id, + task_kind, + ) + }; + + let canister_id = if let Some(v) = canister_id { + if let Err(e) = validator.validate_limited(&domain).await { + return failure(e); + }; + + v + } else { + match validator.validate(&domain).await { + Ok(v) => v, + Err(e) => { + return failure(e); } } + }; - Err(err) => TaskResult::failure( - domain, - TaskFailReason::ValidationFailed(err.to_string()), - task_id, - task_kind, - ), + Span::current().record("canister_id", canister_id.to_string()); + + match issue_certificate(&domain, canister_id, wildcard, acme_client).await { + Ok(output) => TaskResult::success(domain.clone(), output, task_id, task_kind), + Err(err) => { + let failure = match err.downcast_ref::() { + Some(err) if err.rate_limited() => TaskFailReason::RateLimited, + _ => TaskFailReason::GenericFailure(format_error_chain(&err)), + }; + TaskResult::failure(domain, failure, task_id, task_kind) + } } } @@ -722,20 +750,34 @@ async fn update_task( validator: Arc, task_id: UtcTimestamp, task_kind: TaskKind, + canister_id: Option, ) -> TaskResult { - match validator.validate(&domain).await { - Ok(canister_id) => { - Span::current().record("canister_id", canister_id.to_string()); - TaskResult::success(domain, TaskOutput::Update(canister_id), task_id, task_kind) - } - - Err(err) => TaskResult::failure( - domain, - TaskFailReason::ValidationFailed(err.to_string()), + let failure = |e: ValidationError| { + TaskResult::failure( + domain.clone(), + TaskFailReason::ValidationFailed(e.to_string()), task_id, task_kind, - ), - } + ) + }; + + let canister_id = if let Some(v) = canister_id { + if let Err(e) = validator.validate_limited(&domain).await { + return failure(e); + }; + + v + } else { + match validator.validate(&domain).await { + Ok(v) => v, + Err(e) => { + return failure(e); + } + } + }; + + Span::current().record("canister_id", canister_id.to_string()); + TaskResult::success(domain, TaskOutput::Update(canister_id), task_id, task_kind) } /// Revokes a certificate using the ACME protocol. @@ -862,9 +904,9 @@ mod tests { use anyhow; use async_trait::async_trait; use candid::Principal; - use fqdn::FQDN; + use fqdn::fqdn; use prometheus::Registry; - use std::{str::FromStr, sync::Arc, time::Duration}; + use std::{sync::Arc, time::Duration}; use tokio::{spawn, task, time::sleep}; use tokio_util::sync::CancellationToken; @@ -947,7 +989,7 @@ mod tests { Box::pin(async { Ok(Some(ScheduledTask::new( TaskKind::Issue, - FQDN::from_str("example.org").unwrap(), + fqdn!("example.org"), 123, None, false, @@ -963,8 +1005,8 @@ mod tests { let mut validator = MockValidatesDomains::new(); validator - .expect_validate() - .returning(|_| Box::pin(async { Ok(principal!("rrkah-fqaaa-aaaaa-aaaaq-cai")) })); + .expect_validate_limited() + .returning(|_| Box::pin(async { Ok(()) })); // Create worker with short polling interval to speed up the test let config = @@ -1006,9 +1048,12 @@ mod tests { // Arrange let repository = MockRepository::new(); let mut validator = MockValidatesDomains::new(); - validator.expect_validate().returning(|_| { - Box::pin(async { Ok(Principal::from_text("rrkah-fqaaa-aaaaa-aaaaq-cai").unwrap()) }) - }); + validator + .expect_validate() + .returning(|_| Box::pin(async { Ok(principal!("rrkah-fqaaa-aaaaa-aaaaq-cai")) })); + validator + .expect_validate_limited() + .returning(|_| Box::pin(async { Ok(()) })); validator .expect_validate_deletion() .returning(|_| Box::pin(async { Ok(()) })); @@ -1028,19 +1073,22 @@ mod tests { let mut certificate = None; // Test all task kinds - for task_kind in [ - TaskKind::Issue, - TaskKind::Renew, - TaskKind::Update, - TaskKind::Delete, + for (task_kind, canister_id) in [ + ( + TaskKind::Issue, + Some(principal!("rrkah-fqaaa-aaaaa-aaaaq-cai")), + ), + (TaskKind::Renew, None), + (TaskKind::Update, None), + (TaskKind::Delete, None), ] { let task = ScheduledTask::new( task_kind, - FQDN::from_str("example.org").unwrap(), + fqdn!("example.org"), 123, certificate.clone(), false, - Some(principal!("rrkah-fqaaa-aaaaa-aaaaq-cai")), + canister_id, ); // Act @@ -1094,7 +1142,7 @@ mod tests { let task = ScheduledTask::new( TaskKind::Issue, - FQDN::from_str("example.org").unwrap(), + fqdn!("example.org"), 123, None, false, @@ -1123,11 +1171,11 @@ mod tests { // Arrange let repository = MockRepository::new(); let mut validator = MockValidatesDomains::new(); - validator.expect_validate().returning(|_| { + validator.expect_validate_limited().returning(|_| { Box::pin(async { // Slow validation causing task timeout sleep(Duration::from_millis(100)).await; - Ok(Principal::from_text("rrkah-fqaaa-aaaaa-aaaaq-cai").unwrap()) + Ok(()) }) }); @@ -1145,7 +1193,7 @@ mod tests { let task = ScheduledTask::new( TaskKind::Issue, - FQDN::from_str("example.org").unwrap(), + fqdn!("example.org"), 123, None, false, @@ -1190,7 +1238,7 @@ mod tests { let task = ScheduledTask::new( TaskKind::Issue, - FQDN::from_str("example.org").unwrap(), + fqdn!("example.org"), 123, None, false, @@ -1283,7 +1331,7 @@ mod tests { let task = ScheduledTask::new( TaskKind::Update, - FQDN::from_str("example.com").unwrap(), + fqdn!("example.com"), 456, None, false, @@ -1362,7 +1410,7 @@ mod tests { let task = ScheduledTask::new( TaskKind::Issue, - FQDN::from_str("example.net").unwrap(), + fqdn!("example.net"), 123, None, false, @@ -1419,7 +1467,7 @@ mod tests { let task = ScheduledTask::new( TaskKind::Delete, - FQDN::from_str("example.org").unwrap(), + fqdn!("example.org"), 123, Some(vec![]), false, From 5dcfda712be8477393990760ce561853a59999ae Mon Sep 17 00:00:00 2001 From: Igor Novgorodov Date: Fri, 28 Aug 2026 13:30:10 +0000 Subject: [PATCH 5/7] Add pending_canister_id & tests --- custom-domains/canister/src/state.rs | 292 +++++++++++++++++- .../custom_domains/backend/backend_service.rs | 245 ++++++++++++++- .../src/custom_domains/backend/router.rs | 47 ++- .../src/custom_domains/base/types/worker.rs | 4 + .../src/custom_domains/tests/e2e_test.rs | 89 +++++- 5 files changed, 658 insertions(+), 19 deletions(-) diff --git a/custom-domains/canister/src/state.rs b/custom-domains/canister/src/state.rs index 864c09d..94fa4c4 100644 --- a/custom-domains/canister/src/state.rs +++ b/custom-domains/canister/src/state.rs @@ -61,6 +61,16 @@ pub struct DomainEntry { /// after upgrade (they predate this field and default to `false`). #[serde(default)] pub wildcard: bool, + /// Canister ID explicitly requested (with the bypass token) for the currently + /// in-flight task, if any. `None` means the worker must derive and verify the + /// canister ID from the DNS. Set fresh on every `try_add_task` call and cleared + /// once the task finishes so it never leaks into a later task, such as a + /// canister-generated `Renew`. + /// + /// `#[serde(default)]` keeps existing stable-storage entries deserializable + /// after upgrade (they predate this field and default to `None`). + #[serde(default)] + pub pending_canister_id: Option, } #[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] @@ -299,7 +309,7 @@ impl CanisterState { now, enc_cert, Some(domain_entry.wildcard), - None, + domain_entry.pending_canister_id, ))) } None => Ok(None), @@ -506,6 +516,10 @@ impl CanisterState { entry.last_failure_reason = None; entry.rate_limit_failures_count = 0; entry.task_created_at = Some(now); + // Set fresh from this submission -- always overwritten (not preserved + // like `wildcard`) since the bypass decision applies to this specific + // task, not to the domain as a whole. + entry.pending_canister_id = task.canister_id; entry } @@ -518,6 +532,7 @@ impl CanisterState { let mut entry = DomainEntry::new(Some(task.kind), now); entry.task_created_at = Some(now); entry.wildcard = task.wildcard.unwrap_or(false); + entry.pending_canister_id = task.canister_id; entry } }; @@ -593,6 +608,7 @@ impl CanisterState { entry.rate_limit_failures_count = 0; self.last_change.set(now); entry.task_created_at = None; + entry.pending_canister_id = None; match output { TaskOutput::Issue(output) => { @@ -635,6 +651,7 @@ impl CanisterState { // Delete the task if the retry limit is reached if entry.failures_count >= MAX_TASK_FAILURES { entry.task = None; + entry.pending_canister_id = None; } } } @@ -943,6 +960,10 @@ mod tests { !entry.wildcard, "legacy entries must default wildcard to false" ); + assert_eq!( + entry.pending_canister_id, None, + "legacy entries must default pending_canister_id to None" + ); // Sanity check that the rest of the fields round-tripped correctly. assert_eq!(entry.task, Some(TaskKind::Issue)); assert_eq!(entry.failures_count, 3); @@ -1612,6 +1633,58 @@ mod tests { ); } + #[test] + fn test_try_add_task_stores_and_refreshes_pending_canister_id() { + let mut state = create_test_empty_state(); + let now = 1000; + let bypass_canister_id = Principal::from_text("rdmx6-jaaaa-aaaaa-aaadq-cai").unwrap(); + + // Issuing with a caller-supplied canister ID (the bypass path) stores it as the + // pending override for the in-flight task. + let domain = "bypass.example.com".to_string(); + state + .try_add_task( + InputTask { + domain: domain.clone(), + kind: TaskKind::Issue, + wildcard: None, + canister_id: Some(bypass_canister_id), + }, + now, + ) + .expect("issue task should be accepted"); + assert_eq!( + state.domains.get(&domain).unwrap().pending_canister_id, + Some(bypass_canister_id) + ); + + // Unlike `wildcard`, this is NOT preserved across submissions -- it reflects + // only the most recently submitted task, so a later Update without a + // caller-supplied canister ID must clear it back to `None` rather than keep + // trusting the previous task's bypass value. + let mut entry = state.domains.get(&domain).unwrap(); + entry.task = None; // simulate the issue task having completed + entry.not_after = Some(9999); // Update requires an existing certificate + state.domains.insert(domain.clone(), entry); + + state + .try_add_task( + InputTask { + domain: domain.clone(), + kind: TaskKind::Update, + wildcard: Some(false), + canister_id: None, + }, + now, + ) + .expect("update task should be accepted"); + assert_eq!( + state.domains.get(&domain).unwrap().pending_canister_id, + None, + "an Update without a bypass canister_id must clear any stale pending override" + ); + } + #[test] fn test_try_add_task_concurrent_tasks_prevention() { let mut state = create_test_empty_state(); @@ -1983,6 +2056,223 @@ mod tests { assert_eq!(result, expected_task); } + #[test] + fn test_fetch_next_task_propagates_pending_canister_id() { + // Arrange: a task submitted through the bypass path (with a caller-supplied + // canister ID) must hand that canister ID to the worker via `ScheduledTask`, so + // the worker knows to trust it (`validate_limited`) instead of re-deriving and + // re-verifying ownership from DNS (`validate`). + let mut state = create_test_empty_state(); + let now = 1000; + let bypass_canister_id = Principal::from_text("rdmx6-jaaaa-aaaaa-aaadq-cai").unwrap(); + + state + .try_add_task( + InputTask { + domain: "bypass.com".to_string(), + kind: TaskKind::Issue, + wildcard: None, + canister_id: Some(bypass_canister_id), + }, + now, + ) + .expect("failed to add a task"); + + // Act + let task = state.fetch_next_task(now).unwrap(); + + // Assert + let expected_task = Some(ScheduledTask::new( + TaskKindApi::Issue, + "bypass.com".to_string(), + now, + None, + Some(false), + Some(bypass_canister_id), + )); + assert_eq!(task, expected_task); + } + + #[test] + fn test_renewal_does_not_inherit_stale_pending_canister_id_from_prior_bypass_issue() { + // Arrange: register a domain via the bypass path (caller-supplied canister ID, + // skipping ownership verification) and let that task succeed. + let mut state = create_test_empty_state(); + let now = 1000; + let task_id = now; + let bypass_canister_id = Principal::from_text("rdmx6-jaaaa-aaaaa-aaadq-cai").unwrap(); + let domain = "bypass-renew.com".to_string(); + + state + .try_add_task( + InputTask { + domain: domain.clone(), + kind: TaskKind::Issue, + wildcard: None, + canister_id: Some(bypass_canister_id), + }, + now, + ) + .expect("failed to add a task"); + assert_eq!( + state.fetch_next_task(now).unwrap().unwrap().canister_id, + Some(bypass_canister_id), + "sanity check: the Issue task must be handed the bypass canister id" + ); + + let not_before = 0; + let not_after = now + 100; // due for renewal shortly + state + .submit_task_result( + TaskResult { + domain: domain.clone(), + task_id, + task_kind: TaskKind::Issue, + outcome: TaskOutcome::Success(TaskOutput::Issue(IssueCertificateOutput { + canister_id: bypass_canister_id, + enc_cert: b"cert".to_vec(), + enc_priv_key: b"key".to_vec(), + not_before, + not_after, + })), + duration_secs: 1, + }, + now, + ) + .expect("failed to submit task result"); + + assert_eq!( + state.domains.get(&domain).unwrap().pending_canister_id, + None, + "pending_canister_id must be cleared once the task it was submitted for completes" + ); + + // Act: fast-forward to when the certificate is nearing expiration, so the + // canister auto-schedules a `Renew` task for this domain. + let renewal_time = not_after - 1; + let task = state.fetch_next_task(renewal_time).unwrap(); + + // Assert: the internally-generated Renew task must NOT inherit the bypass + // canister id from the earlier, unrelated Issue task -- it never went through + // `try_add_task`, so the worker must fall back to full `validate()` and + // re-verify ownership from DNS, exactly as for any other Renew. + let expected_task = Some(ScheduledTask::new( + TaskKindApi::Renew, + domain, + renewal_time, + Some(b"cert".to_vec()), + Some(false), + None, + )); + assert_eq!(task, expected_task); + } + + #[test] + fn test_submit_task_result_success_clears_pending_canister_id() { + // Arrange + let mut state = create_test_empty_state(); + let now = 1000; + let task_id = 2u64; + let canister_id = Principal::from_text("rdmx6-jaaaa-aaaaa-aaadq-cai").unwrap(); + + let mut domain = DomainEntry::new(Some(TaskKind::Issue), now); + domain.taken_at = Some(task_id); + domain.pending_canister_id = Some(canister_id); + state.domains.insert("test.com".to_string(), domain); + + let task_result = TaskResult { + domain: "test.com".to_string(), + task_id, + task_kind: TaskKind::Issue, + outcome: TaskOutcome::Success(TaskOutput::Issue(IssueCertificateOutput { + canister_id, + enc_cert: b"cert".to_vec(), + enc_priv_key: b"key".to_vec(), + not_before: 0, + not_after: 9999, + })), + duration_secs: 1, + }; + + // Act + state.submit_task_result(task_result, now).unwrap(); + + // Assert + assert_eq!( + state + .domains + .get(&"test.com".to_string()) + .unwrap() + .pending_canister_id, + None + ); + } + + #[test] + fn test_submit_task_result_failure_preserves_pending_canister_id_for_retry() { + // A retried attempt of the SAME task must keep trusting the same bypass + // canister ID -- it's not a new submission, so the override must not be lost. + let mut state = create_test_empty_state(); + let now = 1000; + let task_id = 2u64; + let canister_id = Principal::from_text("rdmx6-jaaaa-aaaaa-aaadq-cai").unwrap(); + + let mut domain = DomainEntry::new(Some(TaskKind::Issue), now); + domain.taken_at = Some(task_id); + domain.pending_canister_id = Some(canister_id); + state.domains.insert("test.com".to_string(), domain); + + let task_result = TaskResult { + domain: "test.com".to_string(), + task_id, + task_kind: TaskKind::Issue, + outcome: TaskOutcome::Failure(TaskFailReason::GenericFailure("boom".to_string())), + duration_secs: 1, + }; + + // Act + state.submit_task_result(task_result, now).unwrap(); + + // Assert + let entry = state.domains.get(&"test.com".to_string()).unwrap(); + assert_eq!( + entry.task, + Some(TaskKind::Issue), + "task should still be retried" + ); + assert_eq!(entry.pending_canister_id, Some(canister_id)); + } + + #[test] + fn test_submit_task_result_failure_clears_pending_canister_id_once_retries_exhausted() { + let mut state = create_test_empty_state(); + let now = 1000; + let task_id = 2u64; + let canister_id = Principal::from_text("rdmx6-jaaaa-aaaaa-aaadq-cai").unwrap(); + + let mut domain = DomainEntry::new(Some(TaskKind::Issue), now); + domain.taken_at = Some(task_id); + domain.pending_canister_id = Some(canister_id); + domain.failures_count = MAX_TASK_FAILURES - 1; + state.domains.insert("test.com".to_string(), domain); + + let task_result = TaskResult { + domain: "test.com".to_string(), + task_id, + task_kind: TaskKind::Issue, + outcome: TaskOutcome::Failure(TaskFailReason::GenericFailure("boom".to_string())), + duration_secs: 1, + }; + + // Act: this failure pushes failures_count to MAX_TASK_FAILURES, dropping the task + state.submit_task_result(task_result, now).unwrap(); + + // Assert + let entry = state.domains.get(&"test.com".to_string()).unwrap(); + assert_eq!(entry.task, None); + assert_eq!(entry.pending_canister_id, None); + } + #[test] fn test_renewal_respects_failure_backoff_after_task_reset() { // A domain whose cert is expiring right now (so renewal is always diff --git a/ic-bn-lib/src/custom_domains/backend/backend_service.rs b/ic-bn-lib/src/custom_domains/backend/backend_service.rs index 31ed1f8..6c6bf99 100644 --- a/ic-bn-lib/src/custom_domains/backend/backend_service.rs +++ b/ic-bn-lib/src/custom_domains/backend/backend_service.rs @@ -42,17 +42,19 @@ impl BackendService { // If the canister ID is provided - use limited validation, otherwise full one // that derives the canister ID from the DNS TXT record - let canister_id = if let Some(canister_id) = canister_id { + let resolved_canister_id = if let Some(canister_id) = canister_id { self.validator.validate_limited(&fqdn).await?; canister_id } else { self.validator.validate(&fqdn).await? }; - let task = InputTask::new(task, fqdn, wildcard, Some(canister_id)); + // Carry forward only the caller-supplied (bypass) canister ID, not the resolved + // one. The worker uses this field to decide whether to trust it or derive it from DNS + let task = InputTask::new(task, fqdn, wildcard, canister_id); self.repository.try_add_task(task).await?; - Ok(canister_id) + Ok(resolved_canister_id) } /// Validates domain can be deleted and submits a delete task @@ -98,3 +100,240 @@ fn parse_domain(domain: &str) -> Result { FQDN::from_str(domain) .map_err(|e| ApiError::BadRequest(format!("Invalid domain format: {e:#}"))) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ + custom_domains::base::traits::{ + repository::{MockRepository, RepositoryError}, + validation::{MockValidatesDomains, ValidationError}, + }, + principal, + }; + + fn service( + repository: MockRepository, + validator: MockValidatesDomains, + bypass_token: Option, + ) -> BackendService { + BackendService::new(Arc::new(repository), Arc::new(validator), bypass_token) + } + + #[tokio::test] + async fn submit_task_with_canister_id_uses_limited_validation_and_preserves_it() { + // Arrange + let canister_id = principal!("rrkah-fqaaa-aaaaa-aaaaq-cai"); + + let mut validator = MockValidatesDomains::new(); + validator + .expect_validate_limited() + .times(1) + .returning(|_| Box::pin(async { Ok(()) })); + // `validate()` is intentionally left unmocked: it must not be called on the bypass path. + + let mut repository = MockRepository::new(); + repository + .expect_try_add_task() + .withf(move |task| task.canister_id == Some(canister_id)) + .times(1) + .returning(|_| Box::pin(async { Ok(()) })); + + let svc = service(repository, validator, None); + + // Act + let result = svc + .submit_task("example.org", TaskKind::Issue, false, Some(canister_id)) + .await + .unwrap(); + + // Assert + assert_eq!(result, canister_id); + } + + #[tokio::test] + async fn submit_task_without_canister_id_uses_full_validation_and_leaves_task_canister_id_none() + { + // Arrange + let derived_canister_id = principal!("rrkah-fqaaa-aaaaa-aaaaq-cai"); + + let mut validator = MockValidatesDomains::new(); + validator + .expect_validate() + .times(1) + .returning(move |_| Box::pin(async move { Ok(derived_canister_id) })); + // `validate_limited()` is intentionally left unmocked: it must not be called without a + // caller-supplied canister ID. + + let mut repository = MockRepository::new(); + repository + .expect_try_add_task() + // Regression guard: the task must carry `None`, not the derived canister_id, + // otherwise the worker would trust it and skip re-verifying ownership from DNS + // at execution time (see the comment on `submit_task`). + .withf(|task| task.canister_id.is_none()) + .times(1) + .returning(|_| Box::pin(async { Ok(()) })); + + let svc = service(repository, validator, None); + + // Act + let result = svc + .submit_task("example.org", TaskKind::Update, false, None) + .await + .unwrap(); + + // Assert + assert_eq!(result, derived_canister_id); + } + + #[tokio::test] + async fn submit_task_propagates_limited_validation_failure_without_adding_task() { + // Arrange + let canister_id = principal!("rrkah-fqaaa-aaaaa-aaaaq-cai"); + + let mut validator = MockValidatesDomains::new(); + validator.expect_validate_limited().times(1).returning(|_| { + Box::pin(async { + Err(ValidationError::ExistingDnsTxtChallenge { + src: "_acme-challenge.example.org.".to_string(), + }) + }) + }); + + let repository = MockRepository::new(); // try_add_task must not be called + + let svc = service(repository, validator, None); + + // Act + let err = svc + .submit_task("example.org", TaskKind::Issue, false, Some(canister_id)) + .await + .unwrap_err(); + + // Assert + assert!(matches!(err, ApiError::BadRequest(_))); + } + + #[tokio::test] + async fn submit_task_propagates_full_validation_failure_without_adding_task() { + // Arrange + let mut validator = MockValidatesDomains::new(); + validator.expect_validate().times(1).returning(|_| { + Box::pin(async { + Err(ValidationError::MissingDnsCname { + src: "_acme-challenge.example.org.".to_string(), + dst: "_acme-challenge.example.org.icp2.io.".to_string(), + }) + }) + }); + + let repository = MockRepository::new(); // try_add_task must not be called + + let svc = service(repository, validator, None); + + // Act + let err = svc + .submit_task("example.org", TaskKind::Issue, false, None) + .await + .unwrap_err(); + + // Assert + assert!(matches!(err, ApiError::BadRequest(_))); + } + + #[tokio::test] + async fn submit_task_propagates_repository_error() { + // Arrange + let mut validator = MockValidatesDomains::new(); + validator + .expect_validate() + .returning(|_| Box::pin(async { Ok(principal!("rrkah-fqaaa-aaaaa-aaaaq-cai")) })); + + let mut repository = MockRepository::new(); + repository.expect_try_add_task().returning(|_| { + Box::pin(async { + Err(RepositoryError::AnotherTaskInProgress( + FQDN::from_str("example.org").unwrap(), + )) + }) + }); + + let svc = service(repository, validator, None); + + // Act + let err = svc + .submit_task("example.org", TaskKind::Issue, false, None) + .await + .unwrap_err(); + + // Assert + assert!(matches!(err, ApiError::Conflict(_))); + } + + #[tokio::test] + async fn submit_task_rejects_invalid_domain_before_validating() { + // Arrange: no expectations set on either mock, so any call panics + let validator = MockValidatesDomains::new(); + let repository = MockRepository::new(); + + let svc = service(repository, validator, None); + + // Act + let err = svc + .submit_task("invalid..domain", TaskKind::Issue, false, None) + .await + .unwrap_err(); + + // Assert + assert!(matches!(err, ApiError::BadRequest(_))); + } + + #[tokio::test] + async fn submit_delete_task_success_carries_no_canister_id() { + // Arrange + let mut validator = MockValidatesDomains::new(); + validator + .expect_validate_deletion() + .times(1) + .returning(|_| Box::pin(async { Ok(()) })); + + let mut repository = MockRepository::new(); + repository + .expect_try_add_task() + .withf(|task| task.kind == TaskKind::Delete && task.canister_id.is_none()) + .times(1) + .returning(|_| Box::pin(async { Ok(()) })); + + let svc = service(repository, validator, None); + + // Act & Assert + svc.submit_delete_task("example.org").await.unwrap(); + } + + #[tokio::test] + async fn submit_delete_task_propagates_validation_failure() { + // Arrange + let mut validator = MockValidatesDomains::new(); + validator + .expect_validate_deletion() + .times(1) + .returning(|_| { + Box::pin(async { + Err(ValidationError::ExistingDnsTxtCanisterId { + src: "_canister-id.example.org.".to_string(), + }) + }) + }); + + let repository = MockRepository::new(); // try_add_task must not be called + + let svc = service(repository, validator, None); + + // Act + let err = svc.submit_delete_task("example.org").await.unwrap_err(); + + // Assert + assert!(matches!(err, ApiError::BadRequest(_))); + } +} diff --git a/ic-bn-lib/src/custom_domains/backend/router.rs b/ic-bn-lib/src/custom_domains/backend/router.rs index d30dbf2..34915ad 100644 --- a/ic-bn-lib/src/custom_domains/backend/router.rs +++ b/ic-bn-lib/src/custom_domains/backend/router.rs @@ -211,17 +211,20 @@ mod tests { let domain_normal = "example.org"; let subdomain_unicode = "тест.unicode.org"; + // The task must carry `canister_id: None` (not the derived value) since this + // request went through full validation, not the bypass path: the worker relies + // on `None` to know it must re-derive & re-verify ownership from DNS itself. let expected_task_normal = InputTask::new( TaskKind::Issue, FQDN::from_str(domain_normal).unwrap(), false, - Some(expected_canister_id), + None, ); let expected_task_unicode = InputTask::new( TaskKind::Issue, FQDN::from_str(subdomain_unicode).unwrap(), false, - Some(expected_canister_id), + None, ); mock_repository .expect_try_add_task() @@ -818,17 +821,19 @@ mod tests { let domain_normal = "example.org"; let subdomain_unicode = "тест.unicode.org"; + // See test_post_domain_success_accepted: without a bypass canister_id the task + // must carry `None` so the worker re-verifies ownership from DNS at execution time. let expected_task_normal = InputTask::new( TaskKind::Update, FQDN::from_str(domain_normal).unwrap(), false, - Some(expected_canister_id), + None, ); let expected_task_unicode = InputTask::new( TaskKind::Update, FQDN::from_str(subdomain_unicode).unwrap(), false, - Some(expected_canister_id), + None, ); mock_repository .expect_try_add_task() @@ -1187,11 +1192,15 @@ mod tests { // provided token doesn't match the configured bypass token. let domain = "example.org"; + // A wrong bypass token must fall back to full validation, and the task must carry + // `canister_id: None` (not the derived value) so the worker re-verifies ownership + // from DNS itself at execution time, rather than trusting a value an attacker with + // a wrong-but-plausible token could otherwise have steered towards `validate_limited`. let expected_task = InputTask::new( TaskKind::Issue, FQDN::from_str(domain).unwrap(), false, - Some(derived_canister_id), + None, ); let mut mock_repository = MockRepository::new(); mock_repository @@ -1238,6 +1247,10 @@ mod tests { let mut mock_repository = MockRepository::new(); mock_repository .expect_try_add_task() + // Full-validation path: the task must not carry the caller's query-string + // canister_id forward, otherwise the worker would skip re-verifying + // ownership from DNS at execution time. + .withf(|task| task.canister_id.is_none()) .times(1) .returning(|_| Box::pin(async { Ok(()) })); @@ -1279,6 +1292,10 @@ mod tests { let mut mock_repository = MockRepository::new(); mock_repository .expect_try_add_task() + // Full-validation path: the task must not carry the caller's query-string + // canister_id forward, otherwise the worker would skip re-verifying + // ownership from DNS at execution time. + .withf(|task| task.canister_id.is_none()) .times(1) .returning(|_| Box::pin(async { Ok(()) })); @@ -1320,6 +1337,10 @@ mod tests { let mut mock_repository = MockRepository::new(); mock_repository .expect_try_add_task() + // Full-validation path: the task must not carry the caller's query-string + // canister_id forward, otherwise the worker would skip re-verifying + // ownership from DNS at execution time. + .withf(|task| task.canister_id.is_none()) .times(1) .returning(|_| Box::pin(async { Ok(()) })); @@ -1472,11 +1493,13 @@ mod tests { // provided token doesn't match the configured bypass token. let domain = "example.org"; + // See test_post_domain_wrong_bypass_token_falls_back_to_full_validation: the task + // must carry `canister_id: None`, not the derived value. let expected_task = InputTask::new( TaskKind::Update, FQDN::from_str(domain).unwrap(), false, - Some(derived_canister_id), + None, ); let mut mock_repository = MockRepository::new(); mock_repository @@ -1523,6 +1546,10 @@ mod tests { let mut mock_repository = MockRepository::new(); mock_repository .expect_try_add_task() + // Full-validation path: the task must not carry the caller's query-string + // canister_id forward, otherwise the worker would skip re-verifying + // ownership from DNS at execution time. + .withf(|task| task.canister_id.is_none()) .times(1) .returning(|_| Box::pin(async { Ok(()) })); @@ -1564,6 +1591,10 @@ mod tests { let mut mock_repository = MockRepository::new(); mock_repository .expect_try_add_task() + // Full-validation path: the task must not carry the caller's query-string + // canister_id forward, otherwise the worker would skip re-verifying + // ownership from DNS at execution time. + .withf(|task| task.canister_id.is_none()) .times(1) .returning(|_| Box::pin(async { Ok(()) })); @@ -1605,6 +1636,10 @@ mod tests { let mut mock_repository = MockRepository::new(); mock_repository .expect_try_add_task() + // Full-validation path: the task must not carry the caller's query-string + // canister_id forward, otherwise the worker would skip re-verifying + // ownership from DNS at execution time. + .withf(|task| task.canister_id.is_none()) .times(1) .returning(|_| Box::pin(async { Ok(()) })); diff --git a/ic-bn-lib/src/custom_domains/base/types/worker.rs b/ic-bn-lib/src/custom_domains/base/types/worker.rs index 32c8c2e..bcb8a0f 100644 --- a/ic-bn-lib/src/custom_domains/base/types/worker.rs +++ b/ic-bn-lib/src/custom_domains/base/types/worker.rs @@ -642,6 +642,8 @@ async fn issue_task( ) }; + // If a canister ID is provided, validate it with limited checks. + // Otherwise, perform full validation to derive the canister ID from DNS. let canister_id = if let Some(v) = canister_id { if let Err(e) = validator.validate_limited(&domain).await { return failure(e); @@ -761,6 +763,8 @@ async fn update_task( ) }; + // If a canister ID is provided, validate it with limited checks. + // Otherwise, perform full validation to derive the canister ID from DNS. let canister_id = if let Some(v) = canister_id { if let Err(e) = validator.validate_limited(&domain).await { return failure(e); diff --git a/ic-bn-lib/src/custom_domains/tests/e2e_test.rs b/ic-bn-lib/src/custom_domains/tests/e2e_test.rs index 838e9d1..0078dcc 100644 --- a/ic-bn-lib/src/custom_domains/tests/e2e_test.rs +++ b/ic-bn-lib/src/custom_domains/tests/e2e_test.rs @@ -45,6 +45,10 @@ const MAX_CANISTER_CALL_RETRY_DELAY: Duration = Duration::from_secs(2); const DOMAINS_COUNT: usize = 160; const WORKERS_COUNT: usize = 4; +/// Bypass token configured on the API server, used by `bypass_canister_id_e2e` to submit a +/// registration with an explicit `canister_id` instead of going through `MockValidator::validate()`. +const BYPASS_TOKEN: &str = "e2e-test-bypass-token"; + // Title: Custom Domains with Pebble ACME test server and multiple workers processing registration requests in parallel // Setup: // - Start Pocket IC and install the Custom Domains canister @@ -57,10 +61,14 @@ const WORKERS_COUNT: usize = 4; // Each worker should start picking up tasks and obtain certificates in parallel // 2. Verify all domains have been registered after all tasks are processed // 3. Download all certificates and verify they match the requested domains -// 4. Submit half of the registered domains for deletion via the API calls -// 5. Verify all these domains are eventually deleted from the canister -// 6. Get canister metrics and verify the expected stats of domain registrations -// 7. Get workers metrics and verify they all have processed more than one task each +// 4. Register a wildcard domain and verify its certificate carries both SANs +// 5. Register a domain via the bypass token with an explicit canister_id and verify +// that canister_id (not the one MockValidator::validate() would derive) is what +// ends up registered +// 6. Submit half of the registered domains for deletion via the API calls +// 7. Verify all these domains are eventually deleted from the canister +// 8. Get canister metrics and verify the expected stats of domain registrations +// 9. Get workers metrics and verify they all have processed more than one task each #[ignore] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] @@ -85,19 +93,25 @@ async fn e2e_pebble_test() -> anyhow::Result<()> { info!("Step 4: Registering a wildcard domain and verifying its certificate carries both SANs"); wildcard_domain_e2e(&ctx).await?; - info!("Step 5: Submitting half of the registered domains for deletion via the API calls"); + info!( + "Step 5: Registering a domain via the bypass token with an explicit canister_id, \ + verifying it skips derivation and the requested canister_id is what gets registered" + ); + bypass_canister_id_e2e(&ctx).await?; + + info!("Step 6: Submitting half of the registered domains for deletion via the API calls"); let deleted_domains = delete_half_domains(&ctx, domains).await?; - info!("Step 6: Verifying all these domains are eventually deleted from the canister"); + info!("Step 7: Verifying all these domains are eventually deleted from the canister"); verify_domains_deletion(&ctx, deleted_domains).await?; info!( - "Step 7: Getting canister metrics and verifying the expected stats of domain registrations" + "Step 8: Getting canister metrics and verifying the expected stats of domain registrations" ); verify_canister_metrics(&ctx).await?; info!( - "Step 8: Getting workers metrics and verifying they all have processed more than one task each" + "Step 9: Getting workers metrics and verifying they all have processed more than one task each" ); verify_workers_metrics_with_retries(workers_metrics).await?; @@ -277,7 +291,7 @@ async fn spawn_api_server( prometheus_registry, RateLimitConfig::default(), true, - None, + Some(BYPASS_TOKEN.to_string()), ); info!("Starting API server at http://{}", api_addr); axum_server::bind(api_addr) @@ -444,6 +458,63 @@ async fn wildcard_domain_e2e(ctx: &TestContext) -> anyhow::Result<()> { Ok(()) } +/// Registers a domain using the bypass token plus an explicit `canister_id`, and verifies +/// that the requested canister_id (not the one `MockValidator::validate()` would derive) is +/// what actually ends up registered against the domain once the async task completes. +/// +/// This exercises the full pipeline end-to-end: the API handler must forward the bypass +/// canister_id into the task, the canister must persist and hand it back to a worker via +/// `fetch_next_task`, and the worker must use it directly (via `validate_limited`) rather +/// than falling back to `MockValidator::validate()`'s fixed canister id. +async fn bypass_canister_id_e2e(ctx: &TestContext) -> anyhow::Result<()> { + let domain = "bypass-domain.example.com".to_string(); + // Distinct from the fixed canister id `MockValidator::validate()` always returns + // ("laqa6-raaaa-aaaam-aehzq-cai"), so a successful bypass is unambiguous. + let bypass_canister_id: Principal = "2vxsx-fae".parse()?; + + let client = reqwest::Client::new(); + let base_url = format!( + "http://{}:{}/v1/{domain}", + ctx.api_server_addr.ip(), + ctx.api_server_addr.port() + ); + + // Register the domain via the bypass path + let response = client + .post(format!("{base_url}?canister_id={bypass_canister_id}")) + .bearer_auth(BYPASS_TOKEN) + .send() + .await?; + assert!( + response.status().is_success(), + "bypass registration request was rejected: {}", + response.status() + ); + + wait_for_all_tasks_completion(&ctx.canister_repository).await?; + verify_domains_registration(ctx, &vec![domain.clone()]).await?; + + let status = ctx + .canister_repository + .get_domain_status(&domain.parse()?) + .await? + .with_context(|| format!("domain {domain} not found after bypass registration"))?; + assert_eq!(status.status, RegistrationStatus::Registered); + assert_eq!( + status.canister_id, + Some(bypass_canister_id), + "domain must be registered against the bypass-supplied canister_id, not the one \ + MockValidator::validate() would have derived" + ); + + // Clean up so later registration-count assertions remain valid + let response = client.delete(&base_url).send().await?; + assert!(response.status().is_success()); + verify_domains_deletion(ctx, vec![domain]).await?; + + Ok(()) +} + async fn delete_half_domains( ctx: &TestContext, domains: Vec, From a6481b3e0965d133d47bcbd45083c513dce09035 Mon Sep 17 00:00:00 2001 From: Igor Novgorodov Date: Fri, 28 Aug 2026 15:27:41 +0000 Subject: [PATCH 6/7] Rework to use `bypass_validation` instead --- custom-domains/canister/src/state.rs | 380 ++++++++++++++++++++------- 1 file changed, 284 insertions(+), 96 deletions(-) diff --git a/custom-domains/canister/src/state.rs b/custom-domains/canister/src/state.rs index 94fa4c4..0573f47 100644 --- a/custom-domains/canister/src/state.rs +++ b/custom-domains/canister/src/state.rs @@ -44,7 +44,10 @@ pub struct DomainEntry { pub failures_count: u32, /// Number of rate limit failures for the current task pub rate_limit_failures_count: u32, - /// Canister ID associated with the domain + /// Canister ID associated with the domain. Normally only confirmed once a task + /// actually succeeds, but when `bypass_validation` is set it is also written + /// at submission time, since `fetch_next_task` needs it *before* the task runs + /// in order to tell the worker to trust it. pub canister_id: Option, /// Timestamp when the domain entry was created (set once and never updated) pub created_at: UtcTimestamp, @@ -61,16 +64,14 @@ pub struct DomainEntry { /// after upgrade (they predate this field and default to `false`). #[serde(default)] pub wildcard: bool, - /// Canister ID explicitly requested (with the bypass token) for the currently - /// in-flight task, if any. `None` means the worker must derive and verify the - /// canister ID from the DNS. Set fresh on every `try_add_task` call and cleared - /// once the task finishes so it never leaks into a later task, such as a - /// canister-generated `Renew`. - /// + /// Whether validation should always be bypassed for this domain: `canister_id` is + /// trusted instead of being derived from DNS TXT record. + /// It can be enabled only once at domain creation time, and can only be revoked + /// by deleting and re-creating the domain. /// `#[serde(default)]` keeps existing stable-storage entries deserializable - /// after upgrade (they predate this field and default to `None`). + /// after upgrade (they predate this field and default to `false`). #[serde(default)] - pub pending_canister_id: Option, + pub bypass_validation: bool, } #[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] @@ -309,7 +310,14 @@ impl CanisterState { now, enc_cert, Some(domain_entry.wildcard), - domain_entry.pending_canister_id, + // Only trust `canister_id` outright when this domain is flagged for + // bypass -- applies uniformly to Issue, Update, and a + // canister-generated Renew alike, since none of them re-derive this + // decision themselves. + domain_entry + .bypass_validation + .then_some(domain_entry.canister_id) + .flatten(), ))) } None => Ok(None), @@ -516,10 +524,19 @@ impl CanisterState { entry.last_failure_reason = None; entry.rate_limit_failures_count = 0; entry.task_created_at = Some(now); - // Set fresh from this submission -- always overwritten (not preserved - // like `wildcard`) since the bypass decision applies to this specific - // task, not to the domain as a whole. - entry.pending_canister_id = task.canister_id; + + // Bypass, once granted, can only be revoked by deleting and + // re-creating the domain. So only the bypass path may touch these + // fields, and only to turn bypass on / retarget the canister_id; a + // plain submission must leave both exactly as they are, + // otherwise anyone without the token could downgrade a bypassed domain + // back to full validation at will. + if let (TaskKind::Issue | TaskKind::Update, Some(canister_id)) = + (task.kind, task.canister_id) + { + entry.bypass_validation = true; + entry.canister_id = Some(canister_id); + } entry } @@ -532,7 +549,8 @@ impl CanisterState { let mut entry = DomainEntry::new(Some(task.kind), now); entry.task_created_at = Some(now); entry.wildcard = task.wildcard.unwrap_or(false); - entry.pending_canister_id = task.canister_id; + entry.bypass_validation = task.canister_id.is_some(); + entry.canister_id = task.canister_id; entry } }; @@ -608,7 +626,6 @@ impl CanisterState { entry.rate_limit_failures_count = 0; self.last_change.set(now); entry.task_created_at = None; - entry.pending_canister_id = None; match output { TaskOutput::Issue(output) => { @@ -651,7 +668,6 @@ impl CanisterState { // Delete the task if the retry limit is reached if entry.failures_count >= MAX_TASK_FAILURES { entry.task = None; - entry.pending_canister_id = None; } } } @@ -960,9 +976,9 @@ mod tests { !entry.wildcard, "legacy entries must default wildcard to false" ); - assert_eq!( - entry.pending_canister_id, None, - "legacy entries must default pending_canister_id to None" + assert!( + !entry.bypass_validation, + "legacy entries must default bypass_validation to false" ); // Sanity check that the rest of the fields round-tripped correctly. assert_eq!(entry.task, Some(TaskKind::Issue)); @@ -1634,13 +1650,13 @@ mod tests { } #[test] - fn test_try_add_task_stores_and_refreshes_pending_canister_id() { + fn test_try_add_task_sets_bypass_validation_and_canister_id_on_issue() { let mut state = create_test_empty_state(); let now = 1000; let bypass_canister_id = Principal::from_text("rdmx6-jaaaa-aaaaa-aaadq-cai").unwrap(); - // Issuing with a caller-supplied canister ID (the bypass path) stores it as the - // pending override for the in-flight task. + // Issuing with a caller-supplied canister ID (the bypass path) flags the domain + // for bypass and records the canister ID immediately, before the task even runs. let domain = "bypass.example.com".to_string(); state .try_add_task( @@ -1653,15 +1669,58 @@ mod tests { now, ) .expect("issue task should be accepted"); - assert_eq!( - state.domains.get(&domain).unwrap().pending_canister_id, - Some(bypass_canister_id) - ); - // Unlike `wildcard`, this is NOT preserved across submissions -- it reflects - // only the most recently submitted task, so a later Update without a - // caller-supplied canister ID must clear it back to `None` rather than keep - // trusting the previous task's bypass value. + let entry = state.domains.get(&domain).unwrap(); + assert!(entry.bypass_validation); + assert_eq!(entry.canister_id, Some(bypass_canister_id)); + } + + #[test] + fn test_try_add_task_issue_without_canister_id_does_not_set_bypass() { + let mut state = create_test_empty_state(); + let now = 1000; + + let domain = "normal.example.com".to_string(); + state + .try_add_task( + InputTask { + domain: domain.clone(), + kind: TaskKind::Issue, + wildcard: None, + canister_id: None, + }, + now, + ) + .expect("issue task should be accepted"); + + let entry = state.domains.get(&domain).unwrap(); + assert!(!entry.bypass_validation); + assert_eq!(entry.canister_id, None); + } + + #[test] + fn test_try_add_task_bypass_update_changes_canister_id() { + // A bypass domain's canister_id must be changeable by a later bypass Update + // (e.g. migrating the domain to a different canister) -- the worker needs the + // NEW target, not the one recorded at creation. + let mut state = create_test_empty_state(); + let now = 1000; + let original_canister_id = Principal::from_text("rdmx6-jaaaa-aaaaa-aaadq-cai").unwrap(); + let new_canister_id = Principal::from_text("rrkah-fqaaa-aaaaa-aaaaq-cai").unwrap(); + let domain = "bypass.example.com".to_string(); + + state + .try_add_task( + InputTask { + domain: domain.clone(), + kind: TaskKind::Issue, + wildcard: None, + canister_id: Some(original_canister_id), + }, + now, + ) + .expect("issue task should be accepted"); + let mut entry = state.domains.get(&domain).unwrap(); entry.task = None; // simulate the issue task having completed entry.not_after = Some(9999); // Update requires an existing certificate @@ -1672,17 +1731,150 @@ mod tests { InputTask { domain: domain.clone(), kind: TaskKind::Update, - wildcard: Some(false), + wildcard: None, + canister_id: Some(new_canister_id), + }, + now, + ) + .expect("update task should be accepted"); + + let entry = state.domains.get(&domain).unwrap(); + assert!(entry.bypass_validation); + assert_eq!(entry.canister_id, Some(new_canister_id)); + } + + #[test] + fn test_try_add_task_non_bypass_update_does_not_clear_bypass_validation() { + // Bypass is one-way: a plain (non-bypass) Update must not be able to downgrade + // a bypassed domain back to full validation, otherwise anyone without the + // bypass token could revoke it. The only way to clear it is to delete and + // re-create the domain. + let mut state = create_test_empty_state(); + let now = 1000; + let bypass_canister_id = Principal::from_text("rdmx6-jaaaa-aaaaa-aaadq-cai").unwrap(); + let domain = "bypass.example.com".to_string(); + + state + .try_add_task( + InputTask { + domain: domain.clone(), + kind: TaskKind::Issue, + wildcard: None, + canister_id: Some(bypass_canister_id), + }, + now, + ) + .expect("issue task should be accepted"); + + let mut entry = state.domains.get(&domain).unwrap(); + entry.task = None; // simulate the issue task having completed + entry.not_after = Some(9999); // Update requires an existing certificate + state.domains.insert(domain.clone(), entry); + + state + .try_add_task( + InputTask { + domain: domain.clone(), + kind: TaskKind::Update, + wildcard: None, canister_id: None, }, now, ) .expect("update task should be accepted"); - assert_eq!( - state.domains.get(&domain).unwrap().pending_canister_id, - None, - "an Update without a bypass canister_id must clear any stale pending override" + + let entry = state.domains.get(&domain).unwrap(); + assert!( + entry.bypass_validation, + "a non-bypass submission must not be able to clear bypass_validation" ); + assert_eq!(entry.canister_id, Some(bypass_canister_id)); + } + + #[test] + fn test_try_add_task_delete_and_recreate_clears_bypass_validation() { + // The only sanctioned way to clear bypass_validation: delete the domain + // entirely (dropping its DomainEntry) and re-create it via a fresh, non-bypass + // Issue. + let mut state = create_test_empty_state(); + let now = 1000; + let bypass_canister_id = Principal::from_text("rdmx6-jaaaa-aaaaa-aaadq-cai").unwrap(); + let domain = "bypass.example.com".to_string(); + + state + .try_add_task( + InputTask { + domain: domain.clone(), + kind: TaskKind::Issue, + wildcard: None, + canister_id: Some(bypass_canister_id), + }, + now, + ) + .expect("issue task should be accepted"); + assert!(state.domains.get(&domain).unwrap().bypass_validation); + + // Simulate the domain having actually been deleted (Delete task succeeded). + state.domains.remove(&domain); + + state + .try_add_task( + InputTask { + domain: domain.clone(), + kind: TaskKind::Issue, + wildcard: None, + canister_id: None, + }, + now, + ) + .expect("issue task on the re-created domain should be accepted"); + + let entry = state.domains.get(&domain).unwrap(); + assert!(!entry.bypass_validation); + assert_eq!(entry.canister_id, None); + } + + #[test] + fn test_try_add_task_delete_does_not_disturb_bypass_validation() { + // A Delete never carries a canister_id; it must not clear bypass_validation for + // a domain it fails to actually delete (e.g. DNS validation rejects it later). + let mut state = create_test_empty_state(); + let now = 1000; + let bypass_canister_id = Principal::from_text("rdmx6-jaaaa-aaaaa-aaadq-cai").unwrap(); + let domain = "bypass.example.com".to_string(); + + state + .try_add_task( + InputTask { + domain: domain.clone(), + kind: TaskKind::Issue, + wildcard: None, + canister_id: Some(bypass_canister_id), + }, + now, + ) + .expect("issue task should be accepted"); + + let mut entry = state.domains.get(&domain).unwrap(); + entry.task = None; // simulate the issue task having completed + entry.not_after = Some(9999); + state.domains.insert(domain.clone(), entry); + + state + .try_add_task( + InputTask { + domain: domain.clone(), + kind: TaskKind::Delete, + wildcard: None, + canister_id: None, + }, + now, + ) + .expect("delete task should be accepted"); + + let entry = state.domains.get(&domain).unwrap(); + assert!(entry.bypass_validation); + assert_eq!(entry.canister_id, Some(bypass_canister_id)); } #[test] @@ -2057,7 +2249,7 @@ mod tests { } #[test] - fn test_fetch_next_task_propagates_pending_canister_id() { + fn test_fetch_next_task_propagates_canister_id_when_bypass_validation_is_set() { // Arrange: a task submitted through the bypass path (with a caller-supplied // canister ID) must hand that canister ID to the worker via `ScheduledTask`, so // the worker knows to trust it (`validate_limited`) instead of re-deriving and @@ -2094,7 +2286,38 @@ mod tests { } #[test] - fn test_renewal_does_not_inherit_stale_pending_canister_id_from_prior_bypass_issue() { + fn test_fetch_next_task_does_not_propagate_canister_id_without_bypass_validation() { + // A domain's `canister_id` may already be confirmed (from a past successful + // task), but without `bypass_validation` set the worker must still re-derive and + // re-verify ownership from DNS rather than trust the stored value outright. + let mut state = create_test_empty_state(); + let now = 1000; + let canister_id = Principal::from_text("rdmx6-jaaaa-aaaaa-aaadq-cai").unwrap(); + + let mut domain = DomainEntry::new(Some(TaskKind::Update), now); + domain.task_created_at = Some(now); + domain.canister_id = Some(canister_id); + domain.not_after = Some(9999); // Update requires an existing certificate + domain.bypass_validation = false; + state.domains.insert("normal.com".to_string(), domain); + + // Act + let task = state.fetch_next_task(now).unwrap(); + + // Assert + let expected_task = Some(ScheduledTask::new( + TaskKindApi::Update, + "normal.com".to_string(), + now, + None, + Some(false), + None, + )); + assert_eq!(task, expected_task); + } + + #[test] + fn test_renewal_inherits_canister_id_when_bypass_validation_is_set() { // Arrange: register a domain via the bypass path (caller-supplied canister ID, // skipping ownership verification) and let that task succeed. let mut state = create_test_empty_state(); @@ -2141,35 +2364,36 @@ mod tests { ) .expect("failed to submit task result"); - assert_eq!( - state.domains.get(&domain).unwrap().pending_canister_id, - None, - "pending_canister_id must be cleared once the task it was submitted for completes" + let entry = state.domains.get(&domain).unwrap(); + assert!( + entry.bypass_validation, + "bypass_validation is a sticky, per-domain decision and must survive task completion" ); + assert_eq!(entry.canister_id, Some(bypass_canister_id)); // Act: fast-forward to when the certificate is nearing expiration, so the // canister auto-schedules a `Renew` task for this domain. let renewal_time = not_after - 1; let task = state.fetch_next_task(renewal_time).unwrap(); - // Assert: the internally-generated Renew task must NOT inherit the bypass - // canister id from the earlier, unrelated Issue task -- it never went through - // `try_add_task`, so the worker must fall back to full `validate()` and - // re-verify ownership from DNS, exactly as for any other Renew. + // Assert: the internally-generated Renew task must inherit the domain's bypass + // canister id -- it never goes through `try_add_task` itself, so it relies + // entirely on the sticky `bypass_validation`/`canister_id` set by the Issue. let expected_task = Some(ScheduledTask::new( TaskKindApi::Renew, domain, renewal_time, Some(b"cert".to_vec()), Some(false), - None, + Some(bypass_canister_id), )); assert_eq!(task, expected_task); } #[test] - fn test_submit_task_result_success_clears_pending_canister_id() { - // Arrange + fn test_submit_task_result_failure_preserves_bypass_validation_and_canister_id() { + // A retried attempt of the same task must keep trusting the same bypass + // canister ID -- `submit_task_result` no longer touches either field on failure. let mut state = create_test_empty_state(); let now = 1000; let task_id = 2u64; @@ -2177,49 +2401,8 @@ mod tests { let mut domain = DomainEntry::new(Some(TaskKind::Issue), now); domain.taken_at = Some(task_id); - domain.pending_canister_id = Some(canister_id); - state.domains.insert("test.com".to_string(), domain); - - let task_result = TaskResult { - domain: "test.com".to_string(), - task_id, - task_kind: TaskKind::Issue, - outcome: TaskOutcome::Success(TaskOutput::Issue(IssueCertificateOutput { - canister_id, - enc_cert: b"cert".to_vec(), - enc_priv_key: b"key".to_vec(), - not_before: 0, - not_after: 9999, - })), - duration_secs: 1, - }; - - // Act - state.submit_task_result(task_result, now).unwrap(); - - // Assert - assert_eq!( - state - .domains - .get(&"test.com".to_string()) - .unwrap() - .pending_canister_id, - None - ); - } - - #[test] - fn test_submit_task_result_failure_preserves_pending_canister_id_for_retry() { - // A retried attempt of the SAME task must keep trusting the same bypass - // canister ID -- it's not a new submission, so the override must not be lost. - let mut state = create_test_empty_state(); - let now = 1000; - let task_id = 2u64; - let canister_id = Principal::from_text("rdmx6-jaaaa-aaaaa-aaadq-cai").unwrap(); - - let mut domain = DomainEntry::new(Some(TaskKind::Issue), now); - domain.taken_at = Some(task_id); - domain.pending_canister_id = Some(canister_id); + domain.bypass_validation = true; + domain.canister_id = Some(canister_id); state.domains.insert("test.com".to_string(), domain); let task_result = TaskResult { @@ -2240,11 +2423,14 @@ mod tests { Some(TaskKind::Issue), "task should still be retried" ); - assert_eq!(entry.pending_canister_id, Some(canister_id)); + assert!(entry.bypass_validation); + assert_eq!(entry.canister_id, Some(canister_id)); } #[test] - fn test_submit_task_result_failure_clears_pending_canister_id_once_retries_exhausted() { + fn test_submit_task_result_max_failures_preserves_bypass_validation_and_canister_id() { + // Bypass is a domain-level property, not per-task state: exhausting retries + // drops the task itself, but must not disturb the sticky decision. let mut state = create_test_empty_state(); let now = 1000; let task_id = 2u64; @@ -2252,7 +2438,8 @@ mod tests { let mut domain = DomainEntry::new(Some(TaskKind::Issue), now); domain.taken_at = Some(task_id); - domain.pending_canister_id = Some(canister_id); + domain.bypass_validation = true; + domain.canister_id = Some(canister_id); domain.failures_count = MAX_TASK_FAILURES - 1; state.domains.insert("test.com".to_string(), domain); @@ -2270,7 +2457,8 @@ mod tests { // Assert let entry = state.domains.get(&"test.com".to_string()).unwrap(); assert_eq!(entry.task, None); - assert_eq!(entry.pending_canister_id, None); + assert!(entry.bypass_validation); + assert_eq!(entry.canister_id, Some(canister_id)); } #[test] From a89f4c9c2c39b4c677782b494306539688bd6961 Mon Sep 17 00:00:00 2001 From: Igor Novgorodov Date: Mon, 31 Aug 2026 12:12:48 +0000 Subject: [PATCH 7/7] Fix nits, rework rate limiter to use generic bypasser instead of a token --- .../src/custom_domains/backend/router.rs | 45 +++++- .../custom_domains/base/traits/validation.rs | 4 +- .../src/custom_domains/base/types/worker.rs | 42 +++--- ic-bn-lib/src/http/middleware/rate_limiter.rs | 141 +++++++++++------- 4 files changed, 151 insertions(+), 81 deletions(-) diff --git a/ic-bn-lib/src/custom_domains/backend/router.rs b/ic-bn-lib/src/custom_domains/backend/router.rs index 34915ad..7173eb6 100644 --- a/ic-bn-lib/src/custom_domains/backend/router.rs +++ b/ic-bn-lib/src/custom_domains/backend/router.rs @@ -9,6 +9,8 @@ use axum::{ routing::{delete, get, patch, post}, }; use axum_extra::middleware::option_layer; +use bytes::Bytes; +use http::header::AUTHORIZATION; use prometheus::Registry; use tower_http::{ LatencyUnit, @@ -26,11 +28,45 @@ use super::{ metrics::{HttpMetrics, metrics_handler, metrics_middleware}, }; use crate::{ + constant_time_eq, custom_domains::base::traits::{repository::Repository, validation::ValidatesDomains}, - http::middleware::rate_limiter::layer_by_ip, + http::middleware::rate_limiter::{BYPASS_TOKEN_HEADER, Bypasser, NeverBypasser, layer_by_ip}, reqwest::StatusCode, }; +/// This bypasser checks AUTHORIZATION header for a token first, +/// then the `x-ratelimit-bypass-token` header if not found. +#[derive(Clone)] +struct AuthTokenBypasser(Bytes); + +impl Bypasser for AuthTokenBypasser { + fn should_bypass(&self, req: &Request) -> bool { + req.headers() + .get(AUTHORIZATION) + .filter(|x| x.as_bytes().starts_with(b"Bearer ")) + .and_then(|x| x.as_bytes().get(7..)) + .or_else(|| req.headers().get(BYPASS_TOKEN_HEADER).map(|x| x.as_bytes())) + .is_some_and(|x| constant_time_eq(x, &self.0)) + } +} + +/// Uses either `AuthTokenBypasser` or `NeverBypasser` depending on whether a +/// bypass token is configured. Needed because bypasser is generic and we need a single type. +#[derive(Clone)] +enum RateLimitBypasser { + Token(AuthTokenBypasser), + Never(NeverBypasser), +} + +impl Bypasser for RateLimitBypasser { + fn should_bypass(&self, req: &Request) -> bool { + match self { + Self::Token(b) => b.should_bypass(req), + Self::Never(b) => b.should_bypass(req), + } + } +} + /// Options for configuring rate limits on various endpoints. #[derive(Clone, Debug, Default)] pub struct RateLimitConfig { @@ -52,10 +88,15 @@ pub fn create_router( let backend_service = BackendService::new(repository, validator, bypass_token.clone()); let response = (StatusCode::TOO_MANY_REQUESTS, "Too many requests"); + let bypasser = bypass_token.map_or_else( + || RateLimitBypasser::Never(NeverBypasser), + |x| RateLimitBypasser::Token(AuthTokenBypasser(x.into())), + ); + // Use ic-bn-lib rate limiting middleware, with key by IP address. let create_rate_limiter = |limit: Option, response| { option_layer( - limit.map(|lim| layer_by_ip(lim, 2 * lim, response, bypass_token.clone()).unwrap()), + limit.map(|lim| layer_by_ip(lim, 2 * lim, response, bypasser.clone()).unwrap()), ) }; diff --git a/ic-bn-lib/src/custom_domains/base/traits/validation.rs b/ic-bn-lib/src/custom_domains/base/traits/validation.rs index 5c0ac11..81703c4 100644 --- a/ic-bn-lib/src/custom_domains/base/traits/validation.rs +++ b/ic-bn-lib/src/custom_domains/base/traits/validation.rs @@ -47,7 +47,9 @@ pub trait ValidatesDomains: Send + Sync { /// Validates that a domain can be registered or updated. /// - /// Skips certain checks compared to validate() + /// Skips certain checks compared to `validate()`: + /// * Canister ownership verification (.well-known/ic-domains) + /// * DNS TXT record verification for canister ID async fn validate_limited(&self, domain: &FQDN) -> Result<(), ValidationError>; /// Validates that a domain can be safely deleted. diff --git a/ic-bn-lib/src/custom_domains/base/types/worker.rs b/ic-bn-lib/src/custom_domains/base/types/worker.rs index bcb8a0f..b78df67 100644 --- a/ic-bn-lib/src/custom_domains/base/types/worker.rs +++ b/ic-bn-lib/src/custom_domains/base/types/worker.rs @@ -11,7 +11,7 @@ use async_trait::async_trait; use candid::Principal; use chrono::{DateTime, Utc}; use derive_new::new; -use fqdn::FQDN; +use fqdn::{FQDN, Fqdn}; use instant_acme::{RevocationReason, RevocationRequest}; use pem::parse_many; use prometheus::{ @@ -624,6 +624,20 @@ impl Run for Worker { } } +fn task_validation_failure( + domain: &Fqdn, + error: ValidationError, + task_id: UtcTimestamp, + task_kind: TaskKind, +) -> TaskResult { + TaskResult::failure( + domain.into(), + TaskFailReason::ValidationFailed(error.to_string()), + task_id, + task_kind, + ) +} + async fn issue_task( domain: FQDN, validator: Arc, @@ -633,20 +647,11 @@ async fn issue_task( wildcard: bool, canister_id: Option, ) -> TaskResult { - let failure = |e: ValidationError| { - TaskResult::failure( - domain.clone(), - TaskFailReason::ValidationFailed(e.to_string()), - task_id, - task_kind, - ) - }; - // If a canister ID is provided, validate it with limited checks. // Otherwise, perform full validation to derive the canister ID from DNS. let canister_id = if let Some(v) = canister_id { if let Err(e) = validator.validate_limited(&domain).await { - return failure(e); + return task_validation_failure(&domain, e, task_id, task_kind); }; v @@ -654,7 +659,7 @@ async fn issue_task( match validator.validate(&domain).await { Ok(v) => v, Err(e) => { - return failure(e); + return task_validation_failure(&domain, e, task_id, task_kind); } } }; @@ -754,20 +759,11 @@ async fn update_task( task_kind: TaskKind, canister_id: Option, ) -> TaskResult { - let failure = |e: ValidationError| { - TaskResult::failure( - domain.clone(), - TaskFailReason::ValidationFailed(e.to_string()), - task_id, - task_kind, - ) - }; - // If a canister ID is provided, validate it with limited checks. // Otherwise, perform full validation to derive the canister ID from DNS. let canister_id = if let Some(v) = canister_id { if let Err(e) = validator.validate_limited(&domain).await { - return failure(e); + return task_validation_failure(&domain, e, task_id, task_kind); }; v @@ -775,7 +771,7 @@ async fn update_task( match validator.validate(&domain).await { Ok(v) => v, Err(e) => { - return failure(e); + return task_validation_failure(&domain, e, task_id, task_kind); } } }; diff --git a/ic-bn-lib/src/http/middleware/rate_limiter.rs b/ic-bn-lib/src/http/middleware/rate_limiter.rs index 2a23e73..13d4de0 100644 --- a/ic-bn-lib/src/http/middleware/rate_limiter.rs +++ b/ic-bn-lib/src/http/middleware/rate_limiter.rs @@ -20,15 +20,12 @@ use governor::{ nanos::Nanos, state::keyed::DashMapStateStore, }; -use http::{ - HeaderName, HeaderValue, StatusCode, - header::{AUTHORIZATION, RETRY_AFTER}, -}; +use http::{HeaderName, HeaderValue, StatusCode, header::RETRY_AFTER}; use tower::{Layer, Service}; use crate::{constant_time_eq, hname, http::middleware::RemoteAddr}; -const BYPASS_TOKEN_HEADER: HeaderName = hname!("x-ratelimit-bypass-token"); +pub const BYPASS_TOKEN_HEADER: HeaderName = hname!("x-ratelimit-bypass-token"); /// The `governor` rate limiter type that backs this middleware, generic over the clock so /// that tests can inject `governor::clock::FakeRelativeClock` instead of the real-time @@ -50,6 +47,11 @@ pub trait KeyExtractor: Clone + Send + Sync + 'static { fn extract(&self, req: &Request) -> Option; } +/// Decides if the rate-limiting for the given request should be bypassed +pub trait Bypasser: Clone + Send + Sync + 'static { + fn should_bypass(&self, req: &Request) -> bool; +} + /// Extracts an IP from the request as a rate-limiting key #[derive(Clone)] pub struct IpKeyExtractor; @@ -74,29 +76,60 @@ impl KeyExtractor for GlobalKeyExtractor { } } -struct RateLimiterState { +/// Bypasser implementation that checks for a token in the request header `x-ratelimit-bypass-token` +#[derive(Clone)] +pub struct TokenBypasser(Bytes); + +impl TokenBypasser { + pub fn new(token: impl Into) -> Self { + Self(token.into()) + } +} + +impl Bypasser for TokenBypasser { + fn should_bypass(&self, req: &Request) -> bool { + req.headers() + .get(BYPASS_TOKEN_HEADER) + .map(|x| x.as_bytes()) + .is_some_and(|x| constant_time_eq(x, &self.0)) + } +} + +/// Default Bypasser implementation that never bypasses any requests +#[derive(Clone)] +pub struct NeverBypasser; + +impl Bypasser for NeverBypasser { + fn should_bypass(&self, _req: &Request) -> bool { + false + } +} + +struct RateLimiterState { key_extractor: K, limiter: GovRateLimiter, rate_limited_response: R, - bypass_token: Option, + bypasser: BP, last_cleanup: ArcSwap, } /// Ratelimiter that implements Tower Service #[derive(Clone)] -pub struct RateLimiter { - state: Arc>, +pub struct RateLimiter +{ + state: Arc>, inner: S, } /// Implement Tower Service for RateLimiter -impl Service for RateLimiter +impl Service for RateLimiter where S: Service + Send + 'static, S::Future: Send + 'static, S::Error: Send + 'static, K: KeyExtractor, R: IntoResponse + Clone + Send + Sync + 'static, + BP: Bypasser, C: Clock + Send + Sync + 'static, { type Response = S::Response; @@ -111,21 +144,8 @@ where /// Stale entries cleanup interval - 5 minutes const CLEANUP_INTERVAL: Nanos = Nanos::new(300_000_000_000); - // Check that bypass token is configured, header was sent and it matches. - // Checks both the custom header and the Authorization. - let bypass = request - .headers() - .get(BYPASS_TOKEN_HEADER) - .map(|x| x.as_bytes()) - .or_else(|| { - request - .headers() - .get(AUTHORIZATION) - .and_then(|v| v.to_str().ok()) - .map(|s| s.trim_start_matches("Bearer ").as_bytes()) - }) - .zip(self.state.bypass_token.as_ref()) - .is_some_and(|(hdr, token)| constant_time_eq(hdr, token.as_bytes())); + // Check that bypasser is configured and tells us to bypass + let bypass = self.state.bypasser.should_bypass(&request); // Clean up stale entries from time to time let now = self.state.limiter.clock().now(); @@ -171,18 +191,24 @@ where /// Layer usable as an Axum middleware #[derive(Clone, derive_new::new)] -pub struct RateLimiterLayer { - state: Arc>, +pub struct RateLimiterLayer< + K: KeyExtractor, + R, + BP: Bypasser = NeverBypasser, + C: Clock = DefaultClock, +> { + state: Arc>, } -impl Layer for RateLimiterLayer +impl Layer for RateLimiterLayer where S: Clone, K: KeyExtractor, R: IntoResponse + Clone + Send + Sync + 'static, + BP: Bypasser, C: Clock + Send + Sync + 'static, { - type Service = RateLimiter; + type Service = RateLimiter; fn layer(&self, inner: S) -> Self::Service { RateLimiter { @@ -193,51 +219,51 @@ where } /// Create unkeyed rate-limiter -pub fn layer_global( +pub fn layer_global( rps: u32, burst_size: u32, rate_limited_response: R, - bypass_token: Option, -) -> Result, Error> { + bypasser: BP, +) -> Result, Error> { layer( rps, burst_size, GlobalKeyExtractor, rate_limited_response, - bypass_token, + bypasser, ) } /// Create ratelimiter keyed by IP -pub fn layer_by_ip( +pub fn layer_by_ip( rps: u32, burst_size: u32, rate_limited_response: R, - bypass_token: Option, -) -> Result, Error> { + bypasser: BP, +) -> Result, Error> { layer( rps, burst_size, IpKeyExtractor, rate_limited_response, - bypass_token, + bypasser, ) } /// Create a ratelimiter with a provided key extractor -pub fn layer( +pub fn layer( rps: u32, burst_size: u32, key_extractor: K, rate_limited_response: R, - bypass_token: Option, -) -> Result, Error> { + bypasser: BP, +) -> Result, Error> { layer_with_clock( rps, burst_size, key_extractor, rate_limited_response, - bypass_token, + bypasser, DefaultClock::default(), ) } @@ -245,14 +271,19 @@ pub fn layer( /// Create a ratelimiter with a provided key extractor and clock. This custom clock is there so /// that tests can supply a `FakeRelativeClock` and drive the rate limiter's /// time deterministically, without depending on real wall-clock delays. -fn layer_with_clock( +fn layer_with_clock< + K: KeyExtractor, + R: IntoResponse + Clone + Send + Sync + 'static, + BP: Bypasser, + C: Clock, +>( rps: u32, burst_size: u32, key_extractor: K, rate_limited_response: R, - bypass_token: Option, + bypasser: BP, clock: C, -) -> Result, Error> { +) -> Result, Error> { let period = Duration::from_secs(1) .checked_div(rps) .ok_or_else(|| anyhow!("RPS is zero"))?; @@ -271,7 +302,7 @@ fn layer_with_clock