Merge pull request #342 from spoo-me/feat/safety-secondary-destinations

Author: Zingzy

repositories/indexes.py

@@ -118,6 +118,10 @@ async def ensure_indexes(
         )
         # Every enforcement query is an equality match on dest.host.
         await _url_col.create_index([("dest.host", 1)], name="dest_host", sparse=True)
+    # Only v2 links carry geo rules or a pre-start page.
+    await urls_v2_col.create_index(
+        [("dest.secondary_hosts", 1)], name="dest_secondary_hosts", sparse=True
+    )
 
     # ── clicks (time-series) ───────────────────────────────────────────────
     # Create the time-series collection if it doesn't exist yet.

repositories/url_repository.py

@@ -22,6 +22,18 @@
 log = get_logger(__name__)
 
 
+def dest_host_filter(host: str) -> dict:
+    """A host verdict reaches links that hide the host in a geo rule or a
+    variant, not only links whose long_url points there. The legacy urls and
+    emojis collections carry no rules, so their host-only queries stay
+    complete without this."""
+    return {"$or": [{"dest.host": host}, {"dest.secondary_hosts": host}]}
+
+
+# Every host a link can route to, for grouping: main plus secondary.
+_ALL_HOSTS = {"$setUnion": [["$dest.host"], {"$ifNull": ["$dest.secondary_hosts", []]}]}
+
+
 # Mongo caps a single query document at 16MB; a host-wide block can name
 # tens of thousands of (alias, domain) pairs, so the $or is chunked.
 _ALIAS_CHUNK = 1_000
