diff --git a/templates/agent/handler.py b/templates/agent/handler.py index f990c50..612df79 100644 --- a/templates/agent/handler.py +++ b/templates/agent/handler.py @@ -2,9 +2,11 @@ from aws_lambda_powertools.event_handler import BedrockAgentFunctionResolver from aws_lambda_powertools.utilities.data_classes import BedrockAgentEvent from aws_lambda_powertools.utilities.typing import LambdaContext +from pydantic import ValidationError from templates.agent.models import Item from templates.agent.settings import Settings +from templates.models import Entity from templates.repository import Repository settings = Settings() # type: ignore @@ -29,11 +31,20 @@ def get_item(item_id: str) -> dict: The item details or an error message. """ logger.info("Retrieving item", extra={"itemId": item_id}) + try: + Entity(id=item_id) + except ValidationError: + logger.warning("Invalid item ID provided", extra={"itemId": item_id}) + return {"error": f"Invalid item ID '{item_id}'"} + try: item = repository.get_item(item_id) if not item: return {"error": f"Item {item_id} not found"} return Item.model_validate(item).dump() + except ValidationError as error: + logger.error("Item validation failed", extra={"itemId": item_id}, exc_info=error) + return {"error": "Internal server error"} except Exception as error: logger.error("Failed to get item", extra={"itemId": item_id}, exc_info=error) return {"error": f"Failed to get item with ID '{item_id}'"} @@ -57,6 +68,9 @@ def create_item(item_id: str, name: str, description: str | None = None) -> dict item = Item(id=item_id, name=name, description=description).dump() repository.put_item(item) return item + except ValidationError as error: + logger.warning("Invalid item data provided", extra={"itemId": item_id}, exc_info=error) + return {"error": "Invalid item data"} except Exception as error: logger.error("Failed to create item", extra={"itemId": item_id}, exc_info=error) return {"error": f"Failed to create item with ID '{item_id}'"} diff --git a/templates/api/handler.py b/templates/api/handler.py index fe2e9b7..8831a8e 100644 --- a/templates/api/handler.py +++ b/templates/api/handler.py @@ -6,8 +6,8 @@ from templates.api.models import Item from templates.api.response import JsonResponse -from templates.models import Entity from templates.api.settings import Settings +from templates.models import Entity from templates.repository import Repository settings = Settings() @@ -33,7 +33,7 @@ def get_item(id: str) -> Response: try: Entity(id=id) except ValidationError: - return JsonResponse({"message": "Invalid item ID length"}, status_code=400) + return JsonResponse({"message": "Invalid item ID"}, status_code=400) try: if (item := repository.get_item(id)) is None: diff --git a/templates/graphql/handler.py b/templates/graphql/handler.py index e000537..92a9416 100644 --- a/templates/graphql/handler.py +++ b/templates/graphql/handler.py @@ -6,6 +6,7 @@ from templates.graphql.models import Item from templates.graphql.settings import Settings +from templates.models import Entity from templates.repository import Repository settings = Settings() # type: ignore @@ -29,13 +30,21 @@ def get_item(id: str) -> dict | None: Returns: The item if found, or None. """ + try: + Entity(id=id) + except ValidationError: + logger.warning("Invalid item ID provided", extra={"itemId": id}) + raise RuntimeError(f"Invalid item ID '{id}'") from None + try: if (item := repository.get_item(id)) is None: return None return Item.model_validate(item).dump() except Exception as error: - logger.error("Failed to get item", extra={"itemId": id}, exc_info=error) - raise RuntimeError(f"Failed to get item with ID '{id}'") from None + is_val_error = isinstance(error, ValidationError) + message = "Item validation failed" if is_val_error else f"Failed to get item with ID '{id}'" + logger.error(message, extra={"itemId": id}, exc_info=error) + raise RuntimeError(message) from None @app.resolver(type_name="Query", field_name="listItems") @@ -68,7 +77,10 @@ def create_item(name: str) -> dict: item = Item(name=name).dump() repository.put_item(item) return item - except (ValidationError, Exception) as error: + except ValidationError as error: + logger.warning("Invalid item data provided", extra={"itemName": name}, exc_info=error) + raise RuntimeError("Invalid item data") from None + except Exception as error: logger.error("Failed to create item", extra={"itemName": name}, exc_info=error) raise RuntimeError(f"Failed to create item with name '{name}'") from None diff --git a/templates/models.py b/templates/models.py index 5769ace..e2d4ad6 100644 --- a/templates/models.py +++ b/templates/models.py @@ -29,4 +29,5 @@ class Entity(Object): default_factory=lambda: str(uuid4()), min_length=1, max_length=50, + pattern=r"^[a-zA-Z0-9-_]+$", ) diff --git a/templates/sqs/handler.py b/templates/sqs/handler.py index 1bff9b4..990e543 100644 --- a/templates/sqs/handler.py +++ b/templates/sqs/handler.py @@ -43,7 +43,7 @@ def handle_record(self, record: SQSRecord) -> None: try: message = SqsMessage.model_validate_json(record.body) processed = ProcessedItem(id=message.id, content=message.content, status="PROCESSED") - self._repository.put_item(processed.model_dump()) + self._repository.put_item(processed.dump()) logger.info("Successfully processed and stored message", extra={"messageId": message.id}) except Exception as exc: logger.error("Failed to process SQS record", exc_info=exc) diff --git a/tests/agent/test_handler.py b/tests/agent/test_handler.py index 6859eb1..a8d5bf2 100644 --- a/tests/agent/test_handler.py +++ b/tests/agent/test_handler.py @@ -47,6 +47,15 @@ def test_handler_get_item_not_found(): assert "not found" in result["error"] +def test_handler_get_item_invalid_id(): + from templates.agent.handler import get_item + + result = get_item("invalid!") + + assert "error" in result + assert "Invalid item ID" in result["error"] + + def test_handler_create_item(repository): from templates.agent.handler import create_item diff --git a/tests/api/test_handler.py b/tests/api/test_handler.py index 4bc8378..1130bea 100644 --- a/tests/api/test_handler.py +++ b/tests/api/test_handler.py @@ -109,7 +109,19 @@ def test_get_item_id_too_long(mock_repo, lambda_context): assert response["statusCode"] == 400 body = loads(response["body"]) - assert body["message"] == "Invalid item ID length" + assert body["message"] == "Invalid item ID" + + +def test_get_item_id_invalid_pattern(mock_repo, lambda_context): + """GET /items/{id} returns 400 when the ID contains invalid characters.""" + import templates.api.handler as handler_module + + event = _apigw_event("GET", "/items/invalid!", path_params={"id": "invalid!"}) + response = handler_module.main(event, lambda_context) + + assert response["statusCode"] == 400 + body = loads(response["body"]) + assert body["message"] == "Invalid item ID" def test_post_item_invalid_body(mock_repo, lambda_context): diff --git a/tests/graphql/test_handler.py b/tests/graphql/test_handler.py index f0abdce..e7babf6 100644 --- a/tests/graphql/test_handler.py +++ b/tests/graphql/test_handler.py @@ -58,6 +58,17 @@ def test_sensitive_data_exposure(repository, lambda_context): assert "internal_secret" not in result_list[0] +def test_get_item_invalid_id(lambda_context): + import pytest + + from templates.graphql.handler import main + + event = {"info": {"parentTypeName": "Query", "fieldName": "getItem"}, "arguments": {"id": "invalid!"}} + with pytest.raises(RuntimeError) as excinfo: + main(event, lambda_context) + assert "Invalid item ID" in str(excinfo.value) + + def test_error_message_information_leakage(lambda_context, mocker): """Verify that internal error details are NOT leaked to the client.""" from templates.graphql import handler