Harden the Edge Cookie withdrawal write path

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

@@ -7,8 +7,10 @@
 use error_stack::{Report, ResultExt};
 use fastly::kv_store::{InsertMode, KVStore};
 use trusted_server_core::ec::kv_backend::{
-    EcKvLookup, EcKvStore, EcKvWrite, EcKvWriteMode, EcKvWriteOutcome,
+    EcKvLookup, EcKvStore, EcKvWrite, EcKvWriteMode, EcKvWriteOutcome, ExactKeyMatch,
+    contains_exact_key,
 };
+use trusted_server_core::ec::log_id;
 use trusted_server_core::error::TrustedServerError;
 
 /// Fastly KV Store backend for the EC identity graph.
@@ -42,6 +44,18 @@ impl FastlyEcKvStore {
     }
 }
 
+/// Keys requested per page when checking for one exact key.
+const EXACT_MATCH_PAGE_SIZE: u32 = 100;
+/// Pages walked before giving up on an exact-key check.
+///
+/// A well-formed identifier is the whole key, so at most a handful of keys can
+/// carry it as a prefix and one page is the normal case. The cap stops a short
+/// prefix from walking the keyspace on the response path. Exceeding it is
+/// reported as an error rather than as absence, so the caller re-checks by
+/// lookup instead of discarding a withdrawal for an identity that may be held
+/// on a page this never read.
+const EXACT_MATCH_MAX_PAGES: usize = 4;
+
 impl EcKvStore for FastlyEcKvStore {
     fn store_name(&self) -> &str {
         &self.store_name
@@ -56,7 +70,7 @@ impl EcKvStore for FastlyEcKvStore {
                 return Err(
                     Report::new(err).change_context(TrustedServerError::KvStore {
                         store_name: self.store_name.clone(),
-                        message: format!("Failed to read key '{key}'"),
+                        message: format!("Failed to read key '{}'", log_id(key),),
                     }),
                 );
             }
@@ -98,7 +112,7 @@ impl EcKvStore for FastlyEcKvStore {
             Err(err) => Err(
                 Report::new(err).change_context(TrustedServerError::KvStore {
                     store_name: self.store_name.clone(),
-                    message: format!("Failed to write entry for key '{key}'"),
+                    message: format!("Failed to write entry for key '{}'", log_id(key)),
                 }),
             ),
         }
@@ -117,24 +131,63 @@ impl EcKvStore for FastlyEcKvStore {
             .execute()
             .change_context(TrustedServerError::KvStore {
                 store_name: self.store_name.clone(),
-                message: format!(
-                    "Failed to list keys with prefix '{}'",
-                    prefix.get(..8).unwrap_or(prefix),
-                ),
+                message: format!("Failed to list keys with prefix '{}'", log_id(prefix),),
             })?;
 
         #[allow(clippy::cast_possible_truncation)]
         let count = page.keys().len() as u32;
         Ok(count)
     }
 
+    fn key_exists(&self, key: &str) -> Result<bool, Report<TrustedServerError>> {
+        let store = self.open_store()?;
+        // The list is strongly consistent, unlike a lookup, but it matches by
+        // prefix, so the decision of what counts as a match belongs to
+        // `contains_exact_key` — which is where that behaviour is tested.
+        let pages = store
+            .build_list()
+            .prefix(key)
+            .limit(EXACT_MATCH_PAGE_SIZE)
+            .iter()
+            // Unbounded here on purpose: `contains_exact_key` stops reading at
+            // its budget, so the page count lives with the code that is tested
+            // against it rather than being recomputed at the call site.
+            .map(|page| match page {
+                // A store that reports nothing under the prefix is telling us
+                // the key is absent, which is what `lookup` already does for a
+                // missing key. Treating it as a failure would send every
+                // withdrawal naming an unissued identity down the error path.
+                Err(fastly::kv_store::KVStoreError::ItemNotFound) => Ok(Vec::new()),
+                other => other
+                    .map(fastly::kv_store::ListPage::into_keys)
+                    .change_context(TrustedServerError::KvStore {
+                        store_name: self.store_name.clone(),
+                        message: format!("Failed to check key '{}'", log_id(key)),
+                    }),
+            });
+
+        match contains_exact_key(pages, key, EXACT_MATCH_MAX_PAGES)? {
+            ExactKeyMatch::Found => Ok(true),
+            ExactKeyMatch::Absent => Ok(false),
+            // Not an answer, so it must not be returned as one. The caller
+            // falls back to a lookup on an error.
+            ExactKeyMatch::Undetermined => Err(Report::new(TrustedServerError::KvStore {
+                store_name: self.store_name.clone(),
+                message: format!(
+                    "Listing for key '{}' exceeded {EXACT_MATCH_MAX_PAGES} pages, so existence is undetermined",
+                    log_id(key)
+                ),
+            })),
+        }
+    }
+
     fn delete(&self, key: &str) -> Result<(), Report<TrustedServerError>> {
         let store = self.open_store()?;
         store
             .delete(key)
             .change_context(TrustedServerError::KvStore {
                 store_name: self.store_name.clone(),
-                message: format!("Failed to delete key '{key}'"),
+                message: format!("Failed to delete key '{}'", log_id(key)),
             })
     }
 }

crates/trusted-server-core/src/ec/admin.rs

@@ -677,6 +677,8 @@ mod tests {
 
     use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64};
 
+    use crate::ec::kv::TombstoneOutcome;
+
     use super::*;
     use crate::ec::kv_backend::test_support::InMemoryEcKv;
     use crate::ec::kv_backend::{EcKvStore as _, EcKvWrite, EcKvWriteMode};
@@ -1114,9 +1116,18 @@ mod tests {
     #[test]
     fn reports_tombstone_entries() {
         let ec_id = test_ec_id();
-        let kv = KvIdentityGraph::in_memory("test-store");
-        kv.write_withdrawal_tombstone(&ec_id)
-            .expect("should write tombstone");
+        // Only an identity the store already holds can be tombstoned, so seed
+        // the live entry the withdrawal replaces.
+        let kv = kv_with_entry(
+            &ec_id,
+            &KvEntry::minimal("bidstream.example", "uid-live", 1_741_824_000),
+        );
+        assert_eq!(
+            kv.write_withdrawal_tombstone(&ec_id)
+                .expect("should write tombstone"),
+            TombstoneOutcome::Written,
+            "should tombstone the seeded identity"
+        );
         let req = get_request(&format!("/_ts/admin/ec/{ec_id}"));
 
         let response = handle_admin_ec_lookup(Some(&kv), &test_registry(), &req)