Thread EC KV read through the request and recover orphaned cookies

crates/trusted-server-adapter-fastly/src/app.rs

@@ -407,6 +407,13 @@ fn build_ec_request_state(
         match EcContext::read_from_request_with_geo(settings, req, services, geo_info.as_ref()) {
             Ok(mut context) => {
                 context.set_device_signals(device_signals);
+                // Orphan-recovery eligibility is intentionally left false here.
+                // Authorizing it during generic pre-routing would let named
+                // routes and request-filter short circuits (e.g. a DataDome
+                // challenge) reach EC finalization and rotate an identity off a
+                // non-publisher response. It is granted only inside the
+                // publisher fallback, after filters pass and the origin start
+                // succeeds — see `dispatch_fallback`.
                 (context, None)
             }
             Err(report) => (EcContext::default(), Some(report)),
@@ -758,8 +765,8 @@ async fn dispatch_fallback(
         // Generate an EC ID if needed — mirrors the legacy catch-all arm.
         // Only for document navigations by recognised browsers; subresource
         // requests may lack consent signals such as Sec-GPC.
-        if ec.is_real_browser
-            && is_navigation_request(&req)
+        let is_publisher_navigation = ec.is_real_browser && is_navigation_request(&req);
+        if is_publisher_navigation
             && let Err(err) = ec
                 .ec_context
                 .generate_if_needed(&state.settings, ec.kv_graph.as_ref())
@@ -798,6 +805,14 @@ async fn dispatch_fallback(
                         .await
                         {
                             Ok(pub_response) => {
+                                // Origin start succeeded on the sole publisher-
+                                // page path: authorize orphan recovery now, and
+                                // only for real-browser document navigations.
+                                // Restricting it here keeps identity rotation
+                                // within the publisher-navigation boundary —
+                                // named routes, integration proxies, and filter
+                                // short circuits never reach this point.
+                                ec.ec_context.set_recovery_eligible(is_publisher_navigation);
                                 buffer_publisher_response_async(
                                     pub_response,
                                     &method,
@@ -2326,4 +2341,80 @@ mod tests {
             "the filter's response-header effect must be threaded out"
         );
     }
+
+    fn recovery_eligible_of(response: &Response) -> bool {
+        response
+            .extensions()
+            .get::<super::EcFinalizeState>()
+            .expect("response should carry EcFinalizeState")
+            .ec_context
+            .recovery_eligible()
+    }
+
+    fn browser_navigation_request(path: &str) -> edgezero_core::http::Request {
+        let uri = format!("https://test-publisher.com{path}");
+        let mut req = request_builder()
+            .method(Method::GET)
+            .uri(uri)
+            .header("sec-fetch-dest", "document")
+            .body(Body::empty())
+            .expect("should build request");
+        req.extensions_mut().insert(DeviceSignals::derive(
+            "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 \
+             (KHTML, like Gecko) Chrome/146.0.0.0 Safari/537.36",
+            Some("t13d1516h2_8daaf6152771_b186095e22b6"),
+            Some("1:65536;2:0;4:6291456;6:262144"),
+        ));
+        req
+    }
+
+    #[test]
+    fn named_route_response_is_not_recovery_eligible() {
+        // Orphan recovery must never be authorized on a named route: it is not a
+        // publisher-page navigation, so a missing KV row must not rotate the
+        // identity there.
+        let router = test_router();
+        let response = route(
+            &router,
+            empty_request(Method::GET, "/.well-known/trusted-server.json"),
+        );
+
+        assert!(
+            !recovery_eligible_of(&response),
+            "named-route responses must not authorize orphan recovery"
+        );
+    }
+
+    #[test]
+    fn filter_short_circuit_response_is_not_recovery_eligible() {
+        // A request-filter short circuit (e.g. a DataDome challenge/block) must
+        // not authorize orphan recovery even for a would-be publisher
+        // navigation: no publisher page was served.
+        let router = router_with_request_filters(vec![Arc::new(ChallengeRequestFilter)]);
+        let response = route(&router, browser_navigation_request("/some-page"));
+
+        assert_eq!(
+            response.status(),
+            StatusCode::FORBIDDEN,
+            "the challenge filter should short-circuit routing"
+        );
+        assert!(
+            !recovery_eligible_of(&response),
+            "a short-circuit filter response must not authorize orphan recovery"
+        );
+    }
+
+    #[test]
+    fn publisher_navigation_origin_start_failure_is_not_recovery_eligible() {
+        // Recovery is authorized only after a successful origin start. With no
+        // live backend the publisher origin fails, so even a real-browser
+        // document navigation must leave recovery unauthorized.
+        let router = test_router();
+        let response = route(&router, browser_navigation_request("/some-page"));
+
+        assert!(
+            !recovery_eligible_of(&response),
+            "an origin-start failure must not authorize orphan recovery"
+        );
+    }
 }

crates/trusted-server-adapter-fastly/src/main.rs

@@ -205,9 +205,9 @@ fn edgezero_main(mut req: FastlyRequest) {
         policy.apply_after_route_finalization(&mut response);
     }
 
-    if let Some(ec_state) = ec_state {
+    if let Some(mut ec_state) = ec_state {
         if let Some(settings) = settings_snapshot.as_deref() {
-            match apply_edgezero_ec_finalize(settings, &ec_state, &mut response) {
+            match apply_edgezero_ec_finalize(settings, &mut ec_state, &mut response) {
                 Ok(partner_registry) => {
                     send_edgezero_response(response, request_filter_effects.as_ref());
                     run_edgezero_pull_sync_after_send(settings, &partner_registry, &ec_state);
@@ -222,7 +222,7 @@ fn edgezero_main(mut req: FastlyRequest) {
         } else {
             match load_settings_from_config_store() {
                 Ok(settings) => {
-                    match apply_edgezero_ec_finalize(&settings, &ec_state, &mut response) {
+                    match apply_edgezero_ec_finalize(&settings, &mut ec_state, &mut response) {
                         Ok(partner_registry) => {
                             send_edgezero_response(response, request_filter_effects.as_ref());
                             run_edgezero_pull_sync_after_send(
@@ -286,7 +286,7 @@ fn apply_entry_point_finalize_headers(
 
 fn apply_edgezero_ec_finalize(
     settings: &Settings,
-    ec_state: &EcFinalizeState,
+    ec_state: &mut EcFinalizeState,
     response: &mut HttpResponse,
 ) -> Result<PartnerRegistry, Report<TrustedServerError>> {
     let partner_registry = PartnerRegistry::from_config(&settings.ec.partners)?;
@@ -297,7 +297,7 @@ fn apply_edgezero_ec_finalize(
     };
     ec_finalize_response(
         settings,
-        &ec_state.ec_context,
+        &mut ec_state.ec_context,
         finalize_kv_graph.as_ref(),
         &partner_registry,
         ec_state.eids_cookie.as_deref(),

crates/trusted-server-core/src/auction/endpoints.rs

@@ -12,10 +12,10 @@ use crate::auction::orchestrator::OrchestrationResult;
 use crate::consent::{consent_allows_server_side_auction, gate_eids_by_consent};
 use crate::constants::COOKIE_TS_EIDS;
 use crate::ec::EcContext;
+use crate::ec::EcKvSnapshot;
 use crate::ec::eids::{resolve_partner_ids, to_eids};
 use crate::ec::kv::KvIdentityGraph;
 use crate::ec::kv_types::MAX_UID_LENGTH;
-use crate::ec::log_id;
 use crate::ec::prebid_eids::parse_prebid_eids_cookie;
 use crate::ec::registry::PartnerRegistry;
 use crate::error::TrustedServerError;
@@ -245,9 +245,15 @@ pub async fn handle_auction(
         None
     };
 
-    // Resolve partner EIDs from the KV identity graph when the user has
-    // a valid EC and both KV and partner stores are available.
-    let eids = resolve_auction_eids(kv, registry, ec_context);
+    // Resolve partner EIDs from the KV identity graph when the user has a valid
+    // EC and both KV and partner stores are available. Gate the read on a
+    // present registry: without one, `resolve_auction_eids` yields no
+    // server-side EIDs, so the snapshot would be an unused billable KV read.
+    let auction_kv_snapshot = match (kv, ec_id, registry) {
+        (Some(graph), Some(ec_id), Some(_)) => graph.load_snapshot(ec_id),
+        _ => EcKvSnapshot::NotRead,
+    };
+    let eids = resolve_auction_eids(&auction_kv_snapshot, registry, ec_context);
 
     // Look up geo for device info.
     let geo = services
@@ -345,11 +351,10 @@ pub async fn handle_auction(
 /// store, no EC, consent denied). On KV or partner-resolution errors, logs a
 /// warning and returns empty EIDs so the auction can proceed in degraded mode.
 pub(crate) fn resolve_auction_eids(
-    kv: Option<&KvIdentityGraph>,
+    snapshot: &EcKvSnapshot,
     registry: Option<&PartnerRegistry>,
     ec_context: &EcContext,
 ) -> Option<Vec<Eid>> {
-    let kv = kv?;
     let registry = registry?;
 
     if !ec_context.ec_allowed() {
@@ -358,19 +363,15 @@ pub(crate) fn resolve_auction_eids(
 
     let ec_id = ec_context.ec_value()?;
 
-    let entry = match kv.get(ec_id) {
-        Ok(Some((entry, _generation))) => entry,
-        Ok(None) => return Some(Vec::new()),
-        Err(err) => {
-            log::warn!(
-                "Auction KV read failed for EC ID '{}': {err:?}",
-                log_id(ec_id)
-            );
-            return Some(Vec::new());
-        }
+    let Some(entry) = snapshot.entry_for(ec_id) else {
+        return Some(Vec::new());
     };
 
-    let resolved = resolve_partner_ids(registry, &entry);
+    if !entry.consent.ok {
+        return Some(Vec::new());
+    }
+
+    let resolved = resolve_partner_ids(registry, entry);
     Some(to_eids(&resolved))
 }
 
@@ -854,22 +855,24 @@ mod tests {
     }
 
     #[test]
-    fn resolve_auction_eids_returns_none_without_kv() {
+    fn resolve_auction_eids_returns_empty_without_snapshot() {
         let registry = PartnerRegistry::empty();
         let ec_id = format!("{}.ABC123", "a".repeat(64));
         let ec_context = make_ec_context(Jurisdiction::NonRegulated, Some(&ec_id));
 
-        let result = resolve_auction_eids(None, Some(&registry), &ec_context);
-        assert!(result.is_none(), "should return None when KV is missing");
+        let result = resolve_auction_eids(&EcKvSnapshot::NotRead, Some(&registry), &ec_context);
+        assert!(
+            result.is_some_and(|eids| eids.is_empty()),
+            "should degrade to empty EIDs without a snapshot"
+        );
     }
 
     #[test]
     fn resolve_auction_eids_returns_none_without_registry() {
-        let kv = KvIdentityGraph::failing("test_store");
         let ec_id = format!("{}.ABC123", "a".repeat(64));
         let ec_context = make_ec_context(Jurisdiction::NonRegulated, Some(&ec_id));
 
-        let result = resolve_auction_eids(Some(&kv), None, &ec_context);
+        let result = resolve_auction_eids(&EcKvSnapshot::NotRead, None, &ec_context);
         assert!(
             result.is_none(),
             "should return None when registry is missing"
@@ -878,12 +881,11 @@ mod tests {
 
     #[test]
     fn resolve_auction_eids_returns_none_when_consent_denied() {
-        let kv = KvIdentityGraph::failing("test_store");
         let registry = PartnerRegistry::empty();
         let ec_id = format!("{}.ABC123", "a".repeat(64));
         let ec_context = make_ec_context(Jurisdiction::Unknown, Some(&ec_id));
 
-        let result = resolve_auction_eids(Some(&kv), Some(&registry), &ec_context);
+        let result = resolve_auction_eids(&EcKvSnapshot::NotRead, Some(&registry), &ec_context);
         assert!(
             result.is_none(),
             "should return None when consent is denied"
@@ -892,11 +894,10 @@ mod tests {
 
     #[test]
     fn resolve_auction_eids_returns_none_when_no_ec() {
-        let kv = KvIdentityGraph::failing("test_store");
         let registry = PartnerRegistry::empty();
         let ec_context = make_ec_context(Jurisdiction::NonRegulated, None);
 
-        let result = resolve_auction_eids(Some(&kv), Some(&registry), &ec_context);
+        let result = resolve_auction_eids(&EcKvSnapshot::NotRead, Some(&registry), &ec_context);
         assert!(
             result.is_none(),
             "should return None when no EC value is present"
@@ -905,14 +906,14 @@ mod tests {
 
     #[test]
     fn resolve_auction_eids_returns_empty_on_kv_miss() {
-        let kv = KvIdentityGraph::failing("nonexistent_store");
         let registry = PartnerRegistry::empty();
         let ec_id = format!("{}.ABC123", "a".repeat(64));
         let ec_context = make_ec_context(Jurisdiction::NonRegulated, Some(&ec_id));
 
-        // KV store doesn't exist, so the get() call will error — should return
-        // empty Vec (degraded mode), not None.
-        let result = resolve_auction_eids(Some(&kv), Some(&registry), &ec_context);
+        let snapshot = EcKvSnapshot::Failed {
+            ec_id: ec_id.clone(),
+        };
+        let result = resolve_auction_eids(&snapshot, Some(&registry), &ec_context);
         let eids = result.expect("should return Some on KV error (degraded mode)");
         assert!(
             eids.is_empty(),