mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-17 02:59:11 +02:00
447 lines
17 KiB
Python
447 lines
17 KiB
Python
"""Tests for SSRF-safe HTTP utilities.
|
|
|
|
This module tests the ssrf.py module which provides SSRF-protected HTTP fetching.
|
|
"""
|
|
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import httpx
|
|
import pytest
|
|
|
|
from fastmcp.server.auth.ssrf import (
|
|
SSRFError,
|
|
SSRFFetchError,
|
|
is_ip_allowed,
|
|
ssrf_safe_fetch,
|
|
validate_url,
|
|
)
|
|
|
|
|
|
class TestIsIPAllowed:
|
|
"""Tests for is_ip_allowed function."""
|
|
|
|
def test_public_ipv4_allowed(self):
|
|
"""Public IPv4 addresses should be allowed."""
|
|
assert is_ip_allowed("8.8.8.8") is True
|
|
assert is_ip_allowed("1.1.1.1") is True
|
|
assert is_ip_allowed("93.184.216.34") is True
|
|
|
|
def test_private_ipv4_blocked(self):
|
|
"""Private IPv4 addresses should be blocked."""
|
|
assert is_ip_allowed("192.168.1.1") is False
|
|
assert is_ip_allowed("10.0.0.1") is False
|
|
assert is_ip_allowed("172.16.0.1") is False
|
|
|
|
def test_loopback_blocked(self):
|
|
"""Loopback addresses should be blocked."""
|
|
assert is_ip_allowed("127.0.0.1") is False
|
|
assert is_ip_allowed("::1") is False
|
|
|
|
def test_link_local_blocked(self):
|
|
"""Link-local addresses (AWS metadata) should be blocked."""
|
|
assert is_ip_allowed("169.254.169.254") is False
|
|
|
|
def test_rfc6598_cgnat_blocked(self):
|
|
"""RFC6598 Carrier-Grade NAT addresses should be blocked."""
|
|
assert is_ip_allowed("100.64.0.1") is False
|
|
assert is_ip_allowed("100.100.100.100") is False
|
|
|
|
def test_ipv4_mapped_ipv6_blocked_if_private(self):
|
|
"""IPv4-mapped IPv6 addresses should check the embedded IPv4."""
|
|
assert is_ip_allowed("::ffff:127.0.0.1") is False
|
|
assert is_ip_allowed("::ffff:192.168.1.1") is False
|
|
|
|
|
|
class TestValidateURL:
|
|
"""Tests for validate_url function."""
|
|
|
|
async def test_http_rejected(self):
|
|
"""HTTP URLs should be rejected (HTTPS required)."""
|
|
with pytest.raises(SSRFError, match="must use HTTPS"):
|
|
await validate_url("http://example.com/path")
|
|
|
|
async def test_missing_host_rejected(self):
|
|
"""URLs without host should be rejected."""
|
|
with pytest.raises(SSRFError, match="must have a host"):
|
|
await validate_url("https:///path")
|
|
|
|
async def test_root_path_rejected_when_required(self):
|
|
"""Root paths should be rejected when require_path=True."""
|
|
with patch(
|
|
"fastmcp.server.auth.ssrf.resolve_hostname",
|
|
return_value=["93.184.216.34"],
|
|
):
|
|
with pytest.raises(SSRFError, match="non-root path"):
|
|
await validate_url("https://example.com/", require_path=True)
|
|
|
|
async def test_private_ip_rejected(self):
|
|
"""URLs resolving to private IPs should be rejected."""
|
|
with patch(
|
|
"fastmcp.server.auth.ssrf.resolve_hostname",
|
|
return_value=["192.168.1.1"],
|
|
):
|
|
with pytest.raises(SSRFError, match="blocked IP"):
|
|
await validate_url("https://example.com/path")
|
|
|
|
|
|
class TestSSRFSafeFetch:
|
|
"""Tests for ssrf_safe_fetch function."""
|
|
|
|
async def test_private_ip_blocked(self):
|
|
"""Fetch to private IP should be blocked."""
|
|
with patch(
|
|
"fastmcp.server.auth.ssrf.resolve_hostname",
|
|
return_value=["192.168.1.1"],
|
|
):
|
|
with pytest.raises(SSRFError, match="blocked IP"):
|
|
await ssrf_safe_fetch("https://internal.example.com/api")
|
|
|
|
async def test_cgnat_blocked(self):
|
|
"""Fetch to RFC6598 CGNAT IP should be blocked."""
|
|
with patch(
|
|
"fastmcp.server.auth.ssrf.resolve_hostname",
|
|
return_value=["100.64.0.1"],
|
|
):
|
|
with pytest.raises(SSRFError, match="blocked IP"):
|
|
await ssrf_safe_fetch("https://cgnat.example.com/api")
|
|
|
|
async def test_connects_to_pinned_ip(self):
|
|
"""Verify connection uses pinned IP, not re-resolved DNS."""
|
|
resolved_ip = "93.184.216.34"
|
|
|
|
with (
|
|
patch(
|
|
"fastmcp.server.auth.ssrf.resolve_hostname",
|
|
return_value=[resolved_ip],
|
|
),
|
|
patch("httpx.AsyncClient") as mock_client_class,
|
|
):
|
|
mock_stream = MagicMock()
|
|
mock_stream.status_code = 200
|
|
mock_stream.headers = {"content-length": "15"}
|
|
mock_stream.__aenter__ = AsyncMock(return_value=mock_stream)
|
|
mock_stream.__aexit__ = AsyncMock(return_value=None)
|
|
|
|
async def aiter_bytes():
|
|
yield b'{"data": "test"}'
|
|
|
|
mock_stream.aiter_bytes = aiter_bytes
|
|
|
|
mock_client = AsyncMock()
|
|
mock_client.stream = MagicMock(return_value=mock_stream)
|
|
mock_client.__aenter__.return_value = mock_client
|
|
mock_client.__aexit__ = AsyncMock(return_value=None)
|
|
mock_client_class.return_value = mock_client
|
|
|
|
await ssrf_safe_fetch("https://example.com/api")
|
|
|
|
# Verify URL contains pinned IP
|
|
call_args = mock_client.stream.call_args
|
|
url_called = call_args[0][1]
|
|
assert resolved_ip in url_called
|
|
|
|
async def test_fallback_to_second_ip(self):
|
|
"""If the first IP fails, the next resolved IP should be tried."""
|
|
resolved_ips = ["2001:4860:4860::8888", "93.184.216.34"]
|
|
|
|
with (
|
|
patch(
|
|
"fastmcp.server.auth.ssrf.resolve_hostname",
|
|
return_value=resolved_ips,
|
|
),
|
|
patch("httpx.AsyncClient") as mock_client_class,
|
|
):
|
|
request = httpx.Request("GET", "https://example.com/api")
|
|
|
|
first_client = AsyncMock()
|
|
first_client.stream = MagicMock(
|
|
side_effect=httpx.RequestError("boom", request=request)
|
|
)
|
|
first_client.__aenter__.return_value = first_client
|
|
first_client.__aexit__ = AsyncMock(return_value=None)
|
|
|
|
mock_stream = MagicMock()
|
|
mock_stream.status_code = 200
|
|
mock_stream.headers = {"content-length": "2"}
|
|
mock_stream.__aenter__ = AsyncMock(return_value=mock_stream)
|
|
mock_stream.__aexit__ = AsyncMock(return_value=None)
|
|
|
|
async def aiter_bytes():
|
|
yield b"ok"
|
|
|
|
mock_stream.aiter_bytes = aiter_bytes
|
|
|
|
second_client = AsyncMock()
|
|
second_client.stream = MagicMock(return_value=mock_stream)
|
|
second_client.__aenter__.return_value = second_client
|
|
second_client.__aexit__ = AsyncMock(return_value=None)
|
|
|
|
mock_client_class.side_effect = [first_client, second_client]
|
|
|
|
content = await ssrf_safe_fetch("https://example.com/api")
|
|
assert content == b"ok"
|
|
|
|
call_args = second_client.stream.call_args
|
|
url_called = call_args[0][1]
|
|
assert resolved_ips[1] in url_called
|
|
|
|
async def test_host_header_set(self):
|
|
"""Verify Host header is set to original hostname."""
|
|
resolved_ip = "93.184.216.34"
|
|
original_host = "example.com"
|
|
|
|
with (
|
|
patch(
|
|
"fastmcp.server.auth.ssrf.resolve_hostname",
|
|
return_value=[resolved_ip],
|
|
),
|
|
patch("httpx.AsyncClient") as mock_client_class,
|
|
):
|
|
mock_stream = MagicMock()
|
|
mock_stream.status_code = 200
|
|
mock_stream.headers = {"content-length": "15"}
|
|
mock_stream.__aenter__ = AsyncMock(return_value=mock_stream)
|
|
mock_stream.__aexit__ = AsyncMock(return_value=None)
|
|
|
|
async def aiter_bytes():
|
|
yield b'{"data": "test"}'
|
|
|
|
mock_stream.aiter_bytes = aiter_bytes
|
|
|
|
mock_client = AsyncMock()
|
|
mock_client.stream = MagicMock(return_value=mock_stream)
|
|
mock_client.__aenter__.return_value = mock_client
|
|
mock_client.__aexit__ = AsyncMock(return_value=None)
|
|
mock_client_class.return_value = mock_client
|
|
|
|
await ssrf_safe_fetch(f"https://{original_host}/api")
|
|
|
|
# Verify Host header
|
|
call_kwargs = mock_client.stream.call_args[1]
|
|
assert call_kwargs["headers"]["Host"] == original_host
|
|
|
|
async def test_response_size_limit(self):
|
|
"""Verify response size limit is enforced via streaming."""
|
|
with (
|
|
patch(
|
|
"fastmcp.server.auth.ssrf.resolve_hostname",
|
|
return_value=["93.184.216.34"],
|
|
),
|
|
patch("httpx.AsyncClient") as mock_client_class,
|
|
):
|
|
# Response larger than default 5KB (no Content-Length, so streaming enforces)
|
|
mock_stream = MagicMock()
|
|
mock_stream.status_code = 200
|
|
mock_stream.headers = {} # No Content-Length to force streaming check
|
|
mock_stream.__aenter__ = AsyncMock(return_value=mock_stream)
|
|
mock_stream.__aexit__ = AsyncMock(return_value=None)
|
|
|
|
async def aiter_bytes():
|
|
# Yield 10KB total
|
|
for _ in range(10):
|
|
yield b"x" * 1024
|
|
|
|
mock_stream.aiter_bytes = aiter_bytes
|
|
|
|
mock_client = AsyncMock()
|
|
mock_client.stream = MagicMock(return_value=mock_stream)
|
|
mock_client.__aenter__.return_value = mock_client
|
|
mock_client.__aexit__ = AsyncMock(return_value=None)
|
|
mock_client_class.return_value = mock_client
|
|
|
|
with pytest.raises(SSRFFetchError, match="too large"):
|
|
await ssrf_safe_fetch("https://example.com/api")
|
|
|
|
|
|
class TestJWKSSSRFProtection:
|
|
"""Tests for SSRF protection in JWTVerifier JWKS fetching."""
|
|
|
|
async def test_jwks_private_ip_blocked(self):
|
|
"""JWKS fetch to private IP should be blocked."""
|
|
from fastmcp.server.auth.providers.jwt import JWTVerifier
|
|
|
|
verifier = JWTVerifier(
|
|
jwks_uri="https://internal.example.com/.well-known/jwks.json",
|
|
issuer="https://issuer.example.com",
|
|
ssrf_safe=True,
|
|
)
|
|
|
|
with patch(
|
|
"fastmcp.server.auth.ssrf.resolve_hostname",
|
|
return_value=["192.168.1.1"],
|
|
):
|
|
with pytest.raises(ValueError, match="Failed to fetch JWKS"):
|
|
# Create a dummy token to trigger JWKS fetch
|
|
await verifier._get_jwks_key("test-kid")
|
|
|
|
async def test_jwks_cgnat_blocked(self):
|
|
"""JWKS fetch to RFC6598 CGNAT IP should be blocked."""
|
|
from fastmcp.server.auth.providers.jwt import JWTVerifier
|
|
|
|
verifier = JWTVerifier(
|
|
jwks_uri="https://cgnat.example.com/.well-known/jwks.json",
|
|
issuer="https://issuer.example.com",
|
|
ssrf_safe=True,
|
|
)
|
|
|
|
with patch(
|
|
"fastmcp.server.auth.ssrf.resolve_hostname",
|
|
return_value=["100.64.0.1"],
|
|
):
|
|
with pytest.raises(ValueError, match="Failed to fetch JWKS"):
|
|
await verifier._get_jwks_key("test-kid")
|
|
|
|
async def test_jwks_loopback_blocked(self):
|
|
"""JWKS fetch to loopback should be blocked."""
|
|
from fastmcp.server.auth.providers.jwt import JWTVerifier
|
|
|
|
verifier = JWTVerifier(
|
|
jwks_uri="https://localhost/.well-known/jwks.json",
|
|
issuer="https://issuer.example.com",
|
|
ssrf_safe=True,
|
|
)
|
|
|
|
with patch(
|
|
"fastmcp.server.auth.ssrf.resolve_hostname",
|
|
return_value=["127.0.0.1"],
|
|
):
|
|
with pytest.raises(ValueError, match="Failed to fetch JWKS"):
|
|
await verifier._get_jwks_key("test-kid")
|
|
|
|
|
|
class TestIPv6URLFormatting:
|
|
"""Tests for proper IPv6 address bracketing in URLs."""
|
|
|
|
def test_format_ip_for_url_ipv4(self):
|
|
"""IPv4 addresses should not be bracketed."""
|
|
from fastmcp.server.auth.ssrf import format_ip_for_url
|
|
|
|
assert format_ip_for_url("8.8.8.8") == "8.8.8.8"
|
|
assert format_ip_for_url("192.168.1.1") == "192.168.1.1"
|
|
|
|
def test_format_ip_for_url_ipv6(self):
|
|
"""IPv6 addresses should be bracketed for URL use."""
|
|
from fastmcp.server.auth.ssrf import format_ip_for_url
|
|
|
|
assert format_ip_for_url("2001:db8::1") == "[2001:db8::1]"
|
|
assert format_ip_for_url("::1") == "[::1]"
|
|
assert format_ip_for_url("fe80::1") == "[fe80::1]"
|
|
|
|
def test_format_ip_for_url_invalid(self):
|
|
"""Invalid IP strings should be returned unchanged."""
|
|
from fastmcp.server.auth.ssrf import format_ip_for_url
|
|
|
|
assert format_ip_for_url("not-an-ip") == "not-an-ip"
|
|
assert format_ip_for_url("") == ""
|
|
|
|
async def test_ipv6_pinned_url_is_valid(self):
|
|
"""Verify IPv6 addresses are properly bracketed in pinned URLs."""
|
|
resolved_ipv6 = "2001:4860:4860::8888"
|
|
|
|
with (
|
|
patch(
|
|
"fastmcp.server.auth.ssrf.resolve_hostname",
|
|
return_value=[resolved_ipv6],
|
|
),
|
|
patch("httpx.AsyncClient") as mock_client_class,
|
|
):
|
|
mock_stream = MagicMock()
|
|
mock_stream.status_code = 200
|
|
mock_stream.headers = {"content-length": "10"}
|
|
mock_stream.__aenter__ = AsyncMock(return_value=mock_stream)
|
|
mock_stream.__aexit__ = AsyncMock(return_value=None)
|
|
|
|
async def aiter_bytes():
|
|
yield b'{"key": 1}'
|
|
|
|
mock_stream.aiter_bytes = aiter_bytes
|
|
|
|
mock_client = AsyncMock()
|
|
mock_client.stream = MagicMock(return_value=mock_stream)
|
|
mock_client.__aenter__.return_value = mock_client
|
|
mock_client.__aexit__ = AsyncMock(return_value=None)
|
|
mock_client_class.return_value = mock_client
|
|
|
|
await ssrf_safe_fetch("https://example.com/api")
|
|
|
|
# Verify the URL contains bracketed IPv6 address
|
|
call_args = mock_client.stream.call_args
|
|
url_called = call_args[0][1]
|
|
|
|
# IPv6 should be bracketed: https://[2001:4860:4860::8888]:443/path
|
|
assert f"[{resolved_ipv6}]" in url_called, (
|
|
f"Expected bracketed IPv6 [{resolved_ipv6}] in URL, got {url_called}"
|
|
)
|
|
|
|
|
|
class TestStreamingResponseSizeLimit:
|
|
"""Tests for streaming-based response size enforcement."""
|
|
|
|
async def test_size_limit_enforced_during_streaming(self):
|
|
"""Verify that size limit is enforced as chunks are received, not after."""
|
|
with (
|
|
patch(
|
|
"fastmcp.server.auth.ssrf.resolve_hostname",
|
|
return_value=["93.184.216.34"],
|
|
),
|
|
patch("httpx.AsyncClient") as mock_client_class,
|
|
):
|
|
chunks_yielded = []
|
|
|
|
async def aiter_bytes():
|
|
# Yield chunks that exceed the limit
|
|
for i in range(10):
|
|
chunk = b"x" * 1024 # 1KB per chunk
|
|
chunks_yielded.append(chunk)
|
|
yield chunk
|
|
|
|
mock_stream = MagicMock()
|
|
mock_stream.status_code = 200
|
|
mock_stream.headers = {} # No content-length to force streaming check
|
|
mock_stream.__aenter__ = AsyncMock(return_value=mock_stream)
|
|
mock_stream.__aexit__ = AsyncMock(return_value=None)
|
|
mock_stream.aiter_bytes = aiter_bytes
|
|
|
|
mock_client = AsyncMock()
|
|
mock_client.stream = MagicMock(return_value=mock_stream)
|
|
mock_client.__aenter__.return_value = mock_client
|
|
mock_client.__aexit__ = AsyncMock(return_value=None)
|
|
mock_client_class.return_value = mock_client
|
|
|
|
with pytest.raises(SSRFFetchError, match="too large"):
|
|
await ssrf_safe_fetch("https://example.com/api", max_size=5120)
|
|
|
|
# Verify we stopped after exceeding the limit (should be ~6 chunks for 5KB limit)
|
|
# This confirms we're enforcing during streaming, not after downloading all
|
|
assert len(chunks_yielded) <= 7, (
|
|
f"Downloaded {len(chunks_yielded)} chunks (expected <=7 for streaming enforcement)"
|
|
)
|
|
|
|
async def test_content_length_header_checked_first(self):
|
|
"""Verify Content-Length header is checked before streaming."""
|
|
with (
|
|
patch(
|
|
"fastmcp.server.auth.ssrf.resolve_hostname",
|
|
return_value=["93.184.216.34"],
|
|
),
|
|
patch("httpx.AsyncClient") as mock_client_class,
|
|
):
|
|
mock_stream = MagicMock()
|
|
mock_stream.status_code = 200
|
|
mock_stream.headers = {"content-length": "10240"} # 10KB
|
|
mock_stream.__aenter__ = AsyncMock(return_value=mock_stream)
|
|
mock_stream.__aexit__ = AsyncMock(return_value=None)
|
|
|
|
# aiter_bytes should never be called if Content-Length is checked
|
|
mock_stream.aiter_bytes = MagicMock(
|
|
side_effect=AssertionError("Should not stream")
|
|
)
|
|
|
|
mock_client = AsyncMock()
|
|
mock_client.stream = MagicMock(return_value=mock_stream)
|
|
mock_client.__aenter__.return_value = mock_client
|
|
mock_client.__aexit__ = AsyncMock(return_value=None)
|
|
mock_client_class.return_value = mock_client
|
|
|
|
with pytest.raises(SSRFFetchError, match="too large"):
|
|
await ssrf_safe_fetch("https://example.com/api", max_size=5120)
|