mirror of
https://github.com/Skyvern-AI/skyvern.git
synced 2026-07-09 16:09:13 +00:00
🔄 synced local 'tests/unit/' with remote 'tests/unit/'
This commit is contained in:
parent
4462aebd16
commit
086b93f19e
6 changed files with 504 additions and 8 deletions
|
|
@ -1,4 +1,5 @@
|
|||
import os
|
||||
import socket
|
||||
import tempfile
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
|
@ -6,7 +7,89 @@ from unittest.mock import AsyncMock, MagicMock, patch
|
|||
import aiohttp
|
||||
import pytest
|
||||
|
||||
from skyvern.forge.sdk.core.aiohttp_helper import aiohttp_request
|
||||
from skyvern.exceptions import BlockedHost, HttpException
|
||||
from skyvern.forge.sdk.core.aiohttp_helper import SSRFGuardedResolver, aiohttp_request
|
||||
from skyvern.utils.url_validators import MAX_SAFE_REDIRECTS, validate_fetch_url
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def public_dns(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
def resolves_public(host: str, port: int | None, *args: object, **kwargs: object) -> list[object]:
|
||||
return [(socket.AF_INET, socket.SOCK_STREAM, 0, "", ("93.184.216.34", port or 0))]
|
||||
|
||||
monkeypatch.setattr("skyvern.utils.url_validators.socket.getaddrinfo", resolves_public)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ssrf_guarded_resolver_blocks_dns_rebind_after_validation(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
dns_answers = iter(
|
||||
[
|
||||
[(socket.AF_INET, socket.SOCK_STREAM, 0, "", ("93.184.216.34", 0))],
|
||||
[(socket.AF_INET, socket.SOCK_STREAM, 0, "", ("169.254.169.254", 0))],
|
||||
]
|
||||
)
|
||||
|
||||
def rebinding_dns(host: str, port: int | None, *args: object, **kwargs: object) -> list[object]:
|
||||
return next(dns_answers)
|
||||
|
||||
monkeypatch.setattr("skyvern.utils.url_validators.socket.getaddrinfo", rebinding_dns)
|
||||
|
||||
assert validate_fetch_url("https://rebind.example.test/file.pdf") is not None
|
||||
with pytest.raises(BlockedHost):
|
||||
await SSRFGuardedResolver().resolve("rebind.example.test", 443)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ssrf_guarded_resolver_uses_pinned_ips_without_resolving(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
def unexpected_dns(host: str, port: int | None, *args: object, **kwargs: object) -> list[object]:
|
||||
raise AssertionError("pinned host should not be re-resolved")
|
||||
|
||||
resolver = SSRFGuardedResolver()
|
||||
resolver.pin_url_ips("https://rebind.example.test/file.pdf", ("93.184.216.34",))
|
||||
monkeypatch.setattr("skyvern.utils.url_validators.socket.getaddrinfo", unexpected_dns)
|
||||
|
||||
results = await resolver.resolve("rebind.example.test", 443)
|
||||
|
||||
assert [result["host"] for result in results] == ["93.184.216.34"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aiohttp_request_trusts_internal_proxy_without_trusting_target() -> None:
|
||||
captured_resolver: SSRFGuardedResolver | None = None
|
||||
|
||||
def capture_connector(resolver: SSRFGuardedResolver | None = None) -> object:
|
||||
nonlocal captured_resolver
|
||||
captured_resolver = resolver
|
||||
return object()
|
||||
|
||||
mock_response = AsyncMock()
|
||||
mock_response.status = 200
|
||||
mock_response.headers = {"Content-Type": "application/json"}
|
||||
mock_response.json = AsyncMock(return_value={"success": True})
|
||||
mock_response.text = AsyncMock(return_value='{"success": true}')
|
||||
mock_response.__aenter__ = AsyncMock(return_value=mock_response)
|
||||
mock_response.__aexit__ = AsyncMock(return_value=None)
|
||||
|
||||
mock_session = MagicMock()
|
||||
mock_session.__aenter__ = AsyncMock(return_value=mock_session)
|
||||
mock_session.__aexit__ = AsyncMock(return_value=None)
|
||||
mock_session.request = MagicMock(return_value=mock_response)
|
||||
|
||||
with (
|
||||
patch("skyvern.forge.sdk.core.aiohttp_helper.ssrf_guarded_tcp_connector", side_effect=capture_connector),
|
||||
patch("skyvern.forge.sdk.core.aiohttp_helper.aiohttp.ClientSession", return_value=mock_session),
|
||||
):
|
||||
await aiohttp_request(
|
||||
method="GET",
|
||||
url="https://example.com/api",
|
||||
proxy="http://127.0.0.1:8080",
|
||||
)
|
||||
|
||||
assert captured_resolver is not None
|
||||
proxy_results = await captured_resolver.resolve("127.0.0.1", 8080)
|
||||
assert proxy_results
|
||||
with pytest.raises(BlockedHost):
|
||||
await SSRFGuardedResolver().resolve("127.0.0.1", 8080)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -543,7 +626,7 @@ async def test_aiohttp_request_response_headers_dict() -> None:
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aiohttp_request_follow_redirects() -> None:
|
||||
"""Test that follow_redirects parameter is passed correctly"""
|
||||
"""Test that redirects are handled manually for SSRF validation."""
|
||||
captured_args: list[Any] = []
|
||||
captured_request_kwargs: dict[str, Any] = {}
|
||||
|
||||
|
|
@ -573,10 +656,202 @@ async def test_aiohttp_request_follow_redirects() -> None:
|
|||
follow_redirects=False,
|
||||
)
|
||||
|
||||
# Verify allow_redirects parameter was set correctly
|
||||
# aiohttp automatic redirects stay disabled so each Location target is SSRF-checked first.
|
||||
assert captured_request_kwargs["allow_redirects"] is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aiohttp_request_follows_safe_redirect() -> None:
|
||||
redirect_response = AsyncMock()
|
||||
redirect_response.status = 302
|
||||
redirect_response.headers = {"Location": "https://example.com/final"}
|
||||
redirect_response.__aenter__ = AsyncMock(return_value=redirect_response)
|
||||
redirect_response.__aexit__ = AsyncMock(return_value=None)
|
||||
|
||||
final_response = AsyncMock()
|
||||
final_response.status = 200
|
||||
final_response.headers = {"Content-Type": "application/json"}
|
||||
final_response.json = AsyncMock(return_value={"success": True})
|
||||
final_response.text = AsyncMock(return_value='{"success": true}')
|
||||
final_response.__aenter__ = AsyncMock(return_value=final_response)
|
||||
final_response.__aexit__ = AsyncMock(return_value=None)
|
||||
|
||||
responses = [redirect_response, final_response]
|
||||
requested_urls: list[str] = []
|
||||
|
||||
def capture_request(*args: Any, **kwargs: Any) -> AsyncMock:
|
||||
requested_urls.append(kwargs["url"])
|
||||
return responses.pop(0)
|
||||
|
||||
mock_session = MagicMock()
|
||||
mock_session.__aenter__ = AsyncMock(return_value=mock_session)
|
||||
mock_session.__aexit__ = AsyncMock(return_value=None)
|
||||
mock_session.request = MagicMock(side_effect=capture_request)
|
||||
|
||||
with patch("skyvern.forge.sdk.core.aiohttp_helper.aiohttp.ClientSession", return_value=mock_session):
|
||||
status, _headers, body = await aiohttp_request(
|
||||
method="GET",
|
||||
url="https://example.com/start",
|
||||
follow_redirects=True,
|
||||
)
|
||||
|
||||
assert status == 200
|
||||
assert body == {"success": True}
|
||||
assert requested_urls == ["https://example.com/start", "https://example.com/final"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aiohttp_request_strips_credentials_on_cross_origin_redirect() -> None:
|
||||
redirect_response = AsyncMock()
|
||||
redirect_response.status = 302
|
||||
redirect_response.headers = {"Location": "https://other.example.com/final"}
|
||||
redirect_response.__aenter__ = AsyncMock(return_value=redirect_response)
|
||||
redirect_response.__aexit__ = AsyncMock(return_value=None)
|
||||
|
||||
final_response = AsyncMock()
|
||||
final_response.status = 200
|
||||
final_response.headers = {"Content-Type": "application/json"}
|
||||
final_response.json = AsyncMock(return_value={"success": True})
|
||||
final_response.text = AsyncMock(return_value='{"success": true}')
|
||||
final_response.__aenter__ = AsyncMock(return_value=final_response)
|
||||
final_response.__aexit__ = AsyncMock(return_value=None)
|
||||
|
||||
responses = [redirect_response, final_response]
|
||||
requested_headers: list[dict[str, str]] = []
|
||||
requested_cookies: list[dict[str, str] | None] = []
|
||||
|
||||
def capture_request(*args: Any, **kwargs: Any) -> AsyncMock:
|
||||
requested_headers.append(kwargs["headers"])
|
||||
requested_cookies.append(kwargs["cookies"])
|
||||
return responses.pop(0)
|
||||
|
||||
mock_session = MagicMock()
|
||||
mock_session.__aenter__ = AsyncMock(return_value=mock_session)
|
||||
mock_session.__aexit__ = AsyncMock(return_value=None)
|
||||
mock_session.request = MagicMock(side_effect=capture_request)
|
||||
|
||||
with patch("skyvern.forge.sdk.core.aiohttp_helper.aiohttp.ClientSession", return_value=mock_session):
|
||||
await aiohttp_request(
|
||||
method="GET",
|
||||
url="https://example.com/start",
|
||||
headers={"Authorization": "Bearer secret", "Cookie": "sid=abc", "X-Keep": "1"},
|
||||
cookies={"session": "abc"},
|
||||
follow_redirects=True,
|
||||
)
|
||||
|
||||
assert requested_headers[0]["Authorization"] == "Bearer secret"
|
||||
assert requested_headers[0]["Cookie"] == "sid=abc"
|
||||
assert requested_cookies[0] == {"session": "abc"}
|
||||
assert requested_headers[1] == {"X-Keep": "1"}
|
||||
assert requested_cookies[1] is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aiohttp_request_rebuilds_multipart_for_307_redirect() -> None:
|
||||
with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".txt") as tmp_file:
|
||||
tmp_file.write("test file content")
|
||||
tmp_file_path = tmp_file.name
|
||||
|
||||
try:
|
||||
redirect_response = AsyncMock()
|
||||
redirect_response.status = 307
|
||||
redirect_response.headers = {"Location": "https://example.com/final"}
|
||||
redirect_response.__aenter__ = AsyncMock(return_value=redirect_response)
|
||||
redirect_response.__aexit__ = AsyncMock(return_value=None)
|
||||
|
||||
final_response = AsyncMock()
|
||||
final_response.status = 200
|
||||
final_response.headers = {"Content-Type": "application/json"}
|
||||
final_response.json = AsyncMock(return_value={"success": True})
|
||||
final_response.text = AsyncMock(return_value='{"success": true}')
|
||||
final_response.__aenter__ = AsyncMock(return_value=final_response)
|
||||
final_response.__aexit__ = AsyncMock(return_value=None)
|
||||
|
||||
responses = [redirect_response, final_response]
|
||||
multipart_bodies: list[aiohttp.FormData] = []
|
||||
|
||||
def capture_request(*args: Any, **kwargs: Any) -> AsyncMock:
|
||||
multipart_bodies.append(kwargs["data"])
|
||||
return responses.pop(0)
|
||||
|
||||
mock_session = MagicMock()
|
||||
mock_session.__aenter__ = AsyncMock(return_value=mock_session)
|
||||
mock_session.__aexit__ = AsyncMock(return_value=None)
|
||||
mock_session.request = MagicMock(side_effect=capture_request)
|
||||
|
||||
with patch("skyvern.forge.sdk.core.aiohttp_helper.aiohttp.ClientSession", return_value=mock_session):
|
||||
status, _headers, body = await aiohttp_request(
|
||||
method="POST",
|
||||
url="https://example.com/start",
|
||||
files={"file": tmp_file_path},
|
||||
follow_redirects=True,
|
||||
)
|
||||
|
||||
assert status == 200
|
||||
assert body == {"success": True}
|
||||
assert len(multipart_bodies) == 2
|
||||
assert multipart_bodies[0] is not multipart_bodies[1]
|
||||
finally:
|
||||
if os.path.exists(tmp_file_path):
|
||||
os.unlink(tmp_file_path)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aiohttp_request_rejects_unsafe_redirect_target() -> None:
|
||||
redirect_response = AsyncMock()
|
||||
redirect_response.status = 302
|
||||
redirect_response.headers = {"Location": "http://169.254.169.254/latest/meta-data"}
|
||||
redirect_response.__aenter__ = AsyncMock(return_value=redirect_response)
|
||||
redirect_response.__aexit__ = AsyncMock(return_value=None)
|
||||
|
||||
requested_urls: list[str] = []
|
||||
|
||||
def capture_request(*args: Any, **kwargs: Any) -> AsyncMock:
|
||||
requested_urls.append(kwargs["url"])
|
||||
return redirect_response
|
||||
|
||||
mock_session = MagicMock()
|
||||
mock_session.__aenter__ = AsyncMock(return_value=mock_session)
|
||||
mock_session.__aexit__ = AsyncMock(return_value=None)
|
||||
mock_session.request = MagicMock(side_effect=capture_request)
|
||||
|
||||
with patch("skyvern.forge.sdk.core.aiohttp_helper.aiohttp.ClientSession", return_value=mock_session):
|
||||
with pytest.raises(BlockedHost):
|
||||
await aiohttp_request(
|
||||
method="GET",
|
||||
url="https://example.com/start",
|
||||
follow_redirects=True,
|
||||
)
|
||||
|
||||
assert requested_urls == ["https://example.com/start"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aiohttp_request_redirect_limit_raises_http_exception() -> None:
|
||||
redirect_response = AsyncMock()
|
||||
redirect_response.status = 302
|
||||
redirect_response.headers = {"Location": "https://example.com/next"}
|
||||
redirect_response.__aenter__ = AsyncMock(return_value=redirect_response)
|
||||
redirect_response.__aexit__ = AsyncMock(return_value=None)
|
||||
|
||||
mock_session = MagicMock()
|
||||
mock_session.__aenter__ = AsyncMock(return_value=mock_session)
|
||||
mock_session.__aexit__ = AsyncMock(return_value=None)
|
||||
mock_session.request = MagicMock(return_value=redirect_response)
|
||||
|
||||
with patch("skyvern.forge.sdk.core.aiohttp_helper.aiohttp.ClientSession", return_value=mock_session):
|
||||
with pytest.raises(HttpException) as exc_info:
|
||||
await aiohttp_request(
|
||||
method="GET",
|
||||
url="https://example.com/start",
|
||||
follow_redirects=True,
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
assert "Too many redirects" in (exc_info.value.error_message or "")
|
||||
assert mock_session.request.call_count == MAX_SAFE_REDIRECTS + 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aiohttp_request_proxy_passed_correctly() -> None:
|
||||
"""Test that proxy parameter is passed correctly"""
|
||||
|
|
@ -646,10 +921,12 @@ async def test_aiohttp_request_with_files_uses_multipart() -> None:
|
|||
mock_session.__aexit__ = AsyncMock(return_value=None)
|
||||
mock_session.request = MagicMock(side_effect=capture_request)
|
||||
|
||||
headers = {"Content-Type": "multipart/form-data", "X-Test": "value"}
|
||||
with patch("skyvern.forge.sdk.core.aiohttp_helper.aiohttp.ClientSession", return_value=mock_session):
|
||||
await aiohttp_request(
|
||||
method="POST",
|
||||
url="https://example.com/api/upload",
|
||||
headers=headers,
|
||||
files={"file": tmp_file_path},
|
||||
)
|
||||
|
||||
|
|
@ -660,6 +937,7 @@ async def test_aiohttp_request_with_files_uses_multipart() -> None:
|
|||
# Verify Content-Type header was removed (aiohttp will set it for multipart)
|
||||
assert "Content-Type" not in captured_request_kwargs["headers"]
|
||||
assert "content-type" not in captured_request_kwargs["headers"]
|
||||
assert headers == {"Content-Type": "multipart/form-data", "X-Test": "value"}
|
||||
finally:
|
||||
# Clean up temporary file
|
||||
if os.path.exists(tmp_file_path):
|
||||
|
|
|
|||
|
|
@ -1,6 +1,9 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import socket
|
||||
from collections.abc import AsyncIterator
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
|
|
@ -8,7 +11,9 @@ import pytest
|
|||
|
||||
from skyvern.config import settings
|
||||
from skyvern.constants import DOWNLOAD_FILE_PREFIX
|
||||
from skyvern.exceptions import BlockedHost, SkyvernHTTPException
|
||||
from skyvern.forge.sdk.api import files
|
||||
from skyvern.utils.url_validators import MAX_SAFE_REDIRECTS
|
||||
|
||||
ATTACKER_ORG_ID = "o_attacker"
|
||||
VICTIM_ORG_ID = "o_victim"
|
||||
|
|
@ -242,3 +247,122 @@ async def test_download_file_routes_gcs_uri_to_managed_storage(storage: SimpleNa
|
|||
assert f.read() == b"tenant-secret-bytes"
|
||||
finally:
|
||||
os.unlink(path)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_file_blocks_hostname_resolving_to_private_ip(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
def resolves_private(host: str, port: int | None, *args: object, **kwargs: object) -> list[object]:
|
||||
return [(socket.AF_INET, socket.SOCK_STREAM, 0, "", ("10.0.0.42", port or 0))]
|
||||
|
||||
client_session = MagicMock()
|
||||
monkeypatch.setattr("skyvern.utils.url_validators.socket.getaddrinfo", resolves_private)
|
||||
monkeypatch.setattr(files.aiohttp, "ClientSession", client_session)
|
||||
|
||||
with pytest.raises(BlockedHost):
|
||||
await files.download_file("https://evil.example.test/secret.pdf")
|
||||
|
||||
client_session.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_file_rejects_unsafe_redirect_target(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
def resolves_public(host: str, port: int | None, *args: object, **kwargs: object) -> list[object]:
|
||||
return [(socket.AF_INET, socket.SOCK_STREAM, 0, "", ("93.184.216.34", port or 0))]
|
||||
|
||||
redirect_response = AsyncMock()
|
||||
redirect_response.status = 302
|
||||
redirect_response.headers = {"Location": "http://169.254.169.254/latest/meta-data"}
|
||||
redirect_response.__aenter__ = AsyncMock(return_value=redirect_response)
|
||||
redirect_response.__aexit__ = AsyncMock(return_value=None)
|
||||
|
||||
mock_session = MagicMock()
|
||||
mock_session.__aenter__ = AsyncMock(return_value=mock_session)
|
||||
mock_session.__aexit__ = AsyncMock(return_value=None)
|
||||
mock_session.get = MagicMock(return_value=redirect_response)
|
||||
|
||||
monkeypatch.setattr("skyvern.utils.url_validators.socket.getaddrinfo", resolves_public)
|
||||
monkeypatch.setattr(files.aiohttp, "ClientSession", MagicMock(return_value=mock_session))
|
||||
|
||||
with pytest.raises(BlockedHost):
|
||||
await files.download_file("https://example.com/start.pdf")
|
||||
|
||||
mock_session.get.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_file_strips_credentials_on_cross_origin_redirect(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
def resolves_public(host: str, port: int | None, *args: object, **kwargs: object) -> list[object]:
|
||||
return [(socket.AF_INET, socket.SOCK_STREAM, 0, "", ("93.184.216.34", port or 0))]
|
||||
|
||||
class Content:
|
||||
async def iter_chunked(self, chunk_size: int) -> AsyncIterator[bytes]:
|
||||
yield b"file-bytes"
|
||||
|
||||
redirect_response = AsyncMock()
|
||||
redirect_response.status = 302
|
||||
redirect_response.headers = {"Location": "https://other.example.com/final.pdf"}
|
||||
redirect_response.__aenter__ = AsyncMock(return_value=redirect_response)
|
||||
redirect_response.__aexit__ = AsyncMock(return_value=None)
|
||||
|
||||
final_response = AsyncMock()
|
||||
final_response.status = 200
|
||||
final_response.headers = {}
|
||||
final_response.content_length = len(b"file-bytes")
|
||||
final_response.content = Content()
|
||||
final_response.__aenter__ = AsyncMock(return_value=final_response)
|
||||
final_response.__aexit__ = AsyncMock(return_value=None)
|
||||
|
||||
responses = [redirect_response, final_response]
|
||||
requested_headers: list[dict[str, str]] = []
|
||||
|
||||
def capture_get(*args: object, **kwargs: object) -> AsyncMock:
|
||||
headers = kwargs["headers"]
|
||||
assert isinstance(headers, dict)
|
||||
requested_headers.append(headers)
|
||||
return responses.pop(0)
|
||||
|
||||
mock_session = MagicMock()
|
||||
mock_session.__aenter__ = AsyncMock(return_value=mock_session)
|
||||
mock_session.__aexit__ = AsyncMock(return_value=None)
|
||||
mock_session.get = MagicMock(side_effect=capture_get)
|
||||
|
||||
monkeypatch.setattr("skyvern.utils.url_validators.socket.getaddrinfo", resolves_public)
|
||||
monkeypatch.setattr(files.aiohttp, "ClientSession", MagicMock(return_value=mock_session))
|
||||
|
||||
result = await files.download_file(
|
||||
"https://example.com/start.pdf",
|
||||
headers={"Authorization": "Bearer secret", "Cookie": "sid=abc", "X-Keep": "1"},
|
||||
output_dir=str(tmp_path),
|
||||
)
|
||||
|
||||
assert Path(result).read_bytes() == b"file-bytes"
|
||||
assert requested_headers[0]["Authorization"] == "Bearer secret"
|
||||
assert requested_headers[0]["Cookie"] == "sid=abc"
|
||||
assert requested_headers[1] == {"X-Keep": "1"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_file_redirect_limit_raises_http_exception(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
def resolves_public(host: str, port: int | None, *args: object, **kwargs: object) -> list[object]:
|
||||
return [(socket.AF_INET, socket.SOCK_STREAM, 0, "", ("93.184.216.34", port or 0))]
|
||||
|
||||
redirect_response = AsyncMock()
|
||||
redirect_response.status = 302
|
||||
redirect_response.headers = {"Location": "https://example.com/next.pdf"}
|
||||
redirect_response.__aenter__ = AsyncMock(return_value=redirect_response)
|
||||
redirect_response.__aexit__ = AsyncMock(return_value=None)
|
||||
|
||||
mock_session = MagicMock()
|
||||
mock_session.__aenter__ = AsyncMock(return_value=mock_session)
|
||||
mock_session.__aexit__ = AsyncMock(return_value=None)
|
||||
mock_session.get = MagicMock(return_value=redirect_response)
|
||||
|
||||
monkeypatch.setattr("skyvern.utils.url_validators.socket.getaddrinfo", resolves_public)
|
||||
monkeypatch.setattr(files.aiohttp, "ClientSession", MagicMock(return_value=mock_session))
|
||||
|
||||
with pytest.raises(SkyvernHTTPException, match="Too many redirects"):
|
||||
await files.download_file("https://example.com/start.pdf")
|
||||
|
||||
assert mock_session.get.call_count == MAX_SAFE_REDIRECTS + 1
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ class _FakeDownloadResponse:
|
|||
self._data = data
|
||||
# aiohttp exposes headers as a case-insensitive CIMultiDictProxy; mirror that.
|
||||
self.headers = CIMultiDictProxy(CIMultiDict(headers or {}))
|
||||
self.status = 200
|
||||
self.content_length = len(data) if advertise_length else None
|
||||
self.content = self
|
||||
self.body_read = False
|
||||
|
|
@ -36,7 +37,9 @@ class _FakeDownloadSession:
|
|||
def __init__(self, response: _FakeDownloadResponse) -> None:
|
||||
self._response = response
|
||||
|
||||
def get(self, url: object, headers: dict[str, str] | None = None) -> _FakeDownloadResponse:
|
||||
def get(
|
||||
self, url: object, headers: dict[str, str] | None = None, allow_redirects: bool = True
|
||||
) -> _FakeDownloadResponse:
|
||||
return self._response
|
||||
|
||||
async def __aenter__(self) -> _FakeDownloadSession:
|
||||
|
|
|
|||
|
|
@ -145,10 +145,10 @@ def test_debug_fields_handles_none_inputs() -> None:
|
|||
async def test_resolve_org_from_api_key_logs_decode_error_reason(monkeypatch) -> None:
|
||||
logged: dict[str, object] = {}
|
||||
|
||||
def fake_error(_message: str, **kwargs: object) -> None:
|
||||
def fake_warning(_message: str, **kwargs: object) -> None:
|
||||
logged.update(kwargs)
|
||||
|
||||
monkeypatch.setattr(org_auth_service.LOG, "error", fake_error)
|
||||
monkeypatch.setattr(org_auth_service.LOG, "warning", fake_warning)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await org_auth_service.resolve_org_from_api_key("definitely-not-a-jwt", SimpleNamespace(), ())
|
||||
|
|
|
|||
|
|
@ -1,8 +1,16 @@
|
|||
import socket
|
||||
|
||||
import pytest
|
||||
|
||||
from skyvern.config import settings
|
||||
from skyvern.exceptions import BlockedHost
|
||||
from skyvern.utils.url_validators import encode_url, is_blocked_host, validate_url
|
||||
from skyvern.utils.url_validators import (
|
||||
encode_url,
|
||||
is_blocked_host,
|
||||
validate_fetch_url,
|
||||
validate_redirect_url,
|
||||
validate_url,
|
||||
)
|
||||
|
||||
|
||||
def test_encode_url_basic():
|
||||
|
|
@ -121,3 +129,60 @@ def test_is_blocked_host_allowed_hosts_normalize_brackets_and_mapped(
|
|||
@pytest.mark.parametrize("host", ["LOCALHOST", "LocalHost", "localhost"])
|
||||
def test_is_blocked_host_blocked_hosts_case_insensitive(host: str) -> None:
|
||||
assert is_blocked_host(host) is True
|
||||
|
||||
|
||||
def test_validate_fetch_url_blocks_hostname_resolving_private_ip(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
def resolves_to_private(host: str, port: int | None, *args: object, **kwargs: object) -> list[object]:
|
||||
return [(socket.AF_INET, socket.SOCK_STREAM, 0, "", ("10.0.0.42", port or 0))]
|
||||
|
||||
monkeypatch.setattr("skyvern.utils.url_validators.socket.getaddrinfo", resolves_to_private)
|
||||
|
||||
with pytest.raises(BlockedHost):
|
||||
validate_fetch_url("https://evil.example.test/file.pdf")
|
||||
|
||||
|
||||
def test_validate_fetch_url_fails_closed_on_dns_error(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
def fails_dns(host: str, port: int | None, *args: object, **kwargs: object) -> list[object]:
|
||||
raise OSError("dns unavailable")
|
||||
|
||||
monkeypatch.setattr("skyvern.utils.url_validators.socket.getaddrinfo", fails_dns)
|
||||
|
||||
with pytest.raises(BlockedHost):
|
||||
validate_fetch_url("https://unresolvable.example.test/file.pdf")
|
||||
|
||||
|
||||
def test_validate_url_does_not_resolve_dns(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
def unexpected_dns(host: str, port: int | None, *args: object, **kwargs: object) -> list[object]:
|
||||
raise AssertionError("general URL validation should not resolve DNS")
|
||||
|
||||
monkeypatch.setattr("skyvern.utils.url_validators.socket.getaddrinfo", unexpected_dns)
|
||||
|
||||
assert validate_url("https://webhook.example.com/receive") is not None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"url",
|
||||
[
|
||||
"https://metadata.google.internal/computeMetadata/v1/",
|
||||
"https://kubernetes.default.svc/api",
|
||||
"https://my-service.namespace.svc.cluster.local/api",
|
||||
"https://internal.local/api",
|
||||
],
|
||||
)
|
||||
def test_validate_url_blocks_internal_hostnames(url: str, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
def unexpected_dns(host: str, port: int | None, *args: object, **kwargs: object) -> list[object]:
|
||||
raise AssertionError("internal hostname should be blocked before DNS")
|
||||
|
||||
monkeypatch.setattr("skyvern.utils.url_validators.socket.getaddrinfo", unexpected_dns)
|
||||
|
||||
with pytest.raises(BlockedHost):
|
||||
validate_url(url)
|
||||
|
||||
|
||||
def test_is_blocked_host_allows_public_svc_subdomain() -> None:
|
||||
assert is_blocked_host("api.svc.example.com") is False
|
||||
|
||||
|
||||
def test_validate_redirect_url_rejects_private_redirect_target() -> None:
|
||||
with pytest.raises(BlockedHost):
|
||||
validate_redirect_url("https://example.com/file.pdf", "http://169.254.169.254/latest/meta-data")
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
import json
|
||||
import socket
|
||||
from datetime import datetime, timezone
|
||||
from unittest.mock import AsyncMock
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -193,6 +194,31 @@ class TestSecretResponsePaths:
|
|||
|
||||
|
||||
class TestHttpRequestBlockSecretResponsePaths:
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_rejects_hostname_resolving_to_private_ip_before_request(
|
||||
self, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
context = _make_context()
|
||||
block = _http_block(url="https://evil.example.test/api")
|
||||
db_mock = AsyncMock()
|
||||
client_session = MagicMock()
|
||||
|
||||
def resolves_private(host: str, port: int | None, *args: object, **kwargs: object) -> list[object]:
|
||||
return [(socket.AF_INET, socket.SOCK_STREAM, 0, "", ("10.0.0.42", port or 0))]
|
||||
|
||||
monkeypatch.setattr(HttpRequestBlock, "get_workflow_run_context", lambda _self, _workflow_run_id: context)
|
||||
monkeypatch.setattr("skyvern.utils.url_validators.socket.getaddrinfo", resolves_private)
|
||||
monkeypatch.setattr("skyvern.forge.sdk.core.aiohttp_helper.aiohttp.ClientSession", client_session)
|
||||
monkeypatch.setattr(block_module.app, "DATABASE", db_mock)
|
||||
|
||||
result = await block.execute(workflow_run_id="wr-1", workflow_run_block_id="wrb-1")
|
||||
|
||||
assert result.success is False
|
||||
assert result.status == BlockStatus.failed
|
||||
assert result.failure_reason is not None
|
||||
assert "blocked" in result.failure_reason
|
||||
client_session.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_records_placeholder_and_masks_duplicate_echo(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
context = _make_context()
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue