Fail closed when withheld country list is missing or TES country-codes Err

visibility-filtering/hydration/tes_hydrator.rs

@@ -1,4 +1,4 @@
-use crate::hydration::batch::TweetHydrationBatch;
+use crate::hydration::batch::{Hydrated, TweetHydrationBatch};
 use crate::hydration::metrics::{record_batch_size, timed_keyed_rpc, timed_results};
 use crate::models::{
     CoreFeature, MediaFeature, NsfwFeature, TweetCandidateInput, TweetFeatures, TweetId,
@@ -12,6 +12,7 @@ use xai_core_entities::tweet_entity_service_client::TESClient;
 
 const CLIENT_TIMEOUT: Duration = Duration::from_millis(150);
 const CLIENT: &str = "tes";
+const WORLDWIDE_WITHHELD_COUNTRY: &str = "xx";
 
 pub struct TesHydrator {
     pub tes_client: Arc<dyn TESClient + Send + Sync>,
@@ -24,6 +25,7 @@ pub(crate) struct TweetHydration {
     pub(crate) nsfw_user: TweetHydrationBatch<bool>,
     pub(crate) nsfw_admin: TweetHydrationBatch<bool>,
     pub(crate) takedown_reasons: TweetHydrationBatch<Vec<TakedownReason>>,
+    pub(crate) takedown_country_codes: TweetHydrationBatch<Vec<String>>,
     pub(crate) edit_control: TweetHydrationBatch<EditControl>,
     pub(crate) media: TweetHydrationBatch<MediaFeature>,
 }
@@ -69,6 +71,7 @@ impl TesHydrator {
             nsfw_user,
             nsfw_admin,
             takedown_reasons,
+            takedown_country_codes,
             edit_control,
             media_entities,
         ) = tokio::join!(
@@ -112,6 +115,14 @@ impl TesHydrator {
                 CLIENT_TIMEOUT,
                 self.tes_client.get_takedown_reasons(raw_ids.clone()),
             ),
+            timed_results(
+                CLIENT,
+                "get_takedown_country_codes",
+                safety_level,
+                &candidate_count_by_key,
+                CLIENT_TIMEOUT,
+                self.tes_client.get_takedown_country_codes(raw_ids.clone()),
+            ),
             timed_results(
                 CLIENT,
                 "get_edit_control",
@@ -136,6 +147,7 @@ impl TesHydrator {
             nsfw_user: nsfw_user.map_keys(TweetId),
             nsfw_admin: nsfw_admin.map_keys(TweetId),
             takedown_reasons: takedown_reasons.map_keys(TweetId),
+            takedown_country_codes: takedown_country_codes.map_keys(TweetId),
             edit_control: edit_control.map_keys(TweetId),
             media: media_entities.map_keys(TweetId).map(media_feature),
         }
@@ -178,6 +190,7 @@ fn build_tweet_features(
     let is_nullcast = tweet_keyed.nullcast.get(&id).copied().unwrap_or(false);
     let is_community_tweet = tweet_keyed.community.get(&id).is_some();
     let takedown_reasons = tweet_keyed.takedown_reasons.get_or_default(&id);
+    let takedown_country_codes = withheld_in_countries(&tweet_keyed.takedown_country_codes, &id);
     let nsfw = NsfwFeature {
         user: tweet_keyed.nsfw_user.get(&id).copied().unwrap_or(false),
         admin: tweet_keyed.nsfw_admin.get(&id).copied().unwrap_or(false),
@@ -193,6 +206,7 @@ fn build_tweet_features(
             },
             media,
             takedown_reasons,
+            takedown_country_codes,
             nsfw,
             is_nullcast,
             is_community_tweet,
@@ -201,6 +215,14 @@ fn build_tweet_features(
         .unwrap_or_default()
 }
 
+fn withheld_in_countries(batch: &TweetHydrationBatch<Vec<String>>, id: &TweetId) -> Vec<String> {
+    match batch.hydrated(id) {
+        Some(Hydrated::Found(codes)) => codes.clone(),
+        Some(Hydrated::Failed(_)) => vec![WORLDWIDE_WITHHELD_COUNTRY.to_string()],
+        Some(Hydrated::NotFound) | None => Vec::new(),
+    }
+}
+
 fn media_feature(entities: MediaEntities) -> MediaFeature {
     let mut feature = MediaFeature {
         has_media: !entities.is_empty(),
@@ -474,4 +496,58 @@ mod tests {
         assert!(f.core.text.is_empty());
         assert!(!f.media.has_media);
     }
+
+    fn failed<V>(id: u64) -> TweetHydrationBatch<V> {
+        TweetHydrationBatch::from_results(
+            [TweetId(id)],
+            HashMap::from([(TweetId(id), Err::<Option<V>, _>("tes unavailable"))]),
+        )
+    }
+
+    fn not_found<V>(id: u64) -> TweetHydrationBatch<V> {
+        TweetHydrationBatch::from_results(
+            [TweetId(id)],
+            HashMap::from([(TweetId(id), Ok::<Option<V>, _>(None))]),
+        )
+    }
+
+    fn assemble_with_country_codes(batch: TweetHydrationBatch<Vec<String>>) -> TweetFeatures {
+        let candidates = vec![candidate(10, 100)];
+        let core_datas = HashMap::from([(
+            TweetId(10),
+            PureCoreData {
+                author_id: 100,
+                ..Default::default()
+            },
+        )]);
+        hydrator()
+            .assemble_tweet_features(
+                &candidates,
+                &core_datas,
+                &TweetHydration {
+                    takedown_country_codes: batch,
+                    ..Default::default()
+                },
+            )
+            .remove(&TweetId(10))
+            .unwrap()
+    }
+
+    #[test]
+    fn assemble_keeps_found_withheld_in_countries() {
+        let f = assemble_with_country_codes(found(10, vec!["de".to_string(), "fr".to_string()]));
+        assert_eq!(f.takedown_country_codes, vec!["de", "fr"]);
+    }
+
+    #[test]
+    fn assemble_treats_missing_withheld_in_countries_as_empty() {
+        let f = assemble_with_country_codes(not_found(10));
+        assert!(f.takedown_country_codes.is_empty());
+    }
+
+    #[test]
+    fn assemble_fail_closes_withheld_in_countries_err_as_worldwide() {
+        let f = assemble_with_country_codes(failed(10));
+        assert_eq!(f.takedown_country_codes, vec!["xx"]);
+    }
 }

visibility-filtering/models/tweet.rs

@@ -27,6 +27,7 @@ pub struct TweetFeatures {
     pub core: CoreFeature,
     pub media: MediaFeature,
     pub takedown_reasons: Vec<xai_core_entities::entities::TakedownReason>,
+    pub takedown_country_codes: Vec<String>,
     pub nsfw: NsfwFeature,
     pub is_nullcast: bool,
     pub is_community_tweet: bool,

visibility-filtering/rules/context.rs

@@ -249,7 +249,25 @@ pub struct TakedownPredicates<'a> {
 impl TakedownPredicates<'_> {
     #[inline]
     pub fn legal_in_viewer_country(&self) -> bool {
-        self.in_viewer_country(legal_takedown_country)
+        let viewer_country = self.ctx.viewer.country_code.as_deref();
+        let reason_countries = self
+            .ctx
+            .candidate
+            .tweet_features
+            .takedown_reasons
+            .iter()
+            .filter_map(legal_takedown_country);
+        let withheld_in_countries = self
+            .ctx
+            .candidate
+            .tweet_features
+            .takedown_country_codes
+            .iter()
+            .map(String::as_str);
+        withheld_countries_apply(
+            viewer_country,
+            reason_countries.chain(withheld_in_countries),
+        )
     }
 
     #[inline]
@@ -259,18 +277,15 @@ impl TakedownPredicates<'_> {
 
     #[inline]
     fn in_viewer_country(&self, extractor: fn(&TakedownReason) -> Option<&str>) -> bool {
-        let viewer_country = self.ctx.viewer.country_code.as_deref();
-        self.ctx
-            .candidate
-            .tweet_features
-            .takedown_reasons
-            .iter()
-            .filter_map(extractor)
-            .any(|c| {
-                c.eq_ignore_ascii_case(WORLDWIDE_COUNTRY_CODE)
-                    || c.eq_ignore_ascii_case(WORLDWIDE_COPYRIGHT_COUNTRY_CODE)
-                    || viewer_country.is_some_and(|v| c.eq_ignore_ascii_case(v))
-            })
+        withheld_countries_apply(
+            self.ctx.viewer.country_code.as_deref(),
+            self.ctx
+                .candidate
+                .tweet_features
+                .takedown_reasons
+                .iter()
+                .filter_map(extractor),
+        )
     }
 
     #[inline]
@@ -291,6 +306,31 @@ impl TakedownPredicates<'_> {
 const WORLDWIDE_COUNTRY_CODE: &str = "xx";
 const WORLDWIDE_COPYRIGHT_COUNTRY_CODE: &str = "xy";
 
+fn withheld_countries_apply<'a>(
+    viewer_country: Option<&str>,
+    countries: impl IntoIterator<Item = &'a str>,
+) -> bool {
+    let mut saw_country_scoped = false;
+    for raw in countries {
+        let country = if raw.is_empty() {
+            WORLDWIDE_COUNTRY_CODE
+        } else {
+            raw
+        };
+        if country.eq_ignore_ascii_case(WORLDWIDE_COUNTRY_CODE)
+            || country.eq_ignore_ascii_case(WORLDWIDE_COPYRIGHT_COUNTRY_CODE)
+        {
+            return true;
+        }
+        saw_country_scoped = true;
+        if viewer_country.is_some_and(|v| country.eq_ignore_ascii_case(v)) {
+            return true;
+        }
+    }
+    // Country-scoped withhold with no viewer country used to Allow (fail-open).
+    saw_country_scoped && viewer_country.is_none()
+}
+
 fn legal_takedown_country(reason: &TakedownReason) -> Option<&str> {
     match reason {
         TakedownReason::LegalRequest { country_code }

visibility-filtering/rules/golden_corpus.rs

@@ -155,6 +155,15 @@ fn takedown_candidate(reason: TakedownReason) -> HydratedTweetCandidate {
         .build()
 }
 
+fn withheld_in_countries_candidate(codes: &[&str]) -> HydratedTweetCandidate {
+    candidate()
+        .with_tweet_features(TweetFeatures {
+            takedown_country_codes: codes.iter().map(|s| s.to_string()).collect(),
+            ..Default::default()
+        })
+        .build()
+}
+
 fn exclusive_candidate(viewer_super_follows_author: bool) -> HydratedTweetCandidate {
     let mut c = candidate().build();
     c.exclusive_content = Some(ExclusiveContentFeatures {
@@ -621,6 +630,50 @@ fn tweet_shape_cases() -> Vec<Case> {
             expected_action: Allow,
             expected_decided_by: None,
         },
+        Case {
+            name: "legal_takedown_country_scoped_drops_without_viewer_country",
+            level: TimelineHome,
+            viewer: viewer(VIEWER_ID),
+            candidate: takedown_candidate(TakedownReason::LegalRequest {
+                country_code: "de".to_string(),
+            }),
+            expected_action: Drop(FilteredReason::UnspecifiedReason),
+            expected_decided_by: Some("DropLegalTakendownPostRule"),
+        },
+        Case {
+            name: "legal_takedown_empty_country_is_worldwide",
+            level: TimelineHome,
+            viewer: viewer_in_country("us"),
+            candidate: takedown_candidate(TakedownReason::LegalRequest {
+                country_code: String::new(),
+            }),
+            expected_action: Drop(FilteredReason::UnspecifiedReason),
+            expected_decided_by: Some("DropLegalTakendownPostRule"),
+        },
+        Case {
+            name: "withheld_in_countries_drops_matching_viewer",
+            level: TimelineHome,
+            viewer: viewer_in_country("de"),
+            candidate: withheld_in_countries_candidate(&["de", "fr"]),
+            expected_action: Drop(FilteredReason::UnspecifiedReason),
+            expected_decided_by: Some("DropLegalTakendownPostRule"),
+        },
+        Case {
+            name: "withheld_in_countries_allows_other_country",
+            level: TimelineHome,
+            viewer: viewer_in_country("us"),
+            candidate: withheld_in_countries_candidate(&["de", "fr"]),
+            expected_action: Allow,
+            expected_decided_by: None,
+        },
+        Case {
+            name: "withheld_in_countries_err_worldwide_drops_oon",
+            level: TimelineHomeRecommendations,
+            viewer: viewer_in_country("us"),
+            candidate: withheld_in_countries_candidate(&["xx"]),
+            expected_action: Drop(FilteredReason::UnspecifiedReason),
+            expected_decided_by: Some("DropLegalTakendownPostRule"),
+        },
     ]
 }
 

visibility-filtering/rules/tweet_rules.rs

@@ -838,7 +838,7 @@ mod tests {
         ]);
         assert_drops(legal, &viewer_with_country("de"), &legal_c, &reason);
         assert_allows(legal, &viewer_with_country("us"), &legal_c);
-        assert_allows(legal, &viewer(VIEWER_ID), &legal_c);
+        assert_drops(legal, &viewer(VIEWER_ID), &legal_c, &reason);
 
         let bystander = takedown_candidate(vec![TakedownReason::BystanderReport {
             country_code: "de".to_string(),
@@ -895,11 +895,46 @@ mod tests {
         let country_scoped = takedown_candidate(vec![TakedownReason::LegalRequest {
             country_code: "de".to_string(),
         }]);
-        assert_allows(legal, &viewer(VIEWER_ID), &country_scoped);
+        assert_drops(legal, &viewer(VIEWER_ID), &country_scoped, &reason);
         let bystander_scoped = takedown_candidate(vec![TakedownReason::BystanderReport {
             country_code: "de".to_string(),
         }]);
-        assert_allows(local, &viewer(VIEWER_ID), &bystander_scoped);
+        assert_drops(local, &viewer(VIEWER_ID), &bystander_scoped, &reason);
+
+        let empty_country = takedown_candidate(vec![TakedownReason::LegalRequest {
+            country_code: String::new(),
+        }]);
+        assert_drops(legal, &viewer_with_country("us"), &empty_country, &reason);
+        assert_drops(legal, &viewer(VIEWER_ID), &empty_country, &reason);
+
+        let withheld_in_countries = candidate()
+            .with_tweet_features(TweetFeatures {
+                takedown_country_codes: vec!["de".to_string()],
+                ..Default::default()
+            })
+            .build();
+        assert_drops(
+            legal,
+            &viewer_with_country("de"),
+            &withheld_in_countries,
+            &reason,
+        );
+        assert_allows(legal, &viewer_with_country("us"), &withheld_in_countries);
+        assert_drops(legal, &viewer(VIEWER_ID), &withheld_in_countries, &reason);
+        assert_allows(local, &viewer_with_country("de"), &withheld_in_countries);
+
+        let withheld_err_worldwide = candidate()
+            .with_tweet_features(TweetFeatures {
+                takedown_country_codes: vec!["xx".to_string()],
+                ..Default::default()
+            })
+            .build();
+        assert_drops(
+            legal,
+            &viewer_with_country("us"),
+            &withheld_err_worldwide,
+            &reason,
+        );
 
         let dmca = takedown_candidate(vec![TakedownReason::Dmca]);
         assert_drops(legal, &viewer_with_country("de"), &dmca, &reason);