feat: Enhance URL safety checks and integrate is_safe_url across image handling #3397

Author: hlohaus

etc/unittest/backend.py

@@ -2,7 +2,8 @@
 
 import unittest
 import asyncio
-from unittest.mock import MagicMock
+import socket
+from unittest.mock import MagicMock, patch
 from g4f.errors import MissingRequirementsError
 try:
     from g4f.gui.server.backend_api import Backend_Api
@@ -43,6 +44,24 @@ def test_get_providers(self):
         self.assertIsInstance(response, list)
         self.assertTrue(len(response) > 0)
 
+    @patch('g4f.gui.server.backend_api.socket.getaddrinfo')
+    def test_is_safe_url_with_backslash_confusion(self, mock_getaddrinfo):
+        mock_getaddrinfo.return_value = [(socket.AF_INET, socket.SOCK_STREAM, 6, '', ('127.0.0.1', 0))]
+        from g4f.gui.server.backend_api import _is_safe_url
+        self.assertFalse(_is_safe_url('http://127.0.0.1:6666\\@www.baidu.com'))
+
+    @patch('g4f.gui.server.backend_api.socket.getaddrinfo')
+    def test_is_safe_url_blocks_private(self, mock_getaddrinfo):
+        mock_getaddrinfo.return_value = [(socket.AF_INET, socket.SOCK_STREAM, 6, '', ('127.0.0.1', 0))]
+        from g4f.gui.server.backend_api import _is_safe_url
+        self.assertFalse(_is_safe_url('http://127.0.0.1'))
+
+    @patch('g4f.gui.server.backend_api.socket.getaddrinfo')
+    def test_is_safe_url_allows_public(self, mock_getaddrinfo):
+        mock_getaddrinfo.return_value = [(socket.AF_INET, socket.SOCK_STREAM, 6, '', ('8.8.8.8', 0))]
+        from g4f.gui.server.backend_api import _is_safe_url
+        self.assertTrue(_is_safe_url('http://example.com'))
+
     def test_search(self):
         if not has_search:
             self.skipTest("import error")

g4f/gui/server/backend_api.py

@@ -11,10 +11,8 @@
 import shutil
 import random
 import datetime
-import ipaddress
-import socket
 from hashlib import sha256
-from urllib.parse import quote_plus, urlparse
+from urllib.parse import quote_plus
 from functools import lru_cache
 from flask import Flask, Response, redirect, request, jsonify, send_from_directory
 from werkzeug.exceptions import NotFound
@@ -46,7 +44,7 @@
 from ...tools.files import supports_filename, get_streaming, get_bucket_dir, get_tempfile
 from ...tools.run_tools import iter_run_tools
 from ...errors import ModelNotFoundError, ProviderNotFoundError, MissingAuthError, RateLimitError
-from ...image import is_allowed_extension, process_image, MEDIA_TYPE_MAP
+from ...image import is_allowed_extension, process_image, MEDIA_TYPE_MAP, is_safe_url as _is_safe_url
 from ...cookies import get_cookies_dir
 from ...image.copy_images import secure_filename, get_source_url, get_media_dir, copy_media
 from ...client.service import get_model_and_provider
@@ -60,30 +58,6 @@
 
 _DATE_RE = re.compile(r'^\d{4}-\d{2}-\d{2}$')
 
-def _is_safe_url(url: str) -> bool:
-    """Return True only for http/https URLs that do not point to private/loopback/reserved addresses."""
-    try:
-        parsed = urlparse(url)
-        if parsed.scheme not in ("http", "https"):
-            return False
-        hostname = parsed.hostname
-        if hostname is None:
-            return False
-        # Resolve all IP addresses for the hostname and reject if any is non-public.
-        # Validating all addresses reduces the window for DNS rebinding attacks.
-        addr_infos = socket.getaddrinfo(hostname, None)
-        if not addr_infos:
-            return False
-        for addr_info in addr_infos:
-            addr = ipaddress.ip_address(addr_info[4][0])
-            if (addr.is_private or addr.is_loopback or addr.is_link_local
-                    or addr.is_reserved or addr.is_multicast or addr.is_unspecified):
-                return False
-    except Exception as e:
-        logger.debug("URL safety check failed for %r: %s", url, e)
-        return False
-    return True
-
 def safe_iter_generator(generator: Generator) -> Generator:
     start = next(generator)
     def iter_generator():

g4f/image/__init__.py

@@ -4,13 +4,20 @@
 import re
 import io
 import base64
+import socket
+import ipaddress
 from io import BytesIO
 from pathlib import Path
 from typing import Optional
 from urllib.parse import urlparse
 
 import requests
 
+try:
+    from urllib3.util import parse_url as urllib3_parse_url
+except ImportError:
+    urllib3_parse_url = None
+
 try:
     from PIL import Image, ImageOps
     has_requirements = True
@@ -103,6 +110,45 @@ def is_allowed_extension(filename: str) -> Optional[str]:
         return None
     return EXTENSIONS_MAP[extension]
 
+
+def is_safe_url(url: str) -> bool:
+    """Return True only for http/https URLs that do not point to private/loopback/reserved addresses."""
+    try:
+        parsed = urlparse(url)
+
+        if parsed.scheme not in ("http", "https"):
+            return False
+
+        if "\\" in url:
+            return False
+
+        hostname = parsed.hostname
+        if hostname is None:
+            return False
+
+        if urllib3_parse_url is not None:
+            parsed_urllib3 = urllib3_parse_url(url)
+            if parsed_urllib3.host and parsed_urllib3.host != hostname:
+                return False
+            hostname = parsed_urllib3.host or hostname
+
+        if hostname is None:
+            return False
+
+        addr_infos = socket.getaddrinfo(hostname, None)
+        if not addr_infos:
+            return False
+
+        for addr_info in addr_infos:
+            addr = ipaddress.ip_address(addr_info[4][0])
+            if (addr.is_private or addr.is_loopback or addr.is_link_local
+                    or addr.is_reserved or addr.is_multicast or addr.is_unspecified):
+                return False
+    except Exception:
+        return False
+    return True
+
+
 def is_data_an_media(data, filename: str = None) -> str:
     content_type = is_data_an_audio(data, filename)
     if content_type is not None:
@@ -378,6 +424,8 @@ def to_bytes(image: ImageType) -> bytes:
             is_data_uri_an_image(image)
             return extract_data_uri(image)
         elif image.startswith("http://") or image.startswith("https://"):
+            if not is_safe_url(image):
+                raise ValueError("Invalid or disallowed media URL")
             path: str = urlparse(image).path
             if path.startswith("/files/"):
                 path = get_bucket_dir(*path.split("/")[2:])

g4f/image/copy_images.py

@@ -16,7 +16,7 @@
 from ..image import MEDIA_TYPE_MAP, EXTENSIONS_MAP
 from ..tools.files import secure_filename
 from ..providers.response import ImageResponse, AudioResponse, VideoResponse, quote_url
-from . import is_accepted_format, extract_data_uri
+from . import is_accepted_format, extract_data_uri, is_safe_url
 from .. import debug
 
 # Directory for storing generated media files
@@ -170,6 +170,8 @@ async def copy_image(image: str, target: str = None) -> str:
                     with open(target_path, "wb") as f:
                         f.write(extract_data_uri(image))
                 elif not os.path.exists(target_path) or os.lstat(target_path).st_size <= 0:
+                    if not is_safe_url(image):
+                        raise ValueError("Invalid or disallowed media URL")
                     # Use aiohttp to fetch the image
                     async with session.get(image, ssl=ssl) as response:
                         response.raise_for_status()