Patched code sample

import socket
import ipaddress
import requests
from urllib.parse import urlparse

def is_safe_url(url: str) -> bool:
    """
    Validates a URL to prevent SSRF attacks.

    This function represents the core of the fix. It ensures that the URL:
    1. Uses an allowed scheme (http, https).
    2. Resolves to a public IP address, not a private, reserved, loopback,
       or link-local address.

    Args:
        url: The URL to validate.

    Returns:
        True if the URL is safe, False otherwise.
    """
    try:
        parsed_url = urlparse(url)

        # 1. Scheme validation: Only allow http and https
        if parsed_url.scheme not in ('http', 'https'):
            return False

        # 2. Hostname validation: Ensure a hostname is present
        hostname = parsed_url.hostname
        if not hostname:
            return False

        # 3. IP address resolution and validation
        # Use socket.getaddrinfo for robust resolution (supports IPv4/IPv6)
        addr_info = socket.getaddrinfo(hostname, parsed_url.port or 0)
        
        # In case of multiple IPs (e.g., DNS round-robin), check all of them
        for family, socktype, proto, canonname, sockaddr in addr_info:
            ip_str = sockaddr[0]
            ip = ipaddress.ip_address(ip_str)

            # The crucial security check: block non-public IPs
            if not ip.is_global:
                # is_global is a convenience method that checks if the address is not
                # private, reserved, loopback, or link-local. This is the main
                # defense against SSRF attacks targeting internal resources.
                return False

        # If all resolved IPs are public/global, the URL is considered safe
        return True

    except (ValueError, socket.gaierror):
        # ValueError for invalid URL, gaierror for DNS resolution failure
        return False


def patched_file_downloader(url: str) -> bytes:
    """
    Represents the fixed file download functionality in BentoML.

    It first calls is_safe_url to validate the target URL before making
    any network request. This prevents the server from being forced to
    request resources from internal or restricted network locations.
    """
    # This is the security control point. The vulnerable version would
    # lack this check and proceed directly to the requests.get() call.
    if not is_safe_url(url):
        raise ValueError(f"URL validation failed. Access to '{url}' is blocked for security reasons.")

    print(f"[INFO] URL '{url}' passed validation. Proceeding with download.")
    try:
        # A timeout is crucial for any production-grade network request
        with requests.get(url, stream=True, timeout=10) as response:
            response.raise_for_status()
            # For demonstration, we'll return the content.
            # A real implementation might stream the file to disk.
            return response.content
    except requests.exceptions.RequestException as e:
        raise IOError(f"Failed to fetch file from URL: {e}") from e

# --- Demonstration of the fix ---
if __name__ == "__main__":
    # A list of URLs to test the fix against.
    test_urls = {
        # Potentially vulnerable URLs that should be blocked
        "http://127.0.0.1/server-info": False,  # Loopback (IPv4)
        "http://[::1]/admin": False,  # Loopback (IPv6)
        "http://localhost/secrets.txt": False, # Resolves to loopback
        "http://169.254.169.254/latest/meta-data/": False,  # AWS/GCP Metadata Service
        "http://192.168.1.1/router_status.html": False,  # Private network
        "http://10.0.0.5/internal_api": False, # Private network
        "file:///etc/passwd": False, # Disallowed scheme
        "ftp://example.com/data": False, # Disallowed scheme

        # Safe, public URLs that should be allowed
        "https://www.bentoml.com/": True,
        "http://example.com/": True,
    }

    print("--- Testing the Patched Downloader with SSRF Prevention ---")
    for url, should_succeed in test_urls.items():
        print(f"\nAttempting to fetch: {url}")
        try:
            patched_file_downloader(url)
            if should_succeed:
                print(f"[SUCCESS] Correctly allowed safe URL.")
            else:
                print(f"[FAILURE] INCORRECTLY allowed a malicious URL.")
        except (ValueError, IOError) as e:
            if not should_succeed:
                print(f"[SUCCESS] Correctly blocked malicious URL. Reason: {e}")
            else:
                print(f"[FAILURE] INCORRECTLY blocked a safe URL. Reason: {e}")