wehub-resource-sync
132 lines
5.6 KiB
Python
132 lines
5.6 KiB
Python
import http.server
|
|
import socket
|
|
import threading
|
|
from collections.abc import Iterator
|
|
from unittest import mock
|
|
|
|
import pytest
|
|
import requests
|
|
from urllib3.util.retry import Retry
|
|
|
|
from mlflow.webhooks.delivery import _create_webhook_session
|
|
from mlflow.webhooks.ssrf import SSRFProtectedHTTPAdapter, SSRFProtectionError
|
|
|
|
|
|
@pytest.fixture
|
|
def loopback_server() -> Iterator[int]:
|
|
"""A real HTTP server on loopback, standing in for an internal/metadata endpoint."""
|
|
|
|
class Handler(http.server.BaseHTTPRequestHandler):
|
|
def do_GET(self):
|
|
self.send_response(200)
|
|
self.end_headers()
|
|
self.wfile.write(b"SECRET")
|
|
|
|
def log_message(self, *args):
|
|
pass
|
|
|
|
server = http.server.HTTPServer(("127.0.0.1", 0), Handler)
|
|
port = server.server_address[1]
|
|
thread = threading.Thread(target=server.serve_forever, name="ssrf-test-loopback", daemon=True)
|
|
thread.start()
|
|
try:
|
|
yield port
|
|
finally:
|
|
server.shutdown()
|
|
|
|
|
|
def _rebind_getaddrinfo(rebind_host: str, target_host: str, target_port: int):
|
|
"""getaddrinfo that resolves rebind_host to (target_host, target_port)."""
|
|
orig = socket.getaddrinfo
|
|
|
|
def resolver(host, port, *args, **kwargs):
|
|
if host == rebind_host:
|
|
return orig(target_host, target_port, *args, **kwargs)
|
|
return orig(host, port, *args, **kwargs)
|
|
|
|
return resolver
|
|
|
|
|
|
def test_webhook_session_uses_ssrf_protected_adapter():
|
|
session = _create_webhook_session()
|
|
for scheme in ("http://", "https://"):
|
|
adapter = session.get_adapter(f"{scheme}example.com")
|
|
assert type(adapter).__name__ == "SSRFProtectedHTTPAdapter"
|
|
|
|
|
|
def test_webhook_session_ignores_env_proxy(loopback_server: int, monkeypatch: pytest.MonkeyPatch):
|
|
# With trust_env left on, an HTTP_PROXY env var would route the request
|
|
# through the proxy, making the validated peer the proxy (not the webhook
|
|
# destination) and bypassing the SSRF check entirely. The session disables
|
|
# env proxy handling, so the loopback "proxy" must never be reached.
|
|
session = _create_webhook_session()
|
|
assert session.trust_env is False
|
|
monkeypatch.setenv("HTTP_PROXY", f"http://127.0.0.1:{loopback_server}")
|
|
resolver = _rebind_getaddrinfo("dest.example.com", "127.0.0.1", loopback_server)
|
|
with mock.patch.object(socket, "getaddrinfo", side_effect=resolver):
|
|
# The real destination still resolves to loopback here, so the peer
|
|
# check fires on the destination — proving the proxy was bypassed and
|
|
# the destination IP is what gets validated.
|
|
with pytest.raises(SSRFProtectionError, match="not a public IP"):
|
|
session.post("http://dest.example.com/hook", timeout=10)
|
|
|
|
|
|
def test_dns_rebinding_to_loopback_is_blocked(loopback_server: int):
|
|
session = _create_webhook_session()
|
|
resolver = _rebind_getaddrinfo("evil.example.com", "127.0.0.1", loopback_server)
|
|
with mock.patch.object(socket, "getaddrinfo", side_effect=resolver):
|
|
with pytest.raises(SSRFProtectionError, match="not a public IP"):
|
|
session.get(f"http://evil.example.com:{loopback_server}/latest/meta-data/", timeout=10)
|
|
|
|
|
|
def test_private_ip_connection_is_blocked_by_default(loopback_server: int):
|
|
session = _create_webhook_session()
|
|
with pytest.raises(SSRFProtectionError, match="not a public IP"):
|
|
session.get(f"http://127.0.0.1:{loopback_server}/", timeout=10)
|
|
|
|
|
|
def test_private_ip_allowed_when_env_var_set(loopback_server: int, monkeypatch: pytest.MonkeyPatch):
|
|
monkeypatch.setenv("MLFLOW_WEBHOOK_ALLOW_PRIVATE_IPS", "true")
|
|
session = _create_webhook_session()
|
|
response = session.get(f"http://127.0.0.1:{loopback_server}/", timeout=10)
|
|
assert response.status_code == 200
|
|
|
|
|
|
def test_ssrf_error_is_not_retried(loopback_server: int):
|
|
# SSRFProtectionError is not a urllib3 exception, so it propagates on the
|
|
# first connection attempt rather than being retried. Use POST to exercise
|
|
# the production retry config (allowed_methods=["POST"]).
|
|
session = _create_webhook_session()
|
|
with mock.patch("mlflow.webhooks.ssrf.ipaddress.ip_address") as mock_ip:
|
|
mock_ip.return_value.is_global = False
|
|
with pytest.raises(SSRFProtectionError, match="not a public IP"):
|
|
session.post(f"http://127.0.0.1:{loopback_server}/", timeout=10)
|
|
# One check per connection attempt; no retry loop.
|
|
mock_ip.assert_called_once()
|
|
|
|
|
|
def test_explicit_proxy_to_private_ip_is_blocked(loopback_server: int):
|
|
# An explicitly configured proxy routes through proxy_manager_for, not the
|
|
# direct pool manager. The adapter's proxy manager must also use protected
|
|
# pools so the proxy endpoint's peer IP is validated.
|
|
session = requests.Session()
|
|
adapter = SSRFProtectedHTTPAdapter(max_retries=Retry(total=0))
|
|
session.mount("http://", adapter)
|
|
session.mount("https://", adapter)
|
|
session.proxies = {"http": f"http://127.0.0.1:{loopback_server}"}
|
|
|
|
assert type(adapter.proxy_manager_for(f"http://127.0.0.1:{loopback_server}")).__name__ == (
|
|
"_SSRFProtectedProxyManager"
|
|
)
|
|
with pytest.raises(SSRFProtectionError, match="not a public IP"):
|
|
session.post("http://example.com/hook", timeout=10)
|
|
|
|
|
|
def test_getpeername_oserror_fails_closed(loopback_server: int):
|
|
# If getpeername() raises OSError (e.g. ENOTCONN), the check must fail closed
|
|
# with SSRFProtectionError rather than leak an uncaught OSError.
|
|
session = _create_webhook_session()
|
|
with mock.patch.object(socket.socket, "getpeername", side_effect=OSError("ENOTCONN")):
|
|
with pytest.raises(SSRFProtectionError, match="peer address"):
|
|
session.post(f"http://127.0.0.1:{loopback_server}/", timeout=10)
|