From 15b9182b6c89cb98bfa5527388a7cabac326b46b Mon Sep 17 00:00:00 2001 From: prk-Jr Date: Wed, 25 Feb 2026 16:05:46 +0530 Subject: [PATCH 1/3] Expose Fastly geo headers on every response from Trusted Server Add GeoInfo::set_response_headers() to set x-geo-* headers on responses and wire it into the centralized response middleware in main.rs. When geo data is unavailable, sets x-geo-info-available: false. Remove unused get_dma_code() function. Add unit tests for the new method. Resolves #280 --- crates/common/src/geo.rs | 226 +++++++++++++++++++++++++++----------- crates/fastly/src/main.rs | 14 ++- 2 files changed, 176 insertions(+), 64 deletions(-) diff --git a/crates/common/src/geo.rs b/crates/common/src/geo.rs index a903c31d..2620bb57 100644 --- a/crates/common/src/geo.rs +++ b/crates/common/src/geo.rs @@ -4,11 +4,11 @@ //! information from incoming requests, particularly DMA (Designated Market Area) codes. use fastly::geo::geo_lookup; -use fastly::Request; +use fastly::{Request, Response}; use crate::constants::{ HEADER_X_GEO_CITY, HEADER_X_GEO_CONTINENT, HEADER_X_GEO_COORDINATES, HEADER_X_GEO_COUNTRY, - HEADER_X_GEO_INFO_AVAILABLE, HEADER_X_GEO_METRO_CODE, + HEADER_X_GEO_INFO_AVAILABLE, HEADER_X_GEO_METRO_CODE, HEADER_X_GEO_REGION, }; /// Geographic information extracted from a request. @@ -81,72 +81,172 @@ impl GeoInfo { pub fn has_metro_code(&self) -> bool { self.metro_code > 0 } + + /// Sets geo information headers on the response. + /// + /// Adds `x-geo-city`, `x-geo-country`, `x-geo-continent`, `x-geo-coordinates`, + /// `x-geo-metro-code`, `x-geo-region` (when available), and + /// `x-geo-info-available: true` to the given response. + pub fn set_response_headers(&self, response: &mut Response) { + response.set_header(HEADER_X_GEO_CITY, &self.city); + response.set_header(HEADER_X_GEO_COUNTRY, &self.country); + response.set_header(HEADER_X_GEO_CONTINENT, &self.continent); + response.set_header(HEADER_X_GEO_COORDINATES, self.coordinates_string()); + response.set_header(HEADER_X_GEO_METRO_CODE, self.metro_code.to_string()); + if let Some(ref region) = self.region { + response.set_header(HEADER_X_GEO_REGION, region); + } + response.set_header(HEADER_X_GEO_INFO_AVAILABLE, "true"); + } } -/// Extracts the DMA (Designated Market Area) code from the request's geolocation data. -/// -/// This function: -/// 1. Checks if running in Fastly environment -/// 2. Performs geo lookup based on client IP -/// 3. Sets various geo headers on the request -/// 4. Returns the metro code (DMA) if available -/// -/// # Arguments -/// -/// * `req` - The request to extract DMA code from -/// -/// # Returns -/// -/// The DMA/metro code as a string if available, None otherwise -pub fn get_dma_code(req: &mut Request) -> Option { - // Debug: Check if we're running in Fastly environment - log::info!("Fastly Environment Check:"); - log::info!( - " FASTLY_POP: {}", - std::env::var("FASTLY_POP").unwrap_or_else(|_| "not in Fastly".to_string()) - ); - log::info!( - " FASTLY_REGION: {}", - std::env::var("FASTLY_REGION").unwrap_or_else(|_| "not in Fastly".to_string()) - ); - - // Get detailed geo information using geo_lookup - if let Some(geo) = req.get_client_ip_addr().and_then(geo_lookup) { - log::info!("Geo Information Found:"); - - // Set all available geo information in headers - let city = geo.city(); - req.set_header(HEADER_X_GEO_CITY, city); - log::info!(" City: {}", city); - - let country = geo.country_code(); - req.set_header(HEADER_X_GEO_COUNTRY, country); - log::info!(" Country: {}", country); - - req.set_header(HEADER_X_GEO_CONTINENT, format!("{:?}", geo.continent())); - log::info!(" Continent: {:?}", geo.continent()); - - req.set_header( - HEADER_X_GEO_COORDINATES, - format!("{},{}", geo.latitude(), geo.longitude()), +#[cfg(test)] +mod tests { + use super::*; + use fastly::Response; + + fn sample_geo_info() -> GeoInfo { + GeoInfo { + city: "San Francisco".to_string(), + country: "US".to_string(), + continent: "NorthAmerica".to_string(), + latitude: 37.7749, + longitude: -122.4194, + metro_code: 807, + region: Some("CA".to_string()), + } + } + + #[test] + fn set_response_headers_sets_all_geo_headers() { + let geo = sample_geo_info(); + let mut response = Response::new(); + + geo.set_response_headers(&mut response); + + assert_eq!( + response + .get_header(HEADER_X_GEO_CITY) + .expect("should have city header") + .to_str() + .expect("should be valid str"), + "San Francisco", + "should set city header" + ); + assert_eq!( + response + .get_header(HEADER_X_GEO_COUNTRY) + .expect("should have country header") + .to_str() + .expect("should be valid str"), + "US", + "should set country header" + ); + assert_eq!( + response + .get_header(HEADER_X_GEO_CONTINENT) + .expect("should have continent header") + .to_str() + .expect("should be valid str"), + "NorthAmerica", + "should set continent header" + ); + assert_eq!( + response + .get_header(HEADER_X_GEO_COORDINATES) + .expect("should have coordinates header") + .to_str() + .expect("should be valid str"), + "37.7749,-122.4194", + "should set coordinates header" + ); + assert_eq!( + response + .get_header(HEADER_X_GEO_METRO_CODE) + .expect("should have metro code header") + .to_str() + .expect("should be valid str"), + "807", + "should set metro code header" + ); + assert_eq!( + response + .get_header(HEADER_X_GEO_REGION) + .expect("should have region header") + .to_str() + .expect("should be valid str"), + "CA", + "should set region header" + ); + assert_eq!( + response + .get_header(HEADER_X_GEO_INFO_AVAILABLE) + .expect("should have info available header") + .to_str() + .expect("should be valid str"), + "true", + "should set geo info available to true" + ); + } + + #[test] + fn set_response_headers_omits_region_when_none() { + let geo = GeoInfo { + region: None, + ..sample_geo_info() + }; + let mut response = Response::new(); + + geo.set_response_headers(&mut response); + + assert!( + response.get_header(HEADER_X_GEO_REGION).is_none(), + "should not set region header when region is None" + ); + // Other headers should still be present + assert!( + response.get_header(HEADER_X_GEO_CITY).is_some(), + "should still set city header" + ); + assert_eq!( + response + .get_header(HEADER_X_GEO_INFO_AVAILABLE) + .expect("should have info available header") + .to_str() + .expect("should be valid str"), + "true", + "should still set geo info available to true" + ); + } + + #[test] + fn coordinates_string_formats_lat_lon() { + let geo = sample_geo_info(); + assert_eq!( + geo.coordinates_string(), + "37.7749,-122.4194", + "should format coordinates as lat,lon" ); - log::info!(" Location: ({}, {})", geo.latitude(), geo.longitude()); - - // Get and set the metro code (DMA) - let metro_code = geo.metro_code(); - req.set_header(HEADER_X_GEO_METRO_CODE, metro_code.to_string()); - log::info!("Found DMA/Metro code: {}", metro_code); - return Some(metro_code.to_string()); - } else { - log::info!("No geo information available for the request"); - req.set_header(HEADER_X_GEO_INFO_AVAILABLE, "false"); } - // If no metro code is found, log all request headers for debugging - log::info!("No DMA/Metro code found. All request headers:"); - for (name, value) in req.get_headers() { - log::info!(" {}: {:?}", name, value); + #[test] + fn has_metro_code_returns_true_for_positive() { + let geo = sample_geo_info(); + assert!( + geo.has_metro_code(), + "should return true for metro code 807" + ); } - None + #[test] + fn has_metro_code_returns_false_for_zero() { + let geo = GeoInfo { + metro_code: 0, + ..sample_geo_info() + }; + assert!( + !geo.has_metro_code(), + "should return false for metro code 0" + ); + } } diff --git a/crates/fastly/src/main.rs b/crates/fastly/src/main.rs index 49156972..69ace53a 100644 --- a/crates/fastly/src/main.rs +++ b/crates/fastly/src/main.rs @@ -7,9 +7,11 @@ use trusted_server_common::auction::endpoints::handle_auction; use trusted_server_common::auction::{build_orchestrator, AuctionOrchestrator}; use trusted_server_common::auth::enforce_basic_auth; use trusted_server_common::constants::{ - ENV_FASTLY_IS_STAGING, ENV_FASTLY_SERVICE_VERSION, HEADER_X_TS_ENV, HEADER_X_TS_VERSION, + ENV_FASTLY_IS_STAGING, ENV_FASTLY_SERVICE_VERSION, HEADER_X_GEO_INFO_AVAILABLE, + HEADER_X_TS_ENV, HEADER_X_TS_VERSION, }; use trusted_server_common::error::TrustedServerError; +use trusted_server_common::geo::GeoInfo; use trusted_server_common::integrations::IntegrationRegistry; use trusted_server_common::proxy::{ handle_first_party_click, handle_first_party_proxy, handle_first_party_proxy_rebuild, @@ -68,6 +70,9 @@ async fn route_request( return Ok(response); } + // Extract geo info before routing consumes the request + let geo_info = GeoInfo::from_request(&req); + // Get path and method for routing let path = req.get_path().to_string(); let method = req.get_method().clone(); @@ -132,6 +137,13 @@ async fn route_request( // Convert any errors to HTTP error responses let mut response = result.unwrap_or_else(|e| to_error_response(&e)); + // Set geo headers on the response + if let Some(ref geo) = geo_info { + geo.set_response_headers(&mut response); + } else { + response.set_header(HEADER_X_GEO_INFO_AVAILABLE, "false"); + } + if let Ok(v) = ::std::env::var(ENV_FASTLY_SERVICE_VERSION) { response.set_header(HEADER_X_TS_VERSION, v); } From 6b24b833594f5e347ed61cb0adb63a05edb6c85c Mon Sep 17 00:00:00 2001 From: prk-Jr Date: Wed, 25 Feb 2026 16:20:50 +0530 Subject: [PATCH 2/3] Ensure geo headers are applied to basic auth responses and encapsulate geo header logic into a new helper function. --- crates/fastly/src/main.rs | 26 ++++++++++++++++---------- 1 file changed, 16 insertions(+), 10 deletions(-) diff --git a/crates/fastly/src/main.rs b/crates/fastly/src/main.rs index 69ace53a..043a5d83 100644 --- a/crates/fastly/src/main.rs +++ b/crates/fastly/src/main.rs @@ -66,13 +66,14 @@ async fn route_request( integration_registry: &IntegrationRegistry, req: Request, ) -> Result { - if let Some(response) = enforce_basic_auth(settings, &req) { + // Extract geo info before auth check or routing consumes the request + let geo_info = GeoInfo::from_request(&req); + + if let Some(mut response) = enforce_basic_auth(settings, &req) { + apply_geo_headers(&geo_info, &mut response); return Ok(response); } - // Extract geo info before routing consumes the request - let geo_info = GeoInfo::from_request(&req); - // Get path and method for routing let path = req.get_path().to_string(); let method = req.get_method().clone(); @@ -137,12 +138,7 @@ async fn route_request( // Convert any errors to HTTP error responses let mut response = result.unwrap_or_else(|e| to_error_response(&e)); - // Set geo headers on the response - if let Some(ref geo) = geo_info { - geo.set_response_headers(&mut response); - } else { - response.set_header(HEADER_X_GEO_INFO_AVAILABLE, "false"); - } + apply_geo_headers(&geo_info, &mut response); if let Ok(v) = ::std::env::var(ENV_FASTLY_SERVICE_VERSION) { response.set_header(HEADER_X_TS_VERSION, v); @@ -158,6 +154,16 @@ async fn route_request( Ok(response) } +/// Sets geo headers on a response. If geo data is available, sets all `x-geo-*` +/// headers; otherwise sets `x-geo-info-available: false`. +fn apply_geo_headers(geo_info: &Option, response: &mut Response) { + if let Some(ref geo) = geo_info { + geo.set_response_headers(response); + } else { + response.set_header(HEADER_X_GEO_INFO_AVAILABLE, "false"); + } +} + fn init_logger() { let logger = Logger::builder() .default_endpoint("tslog") From 04ef49c2f4a03912f548fc188bb82b89d6e12038 Mon Sep 17 00:00:00 2001 From: prk-Jr Date: Wed, 25 Feb 2026 19:11:18 +0530 Subject: [PATCH 3/3] Fix: . Idiomatic signature for apply_geo_headers(), centralize response finalization --- crates/common/src/geo.rs | 31 ------------------------------- crates/fastly/src/main.rs | 30 ++++++++++++++++-------------- 2 files changed, 16 insertions(+), 45 deletions(-) diff --git a/crates/common/src/geo.rs b/crates/common/src/geo.rs index 2620bb57..3a336f79 100644 --- a/crates/common/src/geo.rs +++ b/crates/common/src/geo.rs @@ -218,35 +218,4 @@ mod tests { "should still set geo info available to true" ); } - - #[test] - fn coordinates_string_formats_lat_lon() { - let geo = sample_geo_info(); - assert_eq!( - geo.coordinates_string(), - "37.7749,-122.4194", - "should format coordinates as lat,lon" - ); - } - - #[test] - fn has_metro_code_returns_true_for_positive() { - let geo = sample_geo_info(); - assert!( - geo.has_metro_code(), - "should return true for metro code 807" - ); - } - - #[test] - fn has_metro_code_returns_false_for_zero() { - let geo = GeoInfo { - metro_code: 0, - ..sample_geo_info() - }; - assert!( - !geo.has_metro_code(), - "should return false for metro code 0" - ); - } } diff --git a/crates/fastly/src/main.rs b/crates/fastly/src/main.rs index 043a5d83..a4bc94b1 100644 --- a/crates/fastly/src/main.rs +++ b/crates/fastly/src/main.rs @@ -70,7 +70,7 @@ async fn route_request( let geo_info = GeoInfo::from_request(&req); if let Some(mut response) = enforce_basic_auth(settings, &req) { - apply_geo_headers(&geo_info, &mut response); + finalize_response(settings, geo_info.as_ref(), &mut response); return Ok(response); } @@ -138,7 +138,21 @@ async fn route_request( // Convert any errors to HTTP error responses let mut response = result.unwrap_or_else(|e| to_error_response(&e)); - apply_geo_headers(&geo_info, &mut response); + finalize_response(settings, geo_info.as_ref(), &mut response); + + Ok(response) +} + +/// Applies all standard response headers: geo, version, staging, and configured headers. +/// +/// Called from every response path (including auth early-returns) so that all +/// outgoing responses carry a consistent set of Trusted Server headers. +fn finalize_response(settings: &Settings, geo_info: Option<&GeoInfo>, response: &mut Response) { + if let Some(geo) = geo_info { + geo.set_response_headers(response); + } else { + response.set_header(HEADER_X_GEO_INFO_AVAILABLE, "false"); + } if let Ok(v) = ::std::env::var(ENV_FASTLY_SERVICE_VERSION) { response.set_header(HEADER_X_TS_VERSION, v); @@ -150,18 +164,6 @@ async fn route_request( for (key, value) in &settings.response_headers { response.set_header(key, value); } - - Ok(response) -} - -/// Sets geo headers on a response. If geo data is available, sets all `x-geo-*` -/// headers; otherwise sets `x-geo-info-available: false`. -fn apply_geo_headers(geo_info: &Option, response: &mut Response) { - if let Some(ref geo) = geo_info { - geo.set_response_headers(response); - } else { - response.set_header(HEADER_X_GEO_INFO_AVAILABLE, "false"); - } } fn init_logger() {