diff --git a/doc/code/executor/5_workflow.ipynb b/doc/code/executor/5_workflow.ipynb index 10e49f914e..bbe94a9cc7 100644 --- a/doc/code/executor/5_workflow.ipynb +++ b/doc/code/executor/5_workflow.ipynb @@ -802,7 +802,7 @@ "source": [ "import pathlib\n", "\n", - "from pyrit.common.path import CONVERTER_SEED_PROMPT_PATH\n", + "from pyrit.common.path import CONVERTER_SEED_PROMPT_PATH, DB_DATA_PATH\n", "from pyrit.converter import PDFConverter\n", "from pyrit.executor.core import StrategyConverterConfig\n", "from pyrit.executor.workflow import XPIATestWorkflow\n", @@ -865,7 +865,12 @@ " injection_items=injection_items, # Inject hidden text\n", ")\n", "\n", - "upload_target = HTTPXAPITarget(http_url=f\"http://localhost:8000/upload/\", method=\"POST\", timeout=180)\n", + "upload_target = HTTPXAPITarget(\n", + " http_url=\"http://localhost:8000/upload/\",\n", + " method=\"POST\",\n", + " allowed_upload_directory=DB_DATA_PATH,\n", + " timeout=180,\n", + ")\n", "\n", "http_api_processing_target = HTTPXAPITarget(\n", " http_url=f\"http://localhost:8000/search_candidates/\", method=\"POST\", timeout=180\n", diff --git a/doc/code/executor/5_workflow.py b/doc/code/executor/5_workflow.py index 1780145d9a..a222d40446 100644 --- a/doc/code/executor/5_workflow.py +++ b/doc/code/executor/5_workflow.py @@ -191,7 +191,7 @@ async def processing_callback() -> str: # %% import pathlib -from pyrit.common.path import CONVERTER_SEED_PROMPT_PATH +from pyrit.common.path import CONVERTER_SEED_PROMPT_PATH, DB_DATA_PATH from pyrit.converter import PDFConverter from pyrit.executor.core import StrategyConverterConfig from pyrit.executor.workflow import XPIATestWorkflow @@ -254,7 +254,12 @@ async def processing_callback() -> str: injection_items=injection_items, # Inject hidden text ) -upload_target = HTTPXAPITarget(http_url=f"http://localhost:8000/upload/", method="POST", timeout=180) +upload_target = HTTPXAPITarget( + http_url="http://localhost:8000/upload/", + method="POST", + allowed_upload_directory=DB_DATA_PATH, + timeout=180, +) http_api_processing_target = HTTPXAPITarget( http_url=f"http://localhost:8000/search_candidates/", method="POST", timeout=180 diff --git a/pyrit/prompt_target/http_target/http_target.py b/pyrit/prompt_target/http_target/http_target.py index 98304f1ffa..bc0fbc57d2 100644 --- a/pyrit/prompt_target/http_target/http_target.py +++ b/pyrit/prompt_target/http_target/http_target.py @@ -7,6 +7,7 @@ import re from collections.abc import Callable from typing import Any +from urllib.parse import urlsplit import httpx @@ -42,6 +43,7 @@ def __init__( max_requests_per_minute: int | None = None, client: httpx.AsyncClient | None = None, model_name: str = "", + follow_redirects: bool = True, custom_configuration: TargetConfiguration | None = None, **httpx_client_kwargs: Any, ) -> None: @@ -58,6 +60,8 @@ def __init__( max_requests_per_minute (int, Optional): Maximum number of requests per minute. client (httpx.AsyncClient, Optional): Pre-configured httpx client. model_name (str): The model name. Defaults to empty string. + follow_redirects (bool): Whether to follow HTTP redirects. Defaults to True for backward compatibility; + set to False when redirects are unnecessary or the destination must remain fixed. custom_configuration (TargetConfiguration, Optional): Override the default configuration for this target instance. Defaults to None. **httpx_client_kwargs: Additional keyword arguments for httpx.AsyncClient. @@ -72,6 +76,7 @@ def __init__( # Parse the URL early to use as endpoint identifier # This will fail early if the http_request is malformed _, _, endpoint, _, _ = self.parse_raw_http_request(http_request) + self._destination_origin = self._get_destination_origin(endpoint) super().__init__( max_requests_per_minute=max_requests_per_minute, @@ -82,6 +87,7 @@ def __init__( self.http_request = http_request self.callback_function = callback_function self.prompt_regex_string = prompt_regex_string + self.follow_redirects = follow_redirects self.httpx_client_kwargs = httpx_client_kwargs or {} if client and httpx_client_kwargs: @@ -99,6 +105,7 @@ def _build_identifier(self) -> ComponentIdentifier: "use_tls": self.use_tls, "prompt_regex_string": self.prompt_regex_string, "callback_function": getattr(self.callback_function, "__name__", None), + "follow_redirects": self.follow_redirects, }, ) @@ -110,6 +117,7 @@ def with_client( prompt_regex_string: str = "{PROMPT}", callback_function: Callable[..., Any] | None = None, max_requests_per_minute: int | None = None, + follow_redirects: bool = True, ) -> "HTTPTarget": """ Alternative constructor that accepts a pre-configured httpx client. @@ -120,6 +128,8 @@ def with_client( prompt_regex_string: the placeholder for the prompt callback_function: function to parse HTTP response max_requests_per_minute: Optional rate limiting + follow_redirects: Whether to follow HTTP redirects. Defaults to True for backward compatibility; set to + False when redirects are unnecessary or the destination must remain fixed. Returns: HTTPTarget: an instance of HTTPTarget @@ -130,6 +140,7 @@ def with_client( callback_function=callback_function, max_requests_per_minute=max_requests_per_minute, client=client, + follow_redirects=follow_redirects, ) def _inject_prompt_into_request(self, request: MessagePiece) -> str: @@ -142,8 +153,12 @@ def _inject_prompt_into_request(self, request: MessagePiece) -> str: Returns: str: the http request with the prompt added in + + Raises: + ValueError: If a multiline prompt would be substituted into the request line or headers. """ re_pattern = re.compile(self.prompt_regex_string) + self._validate_prompt_for_template_context(prompt=request.converted_value, pattern=re_pattern) if re.search(self.prompt_regex_string, self.http_request): http_request_w_prompt = re_pattern.sub(lambda m: request.converted_value, self.http_request) else: @@ -169,6 +184,7 @@ async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Me http_request_w_prompt = self._inject_prompt_into_request(request) header_dict, http_body, url, http_method, http_version = self.parse_raw_http_request(http_request_w_prompt) + self._validate_destination(url) if "Content-Length" in header_dict: header_dict["Content-Length"] = str(len(http_body)) @@ -191,7 +207,7 @@ async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Me url=url, headers=header_dict, data=http_body, - follow_redirects=True, + follow_redirects=self.follow_redirects, ) else: response = await client.request( @@ -199,7 +215,7 @@ async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Me url=url, headers=header_dict, content=http_body, - follow_redirects=True, + follow_redirects=self.follow_redirects, ) response_content = response.content @@ -296,3 +312,27 @@ def _infer_full_url_from_host( host = headers_dict["host"] return f"{http_protocol}{host}{path}" + + def _validate_destination(self, url: str) -> None: + destination_origin = self._get_destination_origin(url) + if destination_origin != self._destination_origin: + raise ValueError("Prompt substitution cannot change the configured HTTP destination.") + + def _validate_prompt_for_template_context(self, *, prompt: str, pattern: re.Pattern[str]) -> None: + if "\r" not in prompt and "\n" not in prompt: + return + + separator = re.search(r"\r?\n\r?\n", self.http_request) + header_end = separator.start() if separator else len(self.http_request) + + if any(match.start() < header_end for match in pattern.finditer(self.http_request)): + raise ValueError("Prompts substituted into the HTTP request line or headers cannot contain CR or LF.") + + @staticmethod + def _get_destination_origin(url: str) -> tuple[str, str | None, int | None]: + parsed_url = urlsplit(url) + try: + port = parsed_url.port + except ValueError as exc: + raise ValueError(f"Invalid port in HTTP destination: {url}") from exc + return parsed_url.scheme.lower(), parsed_url.hostname, port diff --git a/pyrit/prompt_target/http_target/httpx_api_target.py b/pyrit/prompt_target/http_target/httpx_api_target.py index b64d25b9ad..465ae0fa1f 100644 --- a/pyrit/prompt_target/http_target/httpx_api_target.py +++ b/pyrit/prompt_target/http_target/httpx_api_target.py @@ -1,6 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +import asyncio import logging import mimetypes from collections.abc import Callable @@ -10,6 +11,7 @@ import aiofiles import httpx +from pyrit.common.deprecation import print_deprecation_message from pyrit.models import ( Message, MessagePiece, @@ -33,6 +35,7 @@ class HTTPXAPITarget(HTTPTarget): it's a local file path generated by a Converter (like PDFConverter). """ + _PATH_TYPES: frozenset[str] = frozenset({"image_path", "audio_path", "video_path", "binary_path"}) _DEFAULT_CONFIGURATION: TargetConfiguration = TargetConfiguration( capabilities=TargetCapabilities( supports_multi_turn=True, @@ -54,11 +57,13 @@ def __init__( http_url: str, method: Literal["GET", "POST", "PUT", "DELETE", "PATCH", "HEAD", "OPTIONS"] = "POST", file_path: str | None = None, + allowed_upload_directory: str | Path | None = None, json_data: dict[str, Any] | None = None, form_data: dict[str, Any] | None = None, params: dict[str, Any] | None = None, headers: dict[str, str] | None = None, http2: bool | None = None, + follow_redirects: bool = True, callback_function: Callable[..., Any] | None = None, max_requests_per_minute: int | None = None, custom_configuration: TargetConfiguration | None = None, @@ -72,11 +77,15 @@ def __init__( method (str): The HTTP method to use (GET, POST, PUT, DELETE, PATCH, HEAD, OPTIONS). Defaults to "POST". file_path (str, Optional): Path to a file to upload. If not provided, we attempt to pull it from the prompt's `converted_value`. + allowed_upload_directory (str | Path | None): Directory containing files this target may upload. + Required when uploading files. Defaults to None. json_data (dict, Optional): JSON data to send in the request body (for POST/PUT/PATCH). form_data (dict, Optional): Form data to send in the request body (for POST/PUT/PATCH). params (dict, Optional): Query parameters to include in the request URL (for GET/HEAD). headers (dict, Optional): Headers to include in the request. http2 (bool, Optional): Whether to use HTTP/2. If None, defaults to False. + follow_redirects (bool): Whether to follow HTTP redirects. Defaults to True for backward compatibility; + set to False when redirects are unnecessary or the destination must remain fixed. callback_function (Callable, Optional): Function to parse the HTTP response. max_requests_per_minute (int, Optional): Maximum number of requests per minute. custom_configuration (TargetConfiguration, Optional): Override the default configuration for this target @@ -87,11 +96,13 @@ def __init__( Raises: ValueError: If the HTTP method is invalid. ValueError: If file uploads are attempted with an HTTP method that does not support them. + ValueError: If the allowed upload directory does not exist or is not a directory. """ super().__init__( http_request="", prompt_regex_string="", use_tls=True, + follow_redirects=follow_redirects, callback_function=callback_function, max_requests_per_minute=max_requests_per_minute, custom_configuration=custom_configuration, @@ -101,6 +112,7 @@ def __init__( self.http_url = http_url self.method = method self.file_path = file_path + self.allowed_upload_directory = self._resolve_allowed_upload_directory(allowed_upload_directory) self.json_data = json_data self.form_data = form_data self.params = params @@ -111,10 +123,6 @@ def __init__( if self.method not in {"GET", "POST", "PUT", "DELETE", "PATCH", "HEAD", "OPTIONS"}: raise ValueError(f"Invalid HTTP method: {self.method}") - # Validate file uploads (only `POST` and `PUT` allow file uploads) - if self.file_path and self.method not in {"POST", "PUT"}: - raise ValueError(f"File uploads are not allowed with HTTP method: {self.method}") - @limit_requests_per_minute async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Message]) -> list[Message]: """ @@ -135,15 +143,9 @@ async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Me """ message = normalized_conversation[-1] message_piece: MessagePiece = message.message_pieces[0] - - # If user didn't set file_path, see if the PDF path is in converted_value - if not self.file_path: - possible_path = message_piece.converted_value - if isinstance(possible_path, str) and Path(possible_path).exists(): - logger.info(f"HTTPXApiTarget: auto-using file_path from {possible_path}") - self.file_path = possible_path - elif not Path(self.file_path).exists(): - raise FileNotFoundError(f"File not found: {self.file_path}") + upload_path = await self._get_upload_path_async(message_piece=message_piece) + if upload_path and self.method not in {"POST", "PUT"}: + raise ValueError(f"File uploads are not allowed with HTTP method: {self.method}") if not self.http_url: raise ValueError("No `http_url` provided for HTTPXApiTarget.") @@ -152,12 +154,12 @@ async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Me async with httpx.AsyncClient(http2=http2_version, **self.httpx_client_kwargs) as client: try: - if self.file_path and Path(self.file_path).exists(): + if upload_path: # Handle file upload (only for POST & PUT) - filename = Path(self.file_path).name + filename = upload_path.name mime_type = mimetypes.guess_type(filename)[0] or "application/octet-stream" - async with aiofiles.open(self.file_path, "rb") as fp: + async with aiofiles.open(upload_path, "rb") as fp: file_bytes = await fp.read() files = {"file": (filename, file_bytes, mime_type)} @@ -170,7 +172,7 @@ async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Me headers=self.headers, params=self.params, files=files, - follow_redirects=True, + follow_redirects=self.follow_redirects, ) else: # No file upload, handle based on HTTP method @@ -182,7 +184,7 @@ async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Me params=self.params, json=self.json_data if self.method in {"POST", "PUT", "PATCH"} else None, data=self.form_data if self.method in {"POST", "PUT", "PATCH"} else None, - follow_redirects=True, + follow_redirects=self.follow_redirects, ) except httpx.TimeoutException: @@ -207,3 +209,53 @@ async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Me ) return [response_entry] + + async def _get_upload_path_async(self, *, message_piece: MessagePiece) -> Path | None: + if self.file_path: + candidate = Path(self.file_path) + elif message_piece.converted_value_data_type in self._PATH_TYPES: + candidate = Path(message_piece.converted_value) + elif await asyncio.to_thread(Path(message_piece.converted_value).is_file): + candidate = Path(message_piece.converted_value) + print_deprecation_message( + old_item="HTTPXAPITarget implicit text-path uploads", + new_item="a path-typed MessagePiece or explicit file_path", + removed_in="1.3.0", + ) + else: + return None + + if self.allowed_upload_directory is None: + print_deprecation_message( + old_item="HTTPXAPITarget file uploads without allowed_upload_directory", + new_item="HTTPXAPITarget(..., allowed_upload_directory=...)", + removed_in="1.3.0", + ) + candidate = await asyncio.to_thread(candidate.resolve) + else: + candidate = await asyncio.to_thread( + self._validate_upload_path, + path=candidate, + allowed_directory=self.allowed_upload_directory, + ) + if not await asyncio.to_thread(candidate.is_file): + raise FileNotFoundError(f"File not found: {candidate}") + return candidate + + def _validate_upload_path(self, *, path: Path, allowed_directory: Path) -> Path: + resolved_path = path.resolve() + try: + resolved_path.relative_to(allowed_directory) + except ValueError as exc: + raise ValueError(f"File upload path is outside the allowed upload directory: {resolved_path}") from exc + return resolved_path + + @staticmethod + def _resolve_allowed_upload_directory(allowed_directory: str | Path | None) -> Path | None: + if allowed_directory is None: + return None + + resolved_directory = Path(allowed_directory).resolve() + if not resolved_directory.is_dir(): + raise ValueError(f"Allowed upload directory does not exist or is not a directory: {allowed_directory}") + return resolved_directory diff --git a/tests/integration/ai_recruiter/test_ai_recruiter.py b/tests/integration/ai_recruiter/test_ai_recruiter.py index 2bb756da6c..864d51bac2 100644 --- a/tests/integration/ai_recruiter/test_ai_recruiter.py +++ b/tests/integration/ai_recruiter/test_ai_recruiter.py @@ -11,7 +11,7 @@ import pytest import requests -from pyrit.common.path import CONVERTER_SEED_PROMPT_PATH, HOME_PATH +from pyrit.common.path import CONVERTER_SEED_PROMPT_PATH, DB_DATA_PATH, HOME_PATH from pyrit.converter import PDFConverter from pyrit.exceptions import PyritException from pyrit.executor.core import StrategyConverterConfig @@ -193,6 +193,7 @@ async def test_ai_recruiter_workflow(): upload_target = HTTPXAPITarget( http_url="http://localhost:8000/upload/", method="POST", + allowed_upload_directory=DB_DATA_PATH, timeout=180, ) diff --git a/tests/unit/prompt_target/target/test_http_api_target.py b/tests/unit/prompt_target/target/test_http_api_target.py index 84c4961f62..111d1e5908 100644 --- a/tests/unit/prompt_target/target/test_http_api_target.py +++ b/tests/unit/prompt_target/target/test_http_api_target.py @@ -2,6 +2,7 @@ # Licensed under the MIT license. import os import tempfile +from pathlib import Path from unittest.mock import MagicMock, patch import pytest @@ -28,7 +29,13 @@ async def test_send_prompt_async_file_upload(mock_request, patch_central_databas mock_request.return_value = mock_response # Create HTTPXAPITarget without passing a transport. - target = HTTPXAPITarget(http_url="http://example.com/upload/", method="POST", timeout=180) + target = HTTPXAPITarget( + http_url="http://example.com/upload/", + method="POST", + file_path=file_path, + allowed_upload_directory=Path(file_path).parent, + timeout=180, + ) response = await target.send_prompt_async(message=message) # Our mock transport returns a JSON string containing "File uploaded successfully". @@ -61,10 +68,12 @@ async def test_send_prompt_async_file_upload_preserves_query_params(mock_request target = HTTPXAPITarget( http_url="http://example.com/upload/", method="POST", + allowed_upload_directory=Path(file_path).parent, params={"alpha": "1"}, timeout=180, ) - await target.send_prompt_async(message=message) + with pytest.warns(DeprecationWarning, match="implicit text-path uploads"): + await target.send_prompt_async(message=message) assert mock_request.call_args.kwargs["params"] == {"alpha": "1"} @@ -125,14 +134,34 @@ async def test_send_prompt_async_preserves_query_params_for_post(mock_request, p @patch("httpx.AsyncClient.request") -async def test_send_prompt_async_missing_explicit_file_path_raises(mock_request, patch_central_database): +async def test_send_prompt_async_follows_redirects_when_enabled(mock_request, patch_central_database): + message_piece = MessagePiece(role="user", original_value="prompt", converted_value="prompt") + message = Message(message_pieces=[message_piece]) + mock_response = MagicMock() + mock_response.content = b'{"status": "ok"}' + mock_request.return_value = mock_response + target = HTTPXAPITarget( + http_url="http://example.com/data/", + method="POST", + follow_redirects=True, + ) + + await target.send_prompt_async(message=message) + + assert mock_request.call_args.kwargs["follow_redirects"] is True + + +@patch("httpx.AsyncClient.request") +async def test_send_prompt_async_missing_explicit_file_path_raises(mock_request, patch_central_database, tmp_path): message_piece = MessagePiece(role="user", original_value="mock", converted_value="trigger") message = Message(message_pieces=[message_piece]) + missing_file = tmp_path / "missing.pdf" target = HTTPXAPITarget( http_url="http://example.com/upload/", method="POST", - file_path="/definitely/missing/file.pdf", + file_path=str(missing_file), + allowed_upload_directory=tmp_path, timeout=180, ) @@ -177,7 +206,12 @@ async def test_send_prompt_async_binary_path_upload(mock_request, patch_central_ mock_response.content = b'{"message": "File uploaded successfully", "filename": "mock.pdf"}' mock_request.return_value = mock_response - target = HTTPXAPITarget(http_url="http://example.com/upload/", method="POST", timeout=180) + target = HTTPXAPITarget( + http_url="http://example.com/upload/", + method="POST", + allowed_upload_directory=Path(file_path).parent, + timeout=180, + ) # Must not raise "This target supports only the following data types: ..." response = await target.send_prompt_async(message=message) @@ -189,3 +223,103 @@ async def test_send_prompt_async_binary_path_upload(mock_request, patch_central_ assert files["file"][0] == os.path.basename(file_path) os.unlink(file_path) + + +@patch("httpx.AsyncClient.request") +async def test_send_prompt_async_file_upload_without_allowed_directory_warns( + mock_request, patch_central_database, tmp_path +): + file_path = tmp_path / "document.pdf" + file_path.write_bytes(b"content") + message_piece = MessagePiece(role="user", original_value=str(file_path), converted_value=str(file_path)) + message = Message(message_pieces=[message_piece]) + mock_response = MagicMock() + mock_response.content = b"uploaded" + mock_request.return_value = mock_response + target = HTTPXAPITarget( + http_url="http://example.com/upload/", + method="POST", + follow_redirects=True, + ) + + with pytest.warns(DeprecationWarning) as warning_records: + await target.send_prompt_async(message=message) + + warning_messages = [str(record.message) for record in warning_records] + assert any("implicit text-path uploads" in message for message in warning_messages) + assert any("without allowed_upload_directory" in message for message in warning_messages) + mock_request.assert_called_once() + + +@patch("httpx.AsyncClient.request") +async def test_send_prompt_async_rejects_upload_outside_allowed_directory( + mock_request, patch_central_database, tmp_path +): + allowed_directory = tmp_path / "allowed" + allowed_directory.mkdir() + file_path = tmp_path / "outside.pdf" + file_path.write_bytes(b"content") + message_piece = MessagePiece(role="user", original_value="prompt", converted_value="prompt") + message = Message(message_pieces=[message_piece]) + target = HTTPXAPITarget( + http_url="http://example.com/upload/", + method="POST", + file_path=str(allowed_directory / ".." / file_path.name), + allowed_upload_directory=allowed_directory, + ) + + with pytest.raises(ValueError, match="outside the allowed upload directory"): + await target.send_prompt_async(message=message) + + mock_request.assert_not_called() + + +@patch("httpx.AsyncClient.request") +async def test_send_prompt_async_validates_path_before_file_exists(mock_request, patch_central_database, tmp_path): + allowed_directory = tmp_path / "allowed" + allowed_directory.mkdir() + outside_path = tmp_path / "missing.pdf" + message_piece = MessagePiece( + role="user", + original_value=str(outside_path), + original_value_data_type="binary_path", + converted_value=str(outside_path), + converted_value_data_type="binary_path", + ) + message = Message(message_pieces=[message_piece]) + target = HTTPXAPITarget( + http_url="http://example.com/upload/", + method="POST", + allowed_upload_directory=allowed_directory, + follow_redirects=True, + ) + + with pytest.raises(ValueError, match="outside the allowed upload directory"): + await target.send_prompt_async(message=message) + + mock_request.assert_not_called() + + +@patch("httpx.AsyncClient.request") +async def test_send_prompt_async_validates_upload_method_after_path(mock_request, patch_central_database, tmp_path): + file_path = tmp_path / "document.pdf" + file_path.write_bytes(b"content") + message_piece = MessagePiece( + role="user", + original_value=str(file_path), + original_value_data_type="binary_path", + converted_value=str(file_path), + converted_value_data_type="binary_path", + ) + message = Message(message_pieces=[message_piece]) + target = HTTPXAPITarget( + http_url="http://example.com/upload/", + method="GET", + allowed_upload_directory=tmp_path, + follow_redirects=True, + ) + + with pytest.raises(ValueError, match="File uploads are not allowed with HTTP method: GET"): + await target.send_prompt_async(message=message) + + mock_request.assert_not_called() diff --git a/tests/unit/prompt_target/target/test_http_target.py b/tests/unit/prompt_target/target/test_http_target.py index 5803ec38ee..efe7f9e1c0 100644 --- a/tests/unit/prompt_target/target/test_http_target.py +++ b/tests/unit/prompt_target/target/test_http_target.py @@ -7,7 +7,7 @@ import httpx import pytest -from pyrit.models import Message +from pyrit.models import Message, MessagePiece from pyrit.prompt_target.http_target.http_target import HTTPTarget from pyrit.prompt_target.http_target.http_target_callback_functions import ( get_http_target_json_response_callback_function, @@ -185,6 +185,132 @@ async def test_send_prompt_async_client_kwargs(patch_central_database): assert http_target._client is None +@patch("httpx.AsyncClient.request", new_callable=AsyncMock) +async def test_send_prompt_async_rejects_prompt_destination_change(mock_request, patch_central_database): + target = HTTPTarget(http_request="GET {PROMPT} HTTP/1.1\nHost: example.com\n\n") + message = Message( + message_pieces=[ + MessagePiece( + role="user", + original_value="https://attacker.example/path", + converted_value="https://attacker.example/path", + converted_value_data_type="text", + ) + ] + ) + + with pytest.raises(ValueError, match="cannot change the configured HTTP destination"): + await target.send_prompt_async(message=message) + + mock_request.assert_not_awaited() + + +@patch("httpx.AsyncClient.request", new_callable=AsyncMock) +async def test_send_prompt_async_allows_configured_internal_destination(mock_request, patch_central_database): + target = HTTPTarget(http_request="POST /api/{PROMPT} HTTP/1.1\nHost: 10.0.0.8:8080\n\n") + message = Message(message_pieces=[MessagePiece(role="user", original_value="jobs", converted_value="jobs")]) + mock_response = MagicMock() + mock_response.content = b"ok" + mock_request.return_value = mock_response + + await target.send_prompt_async(message=message) + + assert mock_request.call_args.kwargs["url"] == "https://10.0.0.8:8080/api/jobs" + + +@patch("httpx.AsyncClient.request", new_callable=AsyncMock) +async def test_send_prompt_async_follows_redirects_when_enabled(mock_request, patch_central_database): + target = HTTPTarget( + http_request="POST /api HTTP/1.1\nHost: example.com\n\n", + follow_redirects=True, + ) + message = Message(message_pieces=[MessagePiece(role="user", original_value="prompt")]) + mock_response = MagicMock() + mock_response.content = b"ok" + mock_request.return_value = mock_response + + await target.send_prompt_async(message=message) + + assert mock_request.call_args.kwargs["follow_redirects"] is True + + +@patch("httpx.AsyncClient.request", new_callable=AsyncMock) +async def test_send_prompt_async_disables_redirects_when_requested(mock_request, patch_central_database): + target = HTTPTarget( + http_request="POST /api HTTP/1.1\nHost: example.com\n\n", + follow_redirects=False, + ) + message = Message(message_pieces=[MessagePiece(role="user", original_value="prompt")]) + mock_response = MagicMock() + mock_response.content = b"ok" + mock_request.return_value = mock_response + + await target.send_prompt_async(message=message) + + assert mock_request.call_args.kwargs["follow_redirects"] is False + + +def test_http_target_omitted_redirect_setting_preserves_behavior(patch_central_database): + target = HTTPTarget(http_request="GET / HTTP/1.1\nHost: example.com\n\n") + assert target.follow_redirects is True + + +@pytest.mark.parametrize( + ("http_request", "prompt"), + [ + ("GET /search?q={PROMPT} HTTP/1.1\nHost: example.com\n\n", "first\nsecond"), + ("GET / HTTP/1.1\nHost: example.com\nX-Prompt: {PROMPT}\n\n", "first\nsecond"), + ("GET /search?q={PROMPT} HTTP/1.1\nHost: example.com\n\n", "first\rsecond"), + ("GET / HTTP/1.1\nHost: example.com\nX-Prompt: {PROMPT}\n\n", "first\rsecond"), + ], +) +@patch("httpx.AsyncClient.request", new_callable=AsyncMock) +async def test_send_prompt_async_rejects_newlines_outside_body( + mock_request, + patch_central_database, + http_request, + prompt, +): + target = HTTPTarget(http_request=http_request) + message = Message(message_pieces=[MessagePiece(role="user", original_value=prompt)]) + + with pytest.raises(ValueError, match="cannot contain CR or LF"): + await target.send_prompt_async(message=message) + + mock_request.assert_not_awaited() + + +@patch("httpx.AsyncClient.request", new_callable=AsyncMock) +async def test_send_prompt_async_rejects_newline_when_placeholder_spans_header_and_body( + mock_request, patch_central_database +): + target = HTTPTarget( + http_request="POST / HTTP/1.1\nHost: example.com\nX-Prompt: {PROMPT_HEADER}\n\n{PROMPT_BODY}", + prompt_regex_string=r"\{PROMPT_HEADER\}\n\n\{PROMPT_BODY\}", + ) + message = Message(message_pieces=[MessagePiece(role="user", original_value="first\nsecond")]) + + with pytest.raises(ValueError, match="cannot contain CR or LF"): + await target.send_prompt_async(message=message) + + mock_request.assert_not_awaited() + + +@patch("httpx.AsyncClient.request", new_callable=AsyncMock) +async def test_send_prompt_async_allows_multiline_body_prompt(mock_request, patch_central_database): + target = HTTPTarget( + http_request="POST / HTTP/1.1\nHost: example.com\nContent-Type: text/plain\n\nbefore:{PROMPT}:after" + ) + message = Message(message_pieces=[MessagePiece(role="user", original_value="first\nsecond")]) + mock_response = MagicMock() + mock_response.content = b"ok" + mock_request.return_value = mock_response + + await target.send_prompt_async(message=message) + + assert mock_request.call_args.kwargs["content"] == "before:first\nsecond:after" + + async def test_send_prompt_async_validation(mock_http_target): # Creating a Message with no pieces raises immediately with pytest.raises(ValueError, match="must have at least one message piece"): @@ -368,12 +494,14 @@ def return_parsed(response): use_tls=False, callback_function=return_parsed, max_requests_per_minute=10, + follow_redirects=True, **client_kwargs, ) assert target.http_request == http_request assert target.prompt_regex_string == "{PLACEHOLDER_PROMPT}" assert target.use_tls is False assert target.callback_function == return_parsed + assert target.follow_redirects is True assert target.httpx_client_kwargs == client_kwargs assert target._client is None