diff --git a/templates/agent/handler.py b/templates/agent/handler.py index c5d4ee6..4ffa510 100644 --- a/templates/agent/handler.py +++ b/templates/agent/handler.py @@ -29,10 +29,14 @@ def get_item(item_id: str) -> dict: The item details or an error message. """ logger.info(f"Retrieving item {item_id}") - item = repository.get_item(item_id) - if not item: - return {"error": f"Item {item_id} not found"} - return Item.model_validate(item).dump() + 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 Exception as error: + logger.error(f"Failed to get item with ID '{item_id}'", exc_info=error) + return {"error": f"Failed to get item with ID '{item_id}'"} @tracer.capture_method @@ -49,9 +53,13 @@ def create_item(item_id: str, name: str, description: str | None = None) -> dict The created item details. """ logger.info(f"Creating item {item_id}") - item = Item(id=item_id, name=name, description=description) - repository.put_item(item.model_dump()) - return item.model_dump(by_alias=True, exclude_none=True) + try: + item = Item(id=item_id, name=name, description=description).dump() + repository.put_item(item) + return item + except Exception as error: + logger.error(f"Failed to create item with ID '{item_id}'", exc_info=error) + return {"error": f"Failed to create item with ID '{item_id}'"} @logger.inject_lambda_context diff --git a/templates/graphql/handler.py b/templates/graphql/handler.py index 37668a1..72b16a7 100644 --- a/templates/graphql/handler.py +++ b/templates/graphql/handler.py @@ -35,7 +35,7 @@ def get_item(id: str) -> dict | None: return Item.model_validate(item).dump() except Exception as error: logger.error(f"Failed to get item with ID '{id}'", exc_info=error) - raise RuntimeError(f"Failed to get item with ID '{id}'") from error + raise RuntimeError(f"Failed to get item with ID '{id}'") from None @app.resolver(type_name="Query", field_name="listItems") @@ -50,7 +50,7 @@ def list_items() -> list[dict]: return [Item.model_validate(item).dump() for item in repository.list_items()] except Exception as error: logger.error("Failed to list items", exc_info=error) - raise RuntimeError("Failed to list items") from error + raise RuntimeError("Failed to list items") from None @app.resolver(type_name="Mutation", field_name="createItem") @@ -70,7 +70,7 @@ def create_item(name: str) -> dict: return item except (ValidationError, Exception) as error: logger.error(f"Failed to create item with name '{name}'", exc_info=error) - raise RuntimeError(f"Failed to create item with name '{name}'") from error + raise RuntimeError(f"Failed to create item with name '{name}'") from None @logger.inject_lambda_context(correlation_id_path=correlation_paths.APPSYNC_RESOLVER) diff --git a/templates/stream/handler.py b/templates/stream/handler.py index 906c657..0ddcbf6 100644 --- a/templates/stream/handler.py +++ b/templates/stream/handler.py @@ -70,7 +70,7 @@ def handle_record(self, record: DynamoDBRecord) -> None: self._repository.put_item(item.model_dump()) elif event_name and event_name.name == "REMOVE": plain_keys = SourceItem.model_validate(record.dynamodb.keys) - self._repository.delete_item(plain_keys.model_dump(exclude_none=True)) + self._repository.delete_item(plain_keys.id) handler = Handler(repository) diff --git a/tests/agent/test_handler.py b/tests/agent/test_handler.py index 109ed7c..6859eb1 100644 --- a/tests/agent/test_handler.py +++ b/tests/agent/test_handler.py @@ -106,5 +106,19 @@ def test_sensitive_data_exposure(repository): assert "internal_secret" not in result +def test_error_handling_sanitization(mocker): + """Verify that internal error details are NOT leaked to the agent.""" + from templates.agent import handler + from templates.agent.handler import get_item + + mocker.patch.object(handler.repository, "get_item", side_effect=Exception("Database connection failed")) + + result = get_item("123") + + assert "error" in result + assert "Database connection failed" not in result["error"] + assert "Failed to get item with ID '123'" in result["error"] + + if __name__ == "__main__": main() diff --git a/tests/graphql/test_handler.py b/tests/graphql/test_handler.py index b84c7bc..f0abdce 100644 --- a/tests/graphql/test_handler.py +++ b/tests/graphql/test_handler.py @@ -72,6 +72,7 @@ def test_error_message_information_leakage(lambda_context, mocker): assert "Database connection failed" not in str(excinfo.value) assert "Cause:" not in str(excinfo.value) + assert excinfo.value.__cause__ is None if __name__ == "__main__": diff --git a/tests/stream/test_handler.py b/tests/stream/test_handler.py index 19fa59a..417dc2e 100644 --- a/tests/stream/test_handler.py +++ b/tests/stream/test_handler.py @@ -104,7 +104,7 @@ def test_remove_record_calls_delete_item(mock_repo, lambda_context): event = _stream_event(_remove_record("abc")) result = handler_module.main(event, lambda_context) - mock_repo.delete_item.assert_called_once_with({"id": "abc"}) + mock_repo.delete_item.assert_called_once_with("abc") mock_repo.put_item.assert_not_called() assert result == {"batchItemFailures": []}