1 files changed, 16 insertions, 49 deletions

diff --git a/webmentions_ssg/tasks/receiver.py b/webmentions_ssg/tasks/receiver.py
index 9475866..39def73 100644
--- a/webmentions_ssg/tasks/receiver.py
+++ b/webmentions_ssg/tasks/receiver.py
@@ -10,11 +10,7 @@ from .. import APP_NAME, VERSION
 from .. import DATABASE as db
 from .. import HUEY as huey
 from ..models import ReceivedWebmention
-from ..url_security import (
-    AddressResolutionError,
-    NonPublicAddressError,
-    ensure_public_url,
-)
+from ..url_security import AddressResolutionError, is_public_url
 
 
 class VerificationError(Exception):
@@ -46,33 +42,20 @@ HTML_URL_ATTRIBUTES = {
         "track",
         "video",
     },
-    "cite": {
-        "blockquote",
-        "del",
-        "ins",
-        "q",
-    },
+    "cite": {"blockquote", "del", "ins", "q"},
 }
 
 
-def html_mentions_target(
-    body: bytes,
-    source_url: str,
-    target_url: str,
-) -> bool:
+def html_mentions_target(body: bytes, source_url: str, target_url: str) -> bool:
     """Check valid HTML URL attributes for the exact target URL."""
 
     document = BeautifulSoup(body, "html.parser")
     base_url = source_url
 
     if (base_element := document.select_one("base[href]")) is not None and isinstance(
-        base_href := base_element.get("href"),
-        str,
+        base_href := base_element.get("href"), str
     ):
-        base_url = urljoin(
-            source_url,
-            base_href.strip(),
-        )
+        base_url = urljoin(source_url, base_href.strip())
 
     for attribute, selectors in HTML_URL_ATTRIBUTES.items():
         selector = ", ".join(
@@ -85,19 +68,10 @@ def html_mentions_target(
         )
 
         for element in document.select(selector):
-            if not isinstance(
-                reference := element.get(attribute),
-                str,
-            ):
+            if not isinstance(reference := element.get(attribute), str):
                 continue
 
-            if (
-                urljoin(
-                    base_url,
-                    reference.strip(),
-                )
-                == target_url
-            ):
+            if urljoin(base_url, reference.strip()) == target_url:
                 return True
 
     return False
@@ -111,9 +85,10 @@ def text_mentions_target(body: str, target_url: str) -> bool:
 def ensure_public_request(request: httpx.Request) -> None:
     """Prevent requests to non-public network addresses."""
     try:
-        ensure_public_url(str(request.url))
-    except NonPublicAddressError as exc:
-        raise VerificationError("Source resolves to a non-public address") from exc
+        if not is_public_url(str(request.url)):
+            raise VerificationError("Source resolves to a non-public address")
+    except ValueError as exc:
+        raise VerificationError("Source URL has no hostname") from exc
     except AddressResolutionError as exc:
         raise TemporaryFetchError("Source hostname could not be resolved") from exc
 
@@ -132,9 +107,7 @@ def fetch_source(source_url: str) -> tuple[httpx.Response, bytes]:
         follow_redirects=True,
         max_redirects=current_app.config.get("WEBMENTIONS_SSG_MAX_REDIRECTS", 20),
         trust_env=False,
-        event_hooks={
-            "request": [ensure_public_request],
-        },
+        event_hooks={"request": [ensure_public_request]},
     ) as client:
         with client.stream("GET", source_url) as response:
             match response.status_code:
@@ -149,8 +122,7 @@ def fetch_source(source_url: str) -> tuple[httpx.Response, bytes]:
                         raise VerificationError(f"Source returned HTTP {status}")
 
             max_source_bytes = current_app.config.get(
-                "WEBMENTIONS_SSG_MAX_SOURCE_BYTES",
-                1_000_000,
+                "WEBMENTIONS_SSG_MAX_SOURCE_BYTES", 1_000_000
             )
 
             if (content_length := response.headers.get("Content-Length")) is not None:
@@ -185,14 +157,10 @@ def source_mentions_target(source_url: str, target_url: str) -> bool:
         case "text/plain":
             try:
                 decoded_body = body.decode(
-                    response.encoding or "utf-8",
-                    errors="replace",
+                    response.encoding or "utf-8", errors="replace"
                 )
             except LookupError:
-                decoded_body = body.decode(
-                    "utf-8",
-                    errors="replace",
-                )
+                decoded_body = body.decode("utf-8", errors="replace")
             return text_mentions_target(decoded_body, target_url)
         case _:
             raise VerificationError(
@@ -208,8 +176,7 @@ def verify_webmention(webmention_uuid: uuid.UUID) -> None:
 
     if webmention is None:
         current_app.logger.warning(
-            "Cannot verify unknown ReceivedWebmention %s",
-            webmention_uuid,
+            "Cannot verify unknown ReceivedWebmention %s", webmention_uuid
         )
         return