diff --git a/pyrit/backend/models/attacks.py b/pyrit/backend/models/attacks.py index 3b280c09d6..3f8e12b386 100644 --- a/pyrit/backend/models/attacks.py +++ b/pyrit/backend/models/attacks.py @@ -16,6 +16,7 @@ from pyrit.backend.models._media import build_filename, infer_mime_type from pyrit.backend.models.common import ( + MAX_IDENTIFIER_LENGTH, MAX_ITEMS, IdentifierStr, LabelDict, @@ -600,7 +601,7 @@ class SaveConversationRequest(_AttackAttributionInput): source_conversation_id: uuid.UUID | None = None expected_objective: str | None = None objective: str | None = Field(None, description="Omit to keep the destination's objective unchanged") - target_registry_name: str | None = None + target_registry_name: IdentifierStr | None = None messages: list[MessageRequest] = Field(default_factory=list) @model_validator(mode="after") @@ -701,10 +702,16 @@ class AddMessageRequest(MessageRequest): description="Ordered registry-backed converter pipelines to apply to the response.", ) start_token: str = Field( - default="⟪", min_length=1, description="Opening marker for request and response converter pipelines" + default="⟪", + min_length=1, + max_length=MAX_IDENTIFIER_LENGTH, + description="Opening marker for request and response converter pipelines", ) end_token: str = Field( - default="⟫", min_length=1, description="Closing marker for request and response converter pipelines" + default="⟫", + min_length=1, + max_length=MAX_IDENTIFIER_LENGTH, + description="Closing marker for request and response converter pipelines", ) target_conversation_id: IdentifierStr = Field( ..., diff --git a/pyrit/backend/models/converters.py b/pyrit/backend/models/converters.py index 28c5a46f40..476d48d31b 100644 --- a/pyrit/backend/models/converters.py +++ b/pyrit/backend/models/converters.py @@ -11,7 +11,12 @@ from pydantic import BaseModel, Field -from pyrit.backend.models.common import MAX_ITEMS, REGISTRY_INSTANCE_NAME_PATTERN, IdentifierStr +from pyrit.backend.models.common import ( + MAX_IDENTIFIER_LENGTH, + MAX_ITEMS, + REGISTRY_INSTANCE_NAME_PATTERN, + IdentifierStr, +) from pyrit.models import ConverterIdentifier, Parameter, PromptDataType __all__ = [ @@ -119,8 +124,18 @@ class ConverterPreviewRequest(BaseModel): original_value: str = Field(..., description="Text to convert") original_value_data_type: PromptDataType = Field(default="text", description="Data type of original value") converter_ids: list[IdentifierStr] = Field(..., max_length=MAX_ITEMS, description="Converter instance IDs to apply") - start_token: str = Field(default="⟪", min_length=1, description="Opening marker for selected text regions") - end_token: str = Field(default="⟫", min_length=1, description="Closing marker for selected text regions") + start_token: str = Field( + default="⟪", + min_length=1, + max_length=MAX_IDENTIFIER_LENGTH, + description="Opening marker for selected text regions", + ) + end_token: str = Field( + default="⟫", + min_length=1, + max_length=MAX_IDENTIFIER_LENGTH, + description="Closing marker for selected text regions", + ) class ConverterPreviewResponse(BaseModel): diff --git a/pyrit/backend/models/scorers.py b/pyrit/backend/models/scorers.py index 2da03a09d8..966eba2779 100644 --- a/pyrit/backend/models/scorers.py +++ b/pyrit/backend/models/scorers.py @@ -5,7 +5,7 @@ from pydantic import BaseModel, Field -from pyrit.backend.models.common import REGISTRY_INSTANCE_NAME_PATTERN, PaginationInfo +from pyrit.backend.models.common import MAX_ITEMS, REGISTRY_INSTANCE_NAME_PATTERN, IdentifierStr, PaginationInfo from pyrit.models import JSONValue, Parameter from pyrit.models.catalog.scorer import ScorerInstance @@ -36,5 +36,7 @@ class CreateScorerRequest(BaseModel): """Request to construct and register a named scorer.""" name: str = Field(..., min_length=1, pattern=REGISTRY_INSTANCE_NAME_PATTERN, description="Unique registry name") - type: str = Field(..., description="Scorer class name") - params: dict[str, JSONValue] = Field(default_factory=dict, description="Scorer constructor parameters") + type: IdentifierStr = Field(..., description="Scorer class name") + params: dict[IdentifierStr, JSONValue] = Field( + default_factory=dict, max_length=MAX_ITEMS, description="Scorer constructor parameters" + ) diff --git a/pyrit/backend/models/techniques.py b/pyrit/backend/models/techniques.py index 1e3aebaf63..35e64d85a4 100644 --- a/pyrit/backend/models/techniques.py +++ b/pyrit/backend/models/techniques.py @@ -7,7 +7,7 @@ from pydantic import BaseModel, ConfigDict, Field, JsonValue, field_validator -from pyrit.backend.models.common import PaginationInfo +from pyrit.backend.models.common import IdentifierStr, PaginationInfo, TextStr from pyrit.models import Parameter from pyrit.models.request_limits import MAX_ITEMS @@ -56,12 +56,12 @@ class CreateTechniqueRequest(BaseModel): name: TechniqueSelector type: str = Field(min_length=1, max_length=256) - params: dict[str, JsonValue] = Field(default_factory=dict, max_length=MAX_ITEMS) - description: str | None = None + params: dict[IdentifierStr, JsonValue] = Field(default_factory=dict, max_length=MAX_ITEMS) + description: TextStr | None = None tags: list[TechniqueSelector] = Field(default_factory=list, max_length=MAX_ITEMS) - request_converters: list[str] | None = Field(default=None, max_length=MAX_ITEMS) - response_converters: list[str] | None = Field(default=None, max_length=MAX_ITEMS) - adversarial_chat: str | None = None + request_converters: list[IdentifierStr] | None = Field(default=None, max_length=MAX_ITEMS) + response_converters: list[IdentifierStr] | None = Field(default=None, max_length=MAX_ITEMS) + adversarial_chat: IdentifierStr | None = None adversarial_system_prompt: str | None = None adversarial_seed_prompt: str | None = None adversarial_prompt_template: str | None = None diff --git a/pyrit/backend/routes/attacks.py b/pyrit/backend/routes/attacks.py index bb75f1f33b..82766a713c 100644 --- a/pyrit/backend/routes/attacks.py +++ b/pyrit/backend/routes/attacks.py @@ -33,6 +33,7 @@ UpdateMainConversationResponse, ) from pyrit.backend.models.common import ( + MAX_IDENTIFIER_LENGTH, MAX_ITEMS, CursorStr, IdentifierStr, @@ -352,7 +353,9 @@ async def remove_human_score(attack_result_id: IdentifierStr) -> AttackSummary: ) async def get_conversation_messages( # pyrit-async-suffix-exempt attack_result_id: IdentifierStr, - conversation_id: IdentifierStr = Query(..., description="The conversation_id whose messages to return"), + conversation_id: str = Query( + ..., max_length=MAX_IDENTIFIER_LENGTH, description="The conversation_id whose messages to return" + ), ) -> ConversationMessagesResponse: """ Get all messages for a conversation belonging to an attack. diff --git a/pyrit/backend/routes/datasets.py b/pyrit/backend/routes/datasets.py index 7dc342be94..5b0ea397d6 100644 --- a/pyrit/backend/routes/datasets.py +++ b/pyrit/backend/routes/datasets.py @@ -8,11 +8,19 @@ discovered from registered ``SeedDatasetProvider`` subclasses and from memory. """ +from typing import Annotated from uuid import UUID from fastapi import APIRouter, HTTPException, Query, status - -from pyrit.backend.models.common import MAX_ITEMS, CursorStr, ProblemDetail +from pydantic import Field + +from pyrit.backend.models.common import ( + MAX_CURSOR_LENGTH, + MAX_ITEMS, + MAX_LABEL_VALUE_LENGTH, + CursorStr, + ProblemDetail, +) from pyrit.backend.models.datasets import ( DatasetListResponse, SeedExampleDetailResponse, @@ -55,7 +63,7 @@ async def list_datasets(loaded_only: bool = False) -> DatasetListResponse: # py }, ) async def list_seed_examples( # pyrit-async-suffix-exempt - selection_key: str = Query(..., description=_SELECTION_KEY_DESCRIPTION), + selection_key: str = Query(..., max_length=MAX_CURSOR_LENGTH, description=_SELECTION_KEY_DESCRIPTION), limit: int = Query(20, ge=1, le=100, description="Maximum examples per page"), cursor: CursorStr | None = Query( None, @@ -68,7 +76,7 @@ async def list_seed_examples( # pyrit-async-suffix-exempt description="Case-insensitive literal text to find in text prompt and objective values", ), modality: list[PromptDataType] | None = Query(None, max_length=MAX_ITEMS, description="Data types, OR-matched"), - harm_category: list[str] | None = Query( + harm_category: list[Annotated[str, Field(max_length=MAX_LABEL_VALUE_LENGTH)]] | None = Query( None, max_length=MAX_ITEMS, description="Whole harm categories, case-insensitive, OR-matched" ), seed_type: list[SeedType] | None = Query(None, max_length=MAX_ITEMS, description="Seed types, OR-matched"), @@ -103,7 +111,7 @@ async def list_seed_examples( # pyrit-async-suffix-exempt ) async def get_seed_example( # pyrit-async-suffix-exempt example_id: UUID, - selection_key: str = Query(..., description=_SELECTION_KEY_DESCRIPTION), + selection_key: str = Query(..., max_length=MAX_CURSOR_LENGTH, description=_SELECTION_KEY_DESCRIPTION), ) -> SeedExampleDetailResponse: """ Get one stored logical seed example with all of its members. diff --git a/pyrit/backend/routes/message_sends.py b/pyrit/backend/routes/message_sends.py index 6099305c10..73460e3b13 100644 --- a/pyrit/backend/routes/message_sends.py +++ b/pyrit/backend/routes/message_sends.py @@ -8,7 +8,7 @@ from fastapi import APIRouter, HTTPException, Query, status -from pyrit.backend.models.common import ProblemDetail +from pyrit.backend.models.common import IdentifierStr, ProblemDetail from pyrit.backend.models.message_sends import MessageSendRequest, MessageSendStatus from pyrit.backend.services.attack_service import get_attack_service from pyrit.backend.services.manual_send_scheduler import ManualSendConflictError, ManualSendQueueFullError @@ -30,7 +30,9 @@ 429: {"model": ProblemDetail, "description": "Manual-message admission unavailable"}, }, ) -async def submit_message_send_async(*, attack_result_id: str, request: MessageSendRequest) -> MessageSendStatus: +async def submit_message_send_async( + *, attack_result_id: IdentifierStr, request: MessageSendRequest +) -> MessageSendStatus: """Return a transient handle after validation and admission, before preparation or target I/O.""" try: return await get_attack_service().submit_message_send_async(attack_result_id=attack_result_id, request=request) @@ -51,7 +53,7 @@ async def submit_message_send_async(*, attack_result_id: str, request: MessageSe responses={404: {"model": ProblemDetail, "description": "Handle expired, missing, or for another attack"}}, ) async def get_message_send_async( - *, attack_result_id: str, send_id: str, wait_ms: Annotated[int, Query(ge=0, le=1000)] = 0 + *, attack_result_id: IdentifierStr, send_id: IdentifierStr, wait_ms: Annotated[int, Query(ge=0, le=1000)] = 0 ) -> MessageSendStatus: """ Read progress, optionally waiting briefly without cancelling or resubmitting the accepted work. diff --git a/pyrit/backend/routes/scorers.py b/pyrit/backend/routes/scorers.py index 6635f10eff..60e8b15dee 100644 --- a/pyrit/backend/routes/scorers.py +++ b/pyrit/backend/routes/scorers.py @@ -5,7 +5,7 @@ from fastapi import APIRouter, HTTPException, Query, status -from pyrit.backend.models.common import ProblemDetail +from pyrit.backend.models.common import CursorStr, IdentifierStr, ProblemDetail from pyrit.backend.models.scorers import ( CreateScorerRequest, ScorerListResponse, @@ -31,7 +31,7 @@ async def list_scorer_types() -> ScorerTypeResponse: # pyrit-async-suffix-exemp @router.get("", response_model=ScorerListResponse) async def list_scorers( limit: int = Query(50, ge=1, le=200, description="Maximum items per page"), - cursor: str | None = Query(None, description="Scorer registry name to start after"), + cursor: CursorStr | None = Query(None, description="Scorer registry name to start after"), ) -> ScorerListResponse: # pyrit-async-suffix-exempt """ List named scorer instances with pagination. @@ -71,7 +71,7 @@ async def create_scorer(request: CreateScorerRequest) -> ScorerInstance: # pyri response_model=ScorerInstance, responses={404: {"model": ProblemDetail, "description": "Scorer not found"}}, ) -async def get_scorer(scorer_registry_name: str) -> ScorerInstance: # pyrit-async-suffix-exempt +async def get_scorer(scorer_registry_name: IdentifierStr) -> ScorerInstance: # pyrit-async-suffix-exempt """ Get a named scorer instance. diff --git a/tests/unit/backend/test_api_routes.py b/tests/unit/backend/test_api_routes.py index 2fe2d9286d..077102d445 100644 --- a/tests/unit/backend/test_api_routes.py +++ b/tests/unit/backend/test_api_routes.py @@ -870,6 +870,15 @@ def test_get_attack_rejects_oversized_id(self, client: TestClient) -> None: assert response.status_code == status.HTTP_422_UNPROCESSABLE_CONTENT mock_get_service.assert_not_called() + def test_get_conversation_messages_rejects_oversized_conversation_id(self, client: TestClient) -> None: + with patch("pyrit.backend.routes.attacks.get_attack_service") as mock_get_service: + response = client.get( + "/api/attacks/attack/messages", params={"conversation_id": "c" * (MAX_IDENTIFIER_LENGTH + 1)} + ) + + assert response.status_code == status.HTTP_422_UNPROCESSABLE_CONTENT + mock_get_service.assert_not_called() + def test_parse_labels_value_with_extra_colons(self, client: TestClient) -> None: """Test that _parse_labels handles values containing colons (split on first only).""" with patch("pyrit.backend.routes.attacks.get_attack_service") as mock_get_service: diff --git a/tests/unit/backend/test_common_models.py b/tests/unit/backend/test_common_models.py index 10c062fec8..0e63b7958a 100644 --- a/tests/unit/backend/test_common_models.py +++ b/tests/unit/backend/test_common_models.py @@ -14,6 +14,7 @@ AddMessageRequest, CreateAttackRequest, MessagePieceRequest, + SaveConversationRequest, UpdateAttackRequest, ) from pyrit.backend.models.common import ( @@ -38,6 +39,7 @@ from pyrit.backend.models.initializers import RegisterInitializerRequest from pyrit.backend.models.scores import ManualScoreRequest from pyrit.backend.models.targets import CreateTargetRequest +from pyrit.backend.models.techniques import CreateTechniqueRequest class TestPaginationInfo: @@ -446,6 +448,8 @@ def test_update_attack_request_expected_objective_is_not_length_limited() -> Non {"converter_ids": ["converter"] * (MAX_ITEMS + 1)}, {"request_converter_configurations": [{"converter_ids": ["converter"]}] * (MAX_ITEMS + 1)}, {"pieces": [{"original_value": "text", "mime_type": "m" * (MAX_IDENTIFIER_LENGTH + 1)}]}, + {"start_token": "s" * (MAX_IDENTIFIER_LENGTH + 1)}, + {"end_token": "e" * (MAX_IDENTIFIER_LENGTH + 1)}, ], ) def test_add_message_request_rejects_values_over_limits(overrides: dict[str, object]) -> None: @@ -478,7 +482,51 @@ def test_message_piece_content_is_not_length_limited() -> None: "params", ), (ConverterPreviewRequest, {"original_value": "x", "converter_ids": ["c"] * (MAX_ITEMS + 1)}, "converter_ids"), + ( + ConverterPreviewRequest, + {"original_value": "x", "converter_ids": [], "start_token": "s" * (MAX_IDENTIFIER_LENGTH + 1)}, + "start_token", + ), + ( + ConverterPreviewRequest, + {"original_value": "x", "converter_ids": [], "end_token": "e" * (MAX_IDENTIFIER_LENGTH + 1)}, + "end_token", + ), + ( + CreateTechniqueRequest, + {"name": "t", "type": "PromptSendingAttack", "description": "d" * (MAX_TEXT_LENGTH + 1)}, + "description", + ), + ( + CreateTechniqueRequest, + {"name": "t", "type": "RedTeamingAttack", "adversarial_chat": "a" * (MAX_IDENTIFIER_LENGTH + 1)}, + "adversarial_chat", + ), + ( + CreateTechniqueRequest, + {"name": "t", "type": "PromptSendingAttack", "request_converters": ["c" * (MAX_IDENTIFIER_LENGTH + 1)]}, + "request_converters", + ), + ( + CreateTechniqueRequest, + {"name": "t", "type": "PromptSendingAttack", "response_converters": ["c" * (MAX_IDENTIFIER_LENGTH + 1)]}, + "response_converters", + ), + ( + CreateTechniqueRequest, + {"name": "t", "type": "PromptSendingAttack", "params": {"p" * (MAX_IDENTIFIER_LENGTH + 1): 1}}, + "params", + ), (CreateTargetRequest, {"type": "t" * (MAX_IDENTIFIER_LENGTH + 1)}, "type"), + ( + SaveConversationRequest, + { + "save_id": str(uuid.uuid4()), + "destination": "new_attack", + "target_registry_name": "t" * (MAX_IDENTIFIER_LENGTH + 1), + }, + "target_registry_name", + ), ( ManualScoreRequest, { @@ -500,6 +548,20 @@ def test_request_models_reject_values_over_limits( assert [item["loc"][0] for item in error.value.errors()] == [field] +def test_create_technique_request_accepts_values_at_limits() -> None: + request = CreateTechniqueRequest( + name="t", + type="RedTeamingAttack", + description="d" * MAX_TEXT_LENGTH, + adversarial_chat="a" * MAX_IDENTIFIER_LENGTH, + request_converters=["c" * MAX_IDENTIFIER_LENGTH], + params={"p" * MAX_IDENTIFIER_LENGTH: 1}, + adversarial_system_prompt="s" * (MAX_TEXT_LENGTH + 1), + ) + + assert len(request.adversarial_system_prompt or "") == MAX_TEXT_LENGTH + 1 + + def test_label_filter_limits_ignore_surrounding_whitespace() -> None: padded = f" {'k' * MAX_LABEL_KEY_LENGTH} : {'v' * MAX_LABEL_VALUE_LENGTH} " diff --git a/tests/unit/backend/test_message_send_routes.py b/tests/unit/backend/test_message_send_routes.py index 43d2b0e3e7..e77127da77 100644 --- a/tests/unit/backend/test_message_send_routes.py +++ b/tests/unit/backend/test_message_send_routes.py @@ -10,6 +10,7 @@ from fastapi.testclient import TestClient from pyrit.backend.main import app +from pyrit.backend.models.common import MAX_IDENTIFIER_LENGTH from pyrit.backend.models.message_sends import MessageSendStatus from pyrit.backend.services.attack_service import AttackService from pyrit.backend.services.manual_send_scheduler import ManualSendConflictError, ManualSendQueueFullError @@ -146,3 +147,14 @@ def test_missing_handle_is_not_recreated(*, client: TestClient, sender: MagicMoc assert response.status_code == 404 assert "do not resend" in response.json()["detail"] sender.submit_async.assert_not_awaited() + + +def test_oversized_identifiers_are_rejected_before_the_services( + *, client: TestClient, sender: MagicMock, payload: dict[str, object] +) -> None: + oversized = "a" * (MAX_IDENTIFIER_LENGTH + 1) + + assert client.post(f"/api/attacks/{oversized}/message-sends", json=payload).status_code == 422 + assert client.get(f"/api/attacks/attack/message-sends/{oversized}").status_code == 422 + sender.submit_async.assert_not_awaited() + sender.get_status_async.assert_not_awaited() diff --git a/tests/unit/backend/test_scorer_routes.py b/tests/unit/backend/test_scorer_routes.py index 25339d3267..b58176f721 100644 --- a/tests/unit/backend/test_scorer_routes.py +++ b/tests/unit/backend/test_scorer_routes.py @@ -10,6 +10,7 @@ from fastapi.testclient import TestClient from pyrit.backend.main import app +from pyrit.backend.models.common import MAX_CURSOR_LENGTH, MAX_IDENTIFIER_LENGTH, MAX_ITEMS from pyrit.backend.services.scorer_service import get_scorer_service from pyrit.models import ComponentIdentifier, Scorable, Score, ScoringExpectation, TargetIdentifier from pyrit.prompt_target import PromptTarget @@ -199,6 +200,17 @@ def test_duplicate_name_returns_400_and_does_not_replace_instance(compatibility_ ("/api/scorers", {"name": "invalid name", "type": "_RouteLeafScorer", "params": {}}), ("/api/scorers?limit=0", None), ("/api/scorers?limit=201", None), + ("/api/scorers", {"name": "valid", "type": "t" * (MAX_IDENTIFIER_LENGTH + 1), "params": {}}), + ( + "/api/scorers", + {"name": "valid", "type": "_RouteLeafScorer", "params": {"p" * (MAX_IDENTIFIER_LENGTH + 1): 1}}, + ), + ( + "/api/scorers", + {"name": "valid", "type": "_RouteLeafScorer", "params": {f"p{i}": i for i in range(MAX_ITEMS + 1)}}, + ), + (f"/api/scorers?cursor={'c' * (MAX_CURSOR_LENGTH + 1)}", None), + (f"/api/scorers/{'s' * (MAX_IDENTIFIER_LENGTH + 1)}", None), ], ) def test_invalid_scorer_request_shapes_return_422( diff --git a/tests/unit/backend/test_seed_example_routes.py b/tests/unit/backend/test_seed_example_routes.py index 0f39b7a9e8..683a7873be 100644 --- a/tests/unit/backend/test_seed_example_routes.py +++ b/tests/unit/backend/test_seed_example_routes.py @@ -11,6 +11,7 @@ from httpx import ASGITransport, AsyncClient from pyrit.backend.main import app +from pyrit.backend.models.common import MAX_CURSOR_LENGTH, MAX_LABEL_VALUE_LENGTH from pyrit.backend.services.dataset_service import get_dataset_service from pyrit.common.utils import to_sha256 from pyrit.memory import MemoryInterface @@ -84,6 +85,9 @@ async def test_list_seed_examples_pages_with_a_filter_bound_cursor(client: Async ("", {"selection_key": DATASET}, 400), ("", {"selection_key": "dataset:named:"}, 400), ("", {"selection_key": NAMED, "limit": 0}, 422), + ("", {"selection_key": "k" * (MAX_CURSOR_LENGTH + 1)}, 422), + ("", {"selection_key": NAMED, "harm_category": "h" * (MAX_LABEL_VALUE_LENGTH + 1)}, 422), + (f"/{MISSING_ID}", {"selection_key": "k" * (MAX_CURSOR_LENGTH + 1)}, 422), (f"/{MISSING_ID}", {"selection_key": "dataset:other"}, 400), (f"/{MISSING_ID}", {"selection_key": NAMED}, 404), ("/not-a-uuid", {"selection_key": NAMED}, 422),