Expose Fastly geo headers on every response (#375)

prk-Jr · web-flow · commit c426a5799427 · 2026-02-27T21:42:00.000+05:30
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. Resolves #280
diff --git 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,163 @@ 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());
+        if self.has_metro_code() {
+            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<String> {
-    // 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"
         );
-        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 set_response_headers_omits_metro_code_when_zero() {
+        let geo = GeoInfo {
+            metro_code: 0,
+            ..sample_geo_info()
+        };
+        let mut response = Response::new();
+
+        geo.set_response_headers(&mut response);
+
+        assert!(
+            response.get_header(HEADER_X_GEO_METRO_CODE).is_none(),
+            "should not set metro code header when metro_code is 0"
+        );
+        assert!(
+            response.get_header(HEADER_X_GEO_CITY).is_some(),
+            "should still set city header"
+        );
     }
 
-    None
+    #[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"
+        );
+    }
 }
diff --git 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,
@@ -64,7 +66,11 @@ async fn route_request(
     integration_registry: &IntegrationRegistry,
     req: Request,
 ) -> Result<Response, Error> {
-    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) {
+        finalize_response(settings, geo_info.as_ref(), &mut response);
         return Ok(response);
     }
 
@@ -132,6 +138,26 @@ async fn route_request(
     // Convert any errors to HTTP error responses
     let mut response = result.unwrap_or_else(|e| to_error_response(&e));
 
+    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.
+///
+/// Header precedence (last write wins): geo headers are set first, then
+/// version/staging, then operator-configured `settings.response_headers`.
+/// This means operators can intentionally override any managed header.
+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);
     }
@@ -142,8 +168,6 @@ async fn route_request(
     for (key, value) in &settings.response_headers {
         response.set_header(key, value);
     }
-
-    Ok(response)
 }
 
 fn init_logger() {
