From a337b32bcd814fa81033a3f86b336069e842f257 Mon Sep 17 00:00:00 2001
From: Thomas Sileo <t@a4.io>
Date: Sun, 4 Dec 2022 11:51:52 +0100
Subject: [PATCH] Blocking server also blocks subdomains

---
 app/boxes.py        |  4 ++--
 app/httpsig.py      |  4 ++--
 app/utils/url.py    | 10 +++++++++-
 tests/test_utils.py | 19 +++++++++++++++++++
 4 files changed, 32 insertions(+), 5 deletions(-)
 create mode 100644 tests/test_utils.py

diff --git a/app/boxes.py b/app/boxes.py
index adcfc0f..1fd7e34 100644
--- a/app/boxes.py
+++ b/app/boxes.py
@@ -28,7 +28,6 @@ from app.actor import save_actor
 from app.actor import update_actor_if_needed
 from app.ap_object import RemoteObject
 from app.config import BASE_URL
-from app.config import BLOCKED_SERVERS
 from app.config import ID
 from app.config import MANUALLY_APPROVES_FOLLOWERS
 from app.config import set_moved_to
@@ -46,6 +45,7 @@ from app.utils.datetime import now
 from app.utils.datetime import parse_isoformat
 from app.utils.facepile import WebmentionReply
 from app.utils.text import slugify
+from app.utils.url import is_hostname_blocked
 
 AnyboxObject = models.InboxObject | models.OutboxObject
 
@@ -2312,7 +2312,7 @@ async def save_to_inbox(
         logger.exception("Failed to fetch actor")
         return
 
-    if actor.server in BLOCKED_SERVERS:
+    if is_hostname_blocked(actor.server):
         logger.warning(f"Server {actor.server} is blocked")
         return
 
diff --git a/app/httpsig.py b/app/httpsig.py
index 9f8ad79..e9f9e69 100644
--- a/app/httpsig.py
+++ b/app/httpsig.py
@@ -23,12 +23,12 @@ from sqlalchemy import select
 
 from app import activitypub as ap
 from app import config
-from app.config import BLOCKED_SERVERS
 from app.config import KEY_PATH
 from app.database import AsyncSession
 from app.database import get_db_session
 from app.key import Key
 from app.utils.datetime import now
+from app.utils.url import is_hostname_blocked
 
 _KEY_CACHE: MutableMapping[str, Key] = LFUCache(256)
 
@@ -184,7 +184,7 @@ async def httpsig_checker(
         )
 
     server = urlparse(key_id).hostname
-    if server in BLOCKED_SERVERS:
+    if is_hostname_blocked(server):
         return HTTPSigInfo(
             has_valid_signature=False,
             server=server,
diff --git a/app/utils/url.py b/app/utils/url.py
index 80630b2..92102c3 100644
--- a/app/utils/url.py
+++ b/app/utils/url.py
@@ -54,7 +54,7 @@ def is_url_valid(url: str) -> bool:
     if not parsed.hostname or parsed.hostname.lower() in ["localhost"]:
         return False
 
-    if parsed.hostname in BLOCKED_SERVERS:
+    if is_hostname_blocked(parsed.hostname):
         logger.warning(f"{parsed.hostname} is blocked")
         return False
 
@@ -81,3 +81,11 @@ def check_url(url: str) -> None:
         raise InvalidURLError(f'"{url}" is invalid')
 
     return None
+
+
+@functools.lru_cache(maxsize=256)
+def is_hostname_blocked(hostname: str) -> bool:
+    for blocked_hostname in BLOCKED_SERVERS:
+        if hostname == blocked_hostname or hostname.endswith(f".{blocked_hostname}"):
+            return True
+    return False
diff --git a/tests/test_utils.py b/tests/test_utils.py
new file mode 100644
index 0000000..ec478d6
--- /dev/null
+++ b/tests/test_utils.py
@@ -0,0 +1,19 @@
+from unittest import mock
+
+import pytest
+
+from app.utils.url import is_hostname_blocked
+
+
+@pytest.mark.parametrize(
+    "hostname,should_be_blocked",
+    [
+        ("example.com", True),
+        ("subdomain.example.com", True),
+        ("example.xyz", False),
+    ],
+)
+def test_is_hostname_blocked(hostname: str, should_be_blocked: bool) -> None:
+    with mock.patch("app.utils.url.BLOCKED_SERVERS", ["example.com"]):
+        is_hostname_blocked.cache_clear()
+        assert is_hostname_blocked(hostname) is should_be_blocked