diff --git a/Cargo.lock b/Cargo.lock index da1b6d5..9b46dae 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2672,9 +2672,9 @@ dependencies = [ [[package]] name = "ic-bn-lib" -version = "0.2.2" +version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eea2c1f3aad0a4402aecc624cda39efbafc7dd6645c7f8ab68fe0b99620dcf02" +checksum = "82b26684028978088a9f670c1fe371c32807728743ca0eaff06acffb2abe2e3f" dependencies = [ "ahash", "anyhow", @@ -2763,9 +2763,9 @@ dependencies = [ [[package]] name = "ic-bn-lib-common" -version = "0.2.2" +version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9ddfc3775a8bd5aff225bc20aaadb4da46fd9087363c7c510bb264fa6a1f480f" +checksum = "a432e07d35f3ab744b55cb39743a592e1115e9ee16d95069d3d54eea3de8bb3d" dependencies = [ "anyhow", "async-trait", @@ -2942,9 +2942,9 @@ dependencies = [ [[package]] name = "ic-custom-domains-backend" -version = "0.2.4" +version = "0.2.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "962142d708266188d2be60dd6a6e27704c3c4363ac4c4a3af843b310a70e3731" +checksum = "86c4dcfc22f1ad7aeb4d4e7bfb1ce9c74eb428da74739cd93a9d7a088e41454b" dependencies = [ "anyhow", "axum", @@ -2974,9 +2974,9 @@ dependencies = [ [[package]] name = "ic-custom-domains-base" -version = "0.2.4" +version = "0.2.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "74464dd81ad40870fa648b41a771f0f6f4ce4b2f824fe9194fe6857d293d8ef2" +checksum = "d817d51182c9174a39fa61b8f5ba30f551c837a5c751aa0a444b810257da7817" dependencies = [ "anyhow", "async-trait", @@ -3005,9 +3005,9 @@ dependencies = [ [[package]] name = "ic-custom-domains-canister-api" -version = "0.2.4" +version = "0.2.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "268df6565445943a8cf4bd1c2d6d86c33c65e23b32e7d3f155a7be8d90a07166" +checksum = "46e74a4ddc1d57e789d8f46e0c2e79e2fdffcd38a75029e3287129bcb2db260b" dependencies = [ "candid", "derive-new", @@ -3018,9 +3018,9 @@ dependencies = [ [[package]] name = "ic-custom-domains-canister-client" -version = "0.2.4" +version = "0.2.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "52850559b366e2095d912b5f57a97cd0288c422e6e24f84a5514988ba2bd84ff" +checksum = "7aa55f92edfedb6bef26fe415ec2b1ab7742cb260195475d6c822443a6963c99" dependencies = [ "anyhow", "arc-swap", diff --git a/Cargo.toml b/Cargo.toml index 26b1094..25e9878 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -35,15 +35,15 @@ http = "1.3.1" http-body = "1.0.1" http-body-util = "0.1.2" humantime = "2.2.0" -ic-bn-lib = { version = "0.2.2", features = [ +ic-bn-lib = { version = "0.2.4", features = [ "acme", "vector", "cert-providers", "clients-hyper", ] } -ic-bn-lib-common = "0.2.2" -ic-custom-domains-backend = "0.2.4" -ic-custom-domains-base = "0.2.4" +ic-bn-lib-common = "0.2.4" +ic-custom-domains-backend = "0.2.6" +ic-custom-domains-base = "0.2.6" ic-http-certification = { version = "3.1.0", optional = true } ic-transport-types = "0.47" ic-http-gateway-protocol = { package = "ic-http-gateway-protocol", git = "https://github.com/dfinity/ic-http-gateway-protocol", tag = "v0.5.1" } diff --git a/benches/domain_lookup.rs b/benches/domain_lookup.rs index 971daff..51ee91f 100644 --- a/benches/domain_lookup.rs +++ b/benches/domain_lookup.rs @@ -32,11 +32,7 @@ fn criterion_benchmark(c: &mut Criterion) { let custom_domains = domains .clone() .into_iter() - .map(|x| CustomDomain { - name: x, - canister_id: principal!("aaaaa-aa"), - timestamp: 0, - }) + .map(|x| CustomDomain::new(x, principal!("aaaaa-aa"))) .collect::>(); let s = CustomDomainStorage::new( diff --git a/src/cli.rs b/src/cli.rs index 9809806..11c0b5c 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -27,7 +27,10 @@ use reqwest::Url; use crate::{ core::{AUTHOR_NAME, SERVICE_NAME}, - routing::{RequestType, domain::CanisterAlias}, + routing::{ + RequestType, + domain::{CanisterAlias, CustomDomainHttpProvider}, + }, }; /// Clap does not support prefixes due to macro limitations. @@ -295,19 +298,19 @@ pub struct Domain { /// List of generic custom domain provider URLs. /// Expects a JSON object in form '{"domain.bar": "aaaaa-aa"}' in response to a GET request. #[clap(env, long, value_delimiter = ',')] - pub domain_custom_provider: Vec, + pub domain_custom_provider: Vec, /// List of generic timestamped custom domain provider URLs. /// Expects a JSON object in form '{"timestamp": 1234, "url": "https://foo/bar"}' in response to a GET request. /// When the timestamp changes - the provider gets the list of domains from the URL provided in response. /// The JSON format there should be the same as for the normal generic provider (see above). #[clap(env, long, value_delimiter = ',')] - pub domain_custom_provider_timestamped: Vec, + pub domain_custom_provider_timestamped: Vec, /// List of generic differential custom domain provider URLs. /// It first downloads the full seed and then only applies incremental updates to it using a timestamp. #[clap(env, long, value_delimiter = ',')] - pub domain_custom_provider_diff: Vec, + pub domain_custom_provider_diff: Vec, /// How frequently to poll custom domain providers for updates #[clap(env, long, default_value = "30s", value_parser = parse_duration)] diff --git a/src/core.rs b/src/core.rs index 2dcdce9..5959344 100644 --- a/src/core.rs +++ b/src/core.rs @@ -226,8 +226,10 @@ pub async fn main( Arc::new(custom_domains::GenericProvider::new( http_client.clone(), - x.clone(), + x.url.clone(), cli.domain.domain_custom_provider_timeout, + x.priority, + x.flags, )) as Arc })); @@ -240,8 +242,10 @@ pub async fn main( Arc::new(custom_domains::GenericProviderTimestamped::new( http_client.clone(), - x.clone(), + x.url.clone(), cli.domain.domain_custom_provider_timeout, + x.priority, + x.flags, )) as Arc }), ); @@ -251,8 +255,11 @@ pub async fn main( Arc::new(custom_domains::GenericProviderDiff::new( http_client.clone(), - x.clone(), + x.url.clone(), cli.domain.domain_custom_provider_timeout, + x.priority, + x.flags, + Duration::from_mins(5), )) as Arc })); @@ -260,10 +267,11 @@ pub async fn main( if let Some(path) = &cli.domain.domain_custom_provider_local_file { warn!("Adding local file custom domain provider: {path}"); - custom_domain_providers.push( - Arc::new(custom_domains::LocalFileProvider::new(path.into())) - as Arc, - ); + custom_domain_providers.push(Arc::new(custom_domains::LocalFileProvider::new( + path.into(), + 0, + None, + )) as Arc); } // Create IC Agent for use by RoutingTableManager / SMTP diff --git a/src/routing/domain.rs b/src/routing/domain.rs index 02aadc9..56b9d8f 100644 --- a/src/routing/domain.rs +++ b/src/routing/domain.rs @@ -1,7 +1,9 @@ use std::{ collections::BTreeMap, + fmt::Display, str::FromStr, - sync::{Arc, RwLock}, + sync::{Arc, LazyLock, RwLock}, + time::Instant, }; use anyhow::{Context, Error, anyhow}; @@ -12,14 +14,71 @@ use fqdn::{FQDN, Fqdn, fqdn}; use ic_bn_lib::custom_domains::LooksUpCustomDomain; use ic_bn_lib_common::{ traits::{Healthy, Run, custom_domains::ProvidesCustomDomains}, - types::CustomDomain, + types::{CustomDomain, DomainFlags}, }; use prometheus::{ IntCounter, IntGauge, Registry, register_int_counter_with_registry, register_int_gauge_with_registry, }; +use regex::Regex; use tokio_util::sync::CancellationToken; use tracing::{debug, warn}; +use url::Url; + +static PROVIDER_RE: LazyLock = LazyLock::new(|| { + Regex::new(r"^(?[^|]+?)(?:\|(?[^<]*))?(?:<(?\d+)>)?$").unwrap() +}); + +/// Custom domain provider URL with optional flags and priority. +/// +/// It is parsed from the following template by Clap: `https://foo.bar/path/to?foo=a|flag1|flag2<0>` +/// The delimiters chosen here (|<>) are forbidden in the URL by the RFC. +#[derive(Debug, Clone, Eq, PartialEq)] +pub struct CustomDomainHttpProvider { + pub url: Url, + pub priority: u8, + pub flags: Option, +} + +impl Display for CustomDomainHttpProvider { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{} (prio {})", self.url, self.priority)?; + if let Some(v) = self.flags { + write!(f, " (flags: {v})")?; + } + + Ok(()) + } +} + +impl FromStr for CustomDomainHttpProvider { + type Err = Error; + + fn from_str(s: &str) -> Result { + let caps = PROVIDER_RE + .captures(s.trim()) + .ok_or_else(|| anyhow!("invalid provider format: {s}"))?; + + Ok(Self { + url: Url::parse(&caps["url"]).context("unable to parse URL")?, + + flags: caps + .name("flags") + .map(|m| m.as_str().trim()) + .filter(|f| !f.is_empty()) + .map(DomainFlags::from_str) + .transpose() + .context("unable to parse flags")?, + + priority: caps + .name("prio") + .map(|m| m.as_str().parse()) + .transpose() + .context("unable to parse priority as integer")? + .unwrap_or(0), + }) + } +} /// Domain entity with certain metadata #[derive(Clone, Debug, Eq, PartialEq)] @@ -40,6 +99,8 @@ pub struct DomainLookup { pub canister_id: Option, pub timestamp: u64, pub verify: bool, + pub priority: u8, + pub flags: Option, } /// Resolves hostname to a canister id @@ -86,6 +147,7 @@ pub struct CustomDomainStorage { snapshot: RwLock>>>, metric_count: IntGauge, metric_dupes: IntGauge, + metric_dupes_overridden: IntGauge, metric_failures: IntCounter, } @@ -118,6 +180,15 @@ impl CustomDomainStorage { ) .unwrap(); + let metric_dupes_overridden = register_int_gauge_with_registry!( + format!("custom_domains_dupes_overridden"), + format!( + "Number of duplicates among custom domains that were overridden (higher prio/ts)" + ), + registry + ) + .unwrap(); + let metric_failures = register_int_counter_with_registry!( format!("custom_domains_failures_total"), format!("Total number of fetch failures"), @@ -131,6 +202,7 @@ impl CustomDomainStorage { providers, metric_count, metric_dupes, + metric_dupes_overridden, metric_failures, } } @@ -159,7 +231,10 @@ impl CustomDomainStorage { snapshot } + #[allow(clippy::cast_possible_wrap)] pub async fn refresh(&self) { + let start = Instant::now(); + let snapshot_old = self.snapshot.read().unwrap().clone(); let snapshot = self.fetch(snapshot_old.clone()).await; @@ -176,42 +251,55 @@ impl CustomDomainStorage { // Convert the snapshot into new lookup structure let domains = snapshot.into_iter().flatten().flatten(); - let mut dupes = 0; + let mut dupes = 0i64; + let mut dupes_overridden = 0i64; - for d in domains { - // Do not add new domain if the same one exists with newer timestamp - if let Some(v) = tree.get(&d.name) { + for new in domains { + // Check if we have the same domain already + if let Some(exists) = tree.get(&new.name) { dupes += 1; - if v.timestamp > d.timestamp { + // Skip if existing prio is higher + if exists.priority > new.priority { continue; } + + // If prio is the same - compare timestamps + if exists.priority == new.priority && exists.timestamp > new.timestamp { + continue; + } + + // Otherwise override + dupes_overridden += 1; } let dl = DomainLookup { domain: Domain { - name: d.name.clone(), + name: new.name.clone(), custom: true, http: true, api: true, }, - timestamp: d.timestamp, - canister_id: Some(d.canister_id), + timestamp: new.timestamp, + canister_id: Some(new.canister_id), verify: true, + priority: new.priority, + flags: new.flags, }; - tree.insert(d.name, dl); + tree.insert(new.name, dl); } warn!( - "{self:?}: got new set of domains: {} ({dupes} duplicates)", + "{self:?}: got new set of domains in {}s: {} ({dupes} dupes, {dupes_overridden} of them overridden)", + start.elapsed().as_secs(), tree.len() ); // Set metrics - #[allow(clippy::cast_possible_wrap)] self.metric_count.set(tree.len() as i64); - self.metric_dupes.set(i64::from(dupes)); + self.metric_dupes.set(dupes); + self.metric_dupes_overridden.set(dupes_overridden); // Store it let inner = CustomDomainStorageInner(tree); @@ -291,6 +379,8 @@ impl DomainResolver { canister_id: Some(alias.1), timestamp: 0, verify: true, + priority: 0, + flags: None, }, ) }) @@ -309,6 +399,8 @@ impl DomainResolver { canister_id: None, timestamp: 0, verify: true, + priority: 0, + flags: None, }, ) }) @@ -324,7 +416,7 @@ impl DomainResolver { // Tries to find the base domain that corresponds to the given host and resolve a canister id fn resolve_domain(&self, host: &Fqdn) -> Option { - // First try to find an exact match + // First try to find an exact match. // This covers base domains and their aliases, plus API domains if let Some(v) = self.domains_all.get(host) { return Some(v.clone()); @@ -380,6 +472,8 @@ impl DomainResolver { canister_id, timestamp: 0, verify: !raw, + priority: 0, + flags: None, }) } } @@ -395,7 +489,10 @@ impl ResolvesDomain for DomainResolver { #[cfg(test)] mod test { use fqdn::fqdn; - use ic_bn_lib_common::principal; + use ic_bn_lib_common::{ + principal, + types::{FLAG_PRERENDER, FLAG_TEST}, + }; use super::*; @@ -466,25 +563,47 @@ mod test { let domains_base = vec![fqdn!("ic0.app"), fqdn!("icp0.io")]; let domains_api = vec![fqdn!("icp-api.io")]; let custom_domain_provider = TestCustomDomainProvider(vec![ + CustomDomain { + name: fqdn!("foo.xyz"), + canister_id: principal!(TEST_CANISTER_ID), + timestamp: 10, + priority: 0, + flags: Some(DomainFlags::new([FLAG_TEST])), + }, + CustomDomain { + name: fqdn!("foo.xyz"), + canister_id: principal!(TEST_CANISTER_ID), + timestamp: 0, + priority: 1, + flags: Some(DomainFlags::new([FLAG_PRERENDER])), + }, CustomDomain { name: fqdn!("foo.bar"), - timestamp: 30, canister_id: principal!(TEST_CANISTER_ID), + timestamp: 30, + priority: 0, + flags: None, }, CustomDomain { name: fqdn!("foo.bar"), - timestamp: 20, canister_id: principal!(TEST_CANISTER_ID_2), + timestamp: 20, + priority: 0, + flags: Some(DomainFlags::new([FLAG_PRERENDER])), }, CustomDomain { name: fqdn!("foo.baz"), - timestamp: 10, canister_id: principal!(TEST_CANISTER_ID), + timestamp: 10, + priority: 0, + flags: None, }, CustomDomain { name: fqdn!("foo.baz"), - timestamp: 20, canister_id: principal!(TEST_CANISTER_ID_3), + timestamp: 20, + priority: 0, + flags: None, }, ]); @@ -501,13 +620,13 @@ mod test { // Verify metrics are tracked correctly assert_eq!( custom_domain_storage.metric_count.get(), - 2, - "should have 2 domains after deduplication" + 3, + "should have 3 domains after deduplication" ); assert_eq!( custom_domain_storage.metric_dupes.get(), - 2, - "should have 2 duplicates (foo.bar and foo.baz each appear twice)" + 3, + "should have 3 duplicates" ); assert_eq!( custom_domain_storage.metric_failures.get(), @@ -539,6 +658,8 @@ mod test { timestamp: 0, canister_id: Some(a.1), verify: true, + priority: 0, + flags: None, }) ); } @@ -573,6 +694,8 @@ mod test { canister_id: None, timestamp: 0, verify: true, + priority: 0, + flags: None, }) ); @@ -584,6 +707,8 @@ mod test { canister_id: None, timestamp: 0, verify: true, + priority: 0, + flags: None, }) ); @@ -598,6 +723,8 @@ mod test { canister_id: None, timestamp: 0, verify: true, + priority: 0, + flags: None, }) ); @@ -609,6 +736,8 @@ mod test { canister_id: Some(canister_id), timestamp: 0, verify: true, + priority: 0, + flags: None, }) ); @@ -620,6 +749,8 @@ mod test { canister_id: Some(canister_id), timestamp: 0, verify: true, + priority: 0, + flags: None, }) ); @@ -631,6 +762,8 @@ mod test { canister_id: Some(canister_id), timestamp: 0, verify: false, + priority: 0, + flags: None, }) ); assert_eq!( @@ -640,6 +773,8 @@ mod test { canister_id: Some(canister_id), timestamp: 0, verify: false, + priority: 0, + flags: None, }) ); @@ -651,6 +786,8 @@ mod test { canister_id: None, timestamp: 0, verify: true, + priority: 0, + flags: None, }) ); @@ -662,6 +799,8 @@ mod test { canister_id: Some(canister_id), timestamp: 0, verify: true, + priority: 0, + flags: None, }) ); assert_eq!( @@ -671,6 +810,8 @@ mod test { canister_id: Some(canister_id), timestamp: 0, verify: true, + priority: 0, + flags: None, }) ); @@ -679,7 +820,7 @@ mod test { assert_eq!(resolver.resolve(&fqdn!("aaaaa-aa.foo.icp0.io")), None,); // Resolve custom domains - // Make sure that newer custom domains are used (with higher timestamp) + // Make sure that domain with higher timestamp are used assert_eq!( resolver.resolve(&fqdn!("foo.bar")), Some(DomainLookup { @@ -692,8 +833,28 @@ mod test { canister_id: Some(principal!(TEST_CANISTER_ID)), timestamp: 30, verify: true, + priority: 0, + flags: None, + }) + ); + // Make sure that the domain with higher priority is used even when the timestamp of the other is higher + assert_eq!( + resolver.resolve(&fqdn!("foo.xyz")), + Some(DomainLookup { + domain: Domain { + name: fqdn!("foo.xyz"), + http: true, + api: true, + custom: true, + }, + canister_id: Some(principal!(TEST_CANISTER_ID)), + timestamp: 0, + verify: true, + priority: 1, + flags: Some(DomainFlags::new([FLAG_PRERENDER])), }) ); + // Equal priority, higher timestamp assert_eq!( resolver.resolve(&fqdn!("foo.baz")), Some(DomainLookup { @@ -706,6 +867,8 @@ mod test { canister_id: Some(principal!(TEST_CANISTER_ID_3)), timestamp: 20, verify: true, + priority: 0, + flags: None, }) ); @@ -771,4 +934,60 @@ mod test { Some(Principal::from_text("aaaaa-aa").unwrap()) ); } + + #[test] + fn test_custom_domain_provider_flags() { + // with prio + assert_eq!( + CustomDomainHttpProvider::from_str("http://foo/bar|prerender<66>").unwrap(), + CustomDomainHttpProvider { + url: "http://foo/bar".parse().unwrap(), + priority: 66, + flags: Some(DomainFlags::new([FLAG_PRERENDER])), + } + ); + // with prio, empty flags + assert_eq!( + CustomDomainHttpProvider::from_str("http://foo/bar|<66>").unwrap(), + CustomDomainHttpProvider { + url: "http://foo/bar".parse().unwrap(), + priority: 66, + flags: None, + } + ); + // with prio, no flags + assert_eq!( + CustomDomainHttpProvider::from_str("http://foo/bar<66>").unwrap(), + CustomDomainHttpProvider { + url: "http://foo/bar".parse().unwrap(), + priority: 66, + flags: None, + } + ); + + // no prio + assert_eq!( + CustomDomainHttpProvider::from_str("http://foo/bar|prerender").unwrap(), + CustomDomainHttpProvider { + url: "http://foo/bar".parse().unwrap(), + priority: 0, + flags: Some(DomainFlags::new([FLAG_PRERENDER])), + } + ); + + // just url + assert_eq!( + CustomDomainHttpProvider::from_str("http://foo/bar").unwrap(), + CustomDomainHttpProvider { + url: "http://foo/bar".parse().unwrap(), + priority: 0, + flags: None, + } + ); + + // error cases + assert!(CustomDomainHttpProvider::from_str("http://foo/bar|prerender").is_err()); + assert!(CustomDomainHttpProvider::from_str("http://foo/bar|blah").is_err()); + assert!(CustomDomainHttpProvider::from_str("|||foo/bar|blah").is_err()); + } } diff --git a/src/routing/middleware/prerender.rs b/src/routing/middleware/prerender.rs index 2e533ff..12d3817 100644 --- a/src/routing/middleware/prerender.rs +++ b/src/routing/middleware/prerender.rs @@ -16,7 +16,10 @@ use http::{ }; use http_body_util::Full; use ic_bn_lib::{hname, hval}; -use ic_bn_lib_common::traits::http::ClientHttp; +use ic_bn_lib_common::{ + traits::http::ClientHttp, + types::{DomainFlags, FLAG_PRERENDER}, +}; use tokio::time::timeout; use tracing::info; @@ -28,6 +31,7 @@ use crate::routing::{ const HEADER_SECRET: HeaderName = hname!("x-worker-secret"); const HEADER_X_PRE_RENDERED: HeaderName = hname!("x-pre-rendered"); +const VALUE_1: HeaderValue = hval!("1"); const HEADERS_TO_REMOVE: [HeaderName; 9] = [ CONTENT_ENCODING, @@ -99,24 +103,29 @@ pub struct PrerenderState { impl PrerenderState { /// Whether we should prerender this request or pass it through - fn should_render(&self, authority: &Fqdn, uri: &Uri, method: &Method, is_bot: bool) -> bool { - if method != Method::GET { + fn should_render(&self, uri: &Uri, request: &Request) -> bool { + if request.method() != Method::GET { return false; } - // Check that authority is a subdomain of any of configured domains. - // It matches only if it's max 1 level deeper than a subdomain. - // E.g. if the configured domain is `bar` then `foo.bar` will match, - // while `baz.foo.bar` will not. - if !self - .domains - .iter() - .any(|x| authority.depth() <= (x.depth() + 1) && authority.is_subdomain_of(x)) - { + if is_static_asset(uri.path()) { return false; } - if authority + let Some(is_bot) = request.extensions().get::() else { + return false; + }; + + if !is_bot.0 { + return false; + } + + let Some(ctx) = request.extensions().get::>() else { + return false; + }; + + if ctx + .authority .labels() .next() .is_some_and(|x| x.contains("-draft")) @@ -124,11 +133,22 @@ impl PrerenderState { return false; } - if !is_bot { - return false; + // Check if there's a pre-render flag present + if let Some(v) = request.extensions().get::() + && v.has_flag(FLAG_PRERENDER) + { + return true; } - if is_static_asset(uri.path()) { + // Check that authority is a subdomain of any of configured domains. + // It matches only if it's max 1 level deeper than a subdomain. + // E.g. if the configured domain is `bar` then `foo.bar` will match, + // while `baz.foo.bar` will not. + if !self + .domains + .iter() + .any(|x| ctx.authority.depth() <= (x.depth() + 1) && ctx.authority.is_subdomain_of(x)) + { return false; } @@ -176,7 +196,7 @@ impl PrerenderState { } // Add marker to show that it was pre-rendered - v.headers_mut().insert(HEADER_X_PRE_RENDERED, hval!("1")); + v.headers_mut().insert(HEADER_X_PRE_RENDERED, VALUE_1); v } @@ -197,12 +217,11 @@ impl PrerenderState { pub async fn middleware( State(state): State>, Extension(ctx): Extension>, - Extension(is_bot): Extension, OriginalUri(uri): OriginalUri, request: Request, next: Next, ) -> Result { - if !state.should_render(&ctx.authority, &uri, request.method(), is_bot.0) { + if !state.should_render(&uri, &request) { return Ok(next.run(request).await); } @@ -230,7 +249,8 @@ mod tests { use crate::{ routing::{ - domain::{CustomDomainStorage, DomainResolver}, + RequestType, + domain::{CustomDomainStorage, Domain, DomainResolver}, middleware::{ is_bot::{self, IsBotState}, validate::{self, ValidateState}, @@ -239,9 +259,12 @@ mod tests { test::{FakeDomainProvider, TestClient}, }; use async_trait::async_trait; - use axum::{Router, body::Body, middleware::from_fn_with_state, response::IntoResponse}; + use axum::{ + Router, body::Body, extract::Request, middleware::from_fn_with_state, + response::IntoResponse, + }; use fqdn::fqdn; - use http::{Request, StatusCode}; + use http::StatusCode; use http_body_util::BodyExt; use ic_bn_lib::hval; use ic_bn_lib_common::types::http::Error as HttpError; @@ -271,6 +294,41 @@ mod tests { } } + fn make_request(is_bot: bool, authority: FQDN, method: Method) -> Request { + let ctx = RequestCtx { + authority: authority.clone(), + domain: Domain { + name: authority, + custom: true, + http: true, + api: false, + }, + verify: false, + request_type: RequestType::Http, + }; + + let mut r = Request::new(Body::empty()); + *r.method_mut() = method; + r.extensions_mut().insert(Arc::new(ctx)); + r.extensions_mut().insert(IsBot(is_bot)); + + r + } + + fn make_request_with_flags( + is_bot: bool, + authority: FQDN, + method: Method, + flags: Option, + ) -> Request { + let mut req = make_request(is_bot, authority, method); + if let Some(v) = flags { + req.extensions_mut().insert(v); + } + + req + } + #[test] fn test_should_render() { let state = PrerenderState::new( @@ -283,83 +341,83 @@ mod tests { // ok assert!(state.should_render( - &fqdn!("caffeine.xyz"), &Uri::from_static("http://foo/script.php"), - &Method::GET, - true + &make_request(true, fqdn!("caffeine.xyz"), Method::GET) )); assert!(state.should_render( - &fqdn!("foo.caffeine.xyz"), &Uri::from_static("http://foo/script.php"), - &Method::GET, - true + &make_request(true, fqdn!("foo.caffeine.xyz"), Method::GET) )); assert!(state.should_render( - &fqdn!("bar.caffeine.abc"), &Uri::from_static("http://foo/script.php"), - &Method::GET, - true + &make_request(true, fqdn!("bar.caffeine.xyz"), Method::GET) )); // wrong domain assert!(!state.should_render( - &fqdn!("foo.caffeine.foo"), &Uri::from_static("http://foo/script.php"), - &Method::GET, - true + &make_request(true, fqdn!("foo.caffeine.foo"), Method::GET) )); assert!(!state.should_render( - &fqdn!("xyz"), &Uri::from_static("http://foo/script.php"), - &Method::GET, - true + &make_request(true, fqdn!("xyz"), Method::GET) )); // too deep domain assert!(!state.should_render( - &fqdn!("bar.foo.caffeine.xyz"), &Uri::from_static("http://foo/script.php"), - &Method::GET, - true + &make_request(true, fqdn!("bar.foo.caffeine.xyz"), Method::GET) )); // static asset assert!(!state.should_render( - &fqdn!("foo.caffeine.xyz"), &Uri::from_static("http://foo/logo.png"), - &Method::GET, - true + &make_request(true, fqdn!("foo.caffeine.xyz"), Method::GET) )); // not bot assert!(!state.should_render( - &fqdn!("foo.caffeine.xyz"), &Uri::from_static("http://foo/script.php"), - &Method::GET, - false + &make_request(false, fqdn!("foo.caffeine.xyz"), Method::GET) )); // wrong method assert!(!state.should_render( - &fqdn!("foo.caffeine.xyz"), &Uri::from_static("http://foo/script.php"), - &Method::POST, - false + &make_request(true, fqdn!("foo.caffeine.xyz"), Method::POST) )); // draft assert!(!state.should_render( - &fqdn!("foo-draft.caffeine.xyz"), &Uri::from_static("http://foo/script.php"), - &Method::GET, - true + &make_request(true, fqdn!("foo-draft.caffeine.xyz"), Method::GET) )); assert!(!state.should_render( - &fqdn!("foo-draft-foo.caffeine.xyz"), &Uri::from_static("http://foo/script.php"), - &Method::GET, - true + &make_request(true, fqdn!("foo-draft-foo.caffeine.xyz"), Method::GET) + )); + + // prerender flag set, but wrong method + assert!(!state.should_render( + &Uri::from_static("http://foo/script.php"), + &make_request_with_flags( + true, + fqdn!("foo.caffeine.xyz"), + Method::POST, + Some(DomainFlags::new([FLAG_PRERENDER])) + ) + )); + + // prerender flag, unconfigured domain + assert!(state.should_render( + &Uri::from_static("http://foo/script.php"), + &make_request_with_flags( + true, + fqdn!("foo.zoo"), + Method::GET, + Some(DomainFlags::new([FLAG_PRERENDER])) + ) )); } diff --git a/src/routing/middleware/validate.rs b/src/routing/middleware/validate.rs index e8945da..c53ed2d 100644 --- a/src/routing/middleware/validate.rs +++ b/src/routing/middleware/validate.rs @@ -48,6 +48,10 @@ pub async fn middleware( .resolve(&authority) .ok_or_else(|| ErrorCause::Client(ClientError::UnknownDomain(authority.clone())))?; + if let Some(v) = lookup.flags { + request.extensions_mut().insert(v); + } + // If configured - try to resolve canister id from query params if state.canister_id_from_query_params && lookup.canister_id.is_none() { lookup.canister_id = canister_id_from_query_params(&request) diff --git a/src/test.rs b/src/test.rs index b9f241f..bdad01d 100644 --- a/src/test.rs +++ b/src/test.rs @@ -151,11 +151,7 @@ pub async fn setup_test_router(tasks: &mut TaskManager) -> (Router, Vec) let custom_domains = domains .clone() .into_iter() - .map(|x| CustomDomain { - name: fqdn!(&x), - canister_id: principal!("aaaaa-aa"), - timestamp: 0, - }) + .map(|x| CustomDomain::new(fqdn!(&x), principal!("aaaaa-aa"))) .collect::>(); let http_client = Arc::new(TestClient(512));