@@ -478,7 +490,7 @@ async def list_by_dest_host_with_urls(
         deliberately status-blind: a re-delivered block must still evict
         entries the first attempt flipped but failed to evict."""
         cursor = self._col.find(
-            {"dest.host": host},
+            dest_host_filter(host),
             {"alias": 1, "domain": 1, "long_url": 1},
         ).limit(limit)
         docs = await cursor.to_list(length=limit)
@@ -492,13 +504,24 @@ async def list_by_dest_host_with_urls(
         return [(d["alias"], d.get("domain", ""), d.get("long_url", "")) for d in docs]
 
     async def unblock_by_dest_host(self, host: str) -> int:
-        """Flip BLOCKED links pointing at *host* back to ACTIVE, scoped to
-        docs carrying ``blocked_reason`` so a manual operator ban is never
-        undone. Stamps stay; ``unblocked_at`` records the reversal."""
+        """Remove *host* as a block cause and reactivate the links that have
+        no cause left. A link also blocked for another of its destinations
+        stays BLOCKED; a per-link block (``blocked_hosts`` null) is never
+        touched by a host unblock. Blocks stamped before ``blocked_hosts``
+        existed were host-of-long_url blocks, so they match on ``dest.host``.
+        Scoped to docs carrying ``blocked_reason`` so a manual operator ban is
+        never undone. Stamps stay; ``unblocked_at`` records the reversal."""
         now = datetime.now(timezone.utc)
+        await self._col.update_many(
+            {"blocked_hosts": host, "status": UrlStatus.BLOCKED.value},
+            {"$pull": {"blocked_hosts": host}},
+        )
         result = await self._col.update_many(
             {
-                "dest.host": host,
+                "$or": [
+                    {**dest_host_filter(host), "blocked_hosts": []},
+                    {"blocked_hosts": {"$exists": False}, "dest.host": host},
+                ],
                 "status": UrlStatus.BLOCKED.value,
                 "blocked_reason": {"$exists": True},
             },
@@ -519,7 +542,7 @@ async def list_active_owned_by_dest_host(
         event set (anonymous links have no possible webhook subscriber)."""
         cursor = self._col.find(
             {
-                "dest.host": host,
+                **dest_host_filter(host),
                 "status": UrlStatus.ACTIVE.value,
                 "owner_id": {"$ne": ANONYMOUS_OWNER_ID},
             }
@@ -543,6 +566,8 @@ async def list_active_hosts_by_registrable(
                     "status": UrlStatus.ACTIVE.value,
                 }
             },
+            # Main host only: a secondary host's own registrable domain is not
+            # the matched one, so fanning it out here would tag it wrongly.
             {
                 "$group": {
                     "_id": "$dest.host",
@@ -570,9 +595,11 @@ async def list_recent_destination_hosts(
         since_id = ObjectId.from_datetime(since)
         pipeline = [
             {"$match": {"_id": {"$gte": since_id}, "dest.host": {"$exists": True}}},
+            {"$project": {"hosts": _ALL_HOSTS, "long_url": 1}},
+            {"$unwind": "$hosts"},
             {
                 "$group": {
-                    "_id": "$dest.host",
+                    "_id": "$hosts",
                     "sample_url": {"$first": "$long_url"},
                 }
             },
@@ -623,6 +650,8 @@ async def block_active_by_aliases(
                         "updated_at": now,
                         "blocked_at": now,
                         "blocked_reason": reason,
+                        # Per-link cause: a host unblock never reactivates it.
+                        "blocked_hosts": None,
                     }
                 },
             )
@@ -635,7 +664,7 @@ async def destination_history(self, host: str) -> dict:
         how many links point here, the anon/owned split, total clicks, and
         the earliest sighting. All facts we already own; no network."""
         pipeline = [
-            {"$match": {"dest.host": host}},
+            {"$match": dest_host_filter(host)},
             {
                 "$group": {
                     "_id": None,
@@ -698,7 +727,7 @@ async def host_breadth(self, host: str, *, sample: int = 8) -> dict:
         benefit of the doubt.
         """
         pipeline = [
-            {"$match": {"dest.host": host}},
+            {"$match": dest_host_filter(host)},
             {
                 "$group": {
                     "_id": None,
@@ -748,16 +777,58 @@ async def block_active_by_dest_host(self, host: str, *, reason: str) -> int:
         audit trail — ``updated_at`` is lossy, these survive later edits."""
         now = datetime.now(timezone.utc)
         result = await self._col.update_many(
-            {"dest.host": host, "status": UrlStatus.ACTIVE.value},
+            {**dest_host_filter(host), "status": UrlStatus.ACTIVE.value},
             {
                 "$set": {
                     "status": UrlStatus.BLOCKED.value,
                     "updated_at": now,
                     "blocked_at": now,
                     "blocked_reason": reason,
+                    "blocked_hosts": [host],
                 }
             },
         )
+        # Already-blocked links gain this host as a second cause; a pre-field
+        # block's first cause is dest.host. Per-link blocks (null) are skipped.
+        await self._col.update_many(
+            {
+                **dest_host_filter(host),
+                "status": UrlStatus.BLOCKED.value,
+                "blocked_reason": {"$exists": True},
+                "$or": [
+                    {"blocked_hosts": {"$type": "array"}},
+                    {"blocked_hosts": {"$exists": False}},
+                ],
+            },
+            [
+                {
+                    "$set": {
+                        "blocked_hosts": {
+                            "$setUnion": [
+                                {
+                                    "$ifNull": [
+                                        "$blocked_hosts",
+                                        {
+                                            "$cond": [
+                                                {
+                                                    "$eq": [
+                                                        {"$type": "$dest.host"},
+                                                        "string",
+                                                    ]
+                                                },
+                                                ["$dest.host"],
+                                                [],
+                                            ]
+                                        },
+                                    ]
+                                },
+                                [host],
+                            ]
+                        }
+                    }
+                }
+            ],
+        )
         return int(result.modified_count)
 
     async def increment_clicks(

schemas/models/url.py

@@ -21,7 +21,7 @@
 
 from schemas.models.base import ANONYMOUS_OWNER_ID, MongoBaseModel, PyObjectId
 from shared.datetime_utils import as_aware_utc
-from shared.url_utils import parse_destination
+from shared.url_utils import link_destination_urls, parse_destination, secondary_hosts
 
 _CONTROL_CHARS_RE = re.compile(r"[\x00-\x1f\x7f]")
 
@@ -143,12 +143,50 @@ class UrlDestination(BaseModel):
     host: str = ""
     subdomain: str = ""
     registrable_domain: str = ""
+    # Hosts the link can ALSO route to (geo rules today, A/B variants
+    # later), main host excluded. Enforcement matches host OR these.
+    secondary_hosts: list[str] = Field(default_factory=list)
 
     @classmethod
     def from_url(cls, url: str | None) -> UrlDestination | None:
         parts = parse_destination(url)
         return cls(**parts) if parts else None
 
+    @classmethod
+    def for_link(
+        cls,
+        long_url: str | None,
+        *,
+        geo_rules: dict[str, str] | None = None,
+        variants: list[str] | None = None,
+        pre_start_url: str | None = None,
+    ) -> UrlDestination | None:
+        """Main parts plus every secondary host. None only when nothing
+        about the link parses."""
+        main = parse_destination(long_url)
+        extra = secondary_hosts(
+            link_destination_urls(
+                None,
+                geo_rules=geo_rules,
+                variants=variants,
+                pre_start_url=pre_start_url,
+            ),
+            exclude=(main or {}).get("host", ""),
+        )
+        if main is None and not extra:
+            return None
+        return cls(**(main or {}), secondary_hosts=extra)
+
+    def to_doc(self) -> dict:
+        """Mongo shape: secondary_hosts absent when empty so the sparse index
+        holds only links that actually have them. The backfill is the one
+        writer of an empty list, as its convergence marker for geo links
+        whose rules add no host (see scripts/backfill_url_dest.py)."""
+        data = self.model_dump()
+        if not data["secondary_hosts"]:
+            data.pop("secondary_hosts")
+        return data
+
 
 class UrlV2Doc(MongoBaseModel):
     """Document model for the `urlsV2` collection.
@@ -217,6 +255,9 @@ def _normalise_domain(cls, v: Any) -> str:
     # deliberately write no host-wide verdict.
     blocked_at: datetime | None = None
     blocked_reason: str | None = None
+    # Host causes of the block; reactivates only when the last one is removed.
+    # None marks a per-link block that no host unblock may undo.
+    blocked_hosts: list[str] | None = None
     # Stamped on reversal; blocked_at/blocked_reason stay.
     unblocked_at: datetime | None = None
     private_stats: bool | None = True  # None for anonymous/unowned URLs