✨ add extensive URL checking

Author: sebastianMindee

mindee/input/url_input_source.py

@@ -10,6 +10,7 @@
 from mindee.error.mindee_error import MindeeSourceError
 from mindee.input.bytes_input import BytesInput
 from mindee.logger import logger
+from mindee.mindee_http.response_validation import validate_url_for_source
 from mindee.parsing.common.string_dict import StringDict
 
 
@@ -25,8 +26,7 @@ def __init__(self, url: str) -> None:
 
         :param url: URL to send, must be HTTPS.
         """
-        if not url.lower().startswith("https"):
-            raise MindeeSourceError("URL must be HTTPS")
+        validate_url_for_source(url)
 
         logger.debug("URL input: %s", url)
 

mindee/mindee_http/cancellation_token.py

@@ -0,0 +1,15 @@
+class CancellationToken:
+    """Custom cancellation token that can be used to cancel a polling request."""
+
+    is_canceled: bool
+    """A cancellation token that can be used to cancel a request."""
+
+    def __init__(
+        self,
+        is_canceled: bool = False,
+    ):
+        self.is_canceled = is_canceled
+
+    def cancel(self):
+        """Cancel the request."""
+        self.is_canceled = True

mindee/mindee_http/response_validation.py

@@ -1,9 +1,79 @@
+import ipaddress
 import json
+from urllib.parse import urlparse
 
 import httpx
 
+from mindee.error.mindee_error import MindeeSourceError
 from mindee.parsing.common.string_dict import StringDict
 
+_CGNAT_BLOCK = ipaddress.IPv4Network("100.64.0.0/10")
+_IPV6_UNIQUE_LOCAL = ipaddress.IPv6Network("fc00::/7")
+
+
+def validate_url_for_source(url: str) -> None:
+    """
+    Validates that a URL is safe to send to the Mindee server.
+
+    Rejects any URL that could be used for Server-Side Request Forgery (SSRF):
+
+    - non-HTTPS schemes,
+    - embedded userinfo (e.g. ``https://user:pass@host``),
+    - loopback hostnames (``localhost``, ``*.localhost``),
+    - literal IP addresses that are loopback, link-local, private (RFC 1918),
+      any-local (``0.0.0.0``), multicast, IPv6 unique-local (``fc00::/7``),
+      or carrier-grade NAT (``100.64.0.0/10``).
+
+    Note: DNS resolution is not performed. A hostname that resolves to a
+    private IP will not be caught here.
+
+    :param url: The URL string to validate.
+    :raises MindeeSourceError: If the URL fails any security check.
+    """
+    try:
+        parsed = urlparse(url)
+    except Exception as exc:
+        raise MindeeSourceError("Invalid URL") from exc
+
+    if parsed.scheme.lower() != "https":
+        raise MindeeSourceError("URL must be HTTPS")
+
+    if parsed.username or parsed.password:
+        raise MindeeSourceError("Source URLs must not embed user credentials")
+
+    host = parsed.hostname
+    if not host:
+        raise MindeeSourceError("Source URL is missing a host")
+
+    lower_host = host.lower()
+    if (
+        lower_host == "localhost"
+        or lower_host.endswith(".localhost")
+        or lower_host == "ip6-localhost"
+        or lower_host == "ip6-loopback"
+    ):
+        raise MindeeSourceError(f"Loopback hostnames are not allowed: {host}")
+
+    try:
+        addr = ipaddress.ip_address(lower_host)
+    except ValueError:
+        return
+
+    if (
+        addr.is_loopback
+        or addr.is_link_local
+        or addr.is_private
+        or addr.is_unspecified
+        or addr.is_multicast
+    ):
+        raise MindeeSourceError(f"URL host resolves to a disallowed address: {addr}")
+
+    if isinstance(addr, ipaddress.IPv4Address) and addr in _CGNAT_BLOCK:
+        raise MindeeSourceError(f"URL host resolves to a disallowed address: {addr}")
+
+    if isinstance(addr, ipaddress.IPv6Address) and addr in _IPV6_UNIQUE_LOCAL:
+        raise MindeeSourceError(f"URL host resolves to a disallowed address: {addr}")
+
 
 def is_valid_sync_response(response: httpx.Response) -> bool:
     """

mindee/v2/client.py

@@ -9,6 +9,7 @@
 from mindee.input import URLInputSource
 from mindee.input.local_input_source import LocalInputSource
 from mindee.logger import logger
+from mindee.mindee_http.cancellation_token import CancellationToken
 from mindee.parsing.common.common_response import CommonStatus
 from mindee.v2.client_options.base_parameters import BaseParameters
 from mindee.v2.mindee_http.mindee_api_v2 import MindeeAPIV2
@@ -104,13 +105,16 @@ def enqueue_and_get_result(
         response_type: type[TypeBaseResponse],
         input_source: LocalInputSource | URLInputSource,
         params: BaseParameters,
+        cancellation_token: CancellationToken | None = None,
     ) -> TypeBaseResponse:
         """
         Enqueues to an asynchronous endpoint and automatically polls for a response.
 
         :param input_source: The document/source file to use. Can be local or remote.
         :param params: Parameters to set when sending a file.
         :param response_type: The product class to use for the response object.
+        :param cancellation_token: A cancellation token that can be used to cancel the
+        request.
 
         :return: A valid inference response.
         """
@@ -128,6 +132,8 @@ def enqueue_and_get_result(
         sleep(params.polling_options.initial_delay_sec)
         try_counter = 0
         while try_counter < params.polling_options.max_retries:
+            if cancellation_token and cancellation_token.is_canceled:
+                raise MindeeError("Request canceled through cancellation token.")
             job_response = self.get_job(enqueue_response.job.id)
             assert isinstance(job_response, JobResponse)
             if job_response.job.status == CommonStatus.FAILED.value:

tests/data

@@ -27,7 +27,6 @@ def check_findoc_return(findoc_response: ExtractionResponse):
 @pytest.mark.pypdfium2
 @pytest.mark.integration
 def test_pdf_should_extract_splits():
-
     client = Client()
     split_input = PathInput(V2_PRODUCT_DATA_DIR / "split" / "default_sample.pdf")
     response = client.enqueue_and_get_result(