Enhance service bus handling of invalid JSON in receive_message function

CHANGELOG.md

@@ -18,6 +18,7 @@ ENHANCEMENTS:
 * Migration to Pydantic v2: Updates codebase to be compatible with Pydantic v2 for future FastAPI upgrades ([#4637](https://github.com/microsoft/AzureTRE/issues/4637))
 
 BUG FIXES:
+* Fix to enhance service bus handling of invalid JSON in receive_message function ([#4932](https://github.com/microsoft/AzureTRE/pull/4932))
 * Ignore changes to `ip_tags` on public IP resources to unblock deployments where these tags are set by Azure policy. (`core` 0.16.17, `tre-shared-service-certs` 0.7.11) ([#5019](https://github.com/microsoft/AzureTRE/issues/5019))
 * Fix workspace deletion when backup is enabled for the base, unrestricted and airlock-import-review workspaces by adding a `delete_backups_on_uninstall` flag and a pre-teardown backup cleanup (`remove_backup.sh`) that stops protection and either deletes or retains the Recovery Services Vault, so deletion works with Azure secure-by-default soft delete ([#4962](https://github.com/microsoft/AzureTRE/issues/4962))
 * Fix Nexus shared service security: fetch admin password from Key Vault at runtime via managed identity (IMDS) instead of embedding it in the VM Run Command script content. Fix `deploy_nexus_container.sh` short-circuit path to fail loudly if the container does not start. (`sonatype-nexus` 3.10.0) ([#4983](https://github.com/microsoft/AzureTRE/pull/4983))

resource_processor/_version.py

@@ -1 +1 @@
-__version__ = "0.13.5"
+__version__ = "0.13.6"

resource_processor/tests_rp/test_runner.py

@@ -115,7 +115,15 @@ async def test_receive_message(mock_invoke_porter_action, mock_service_bus_clien
     mock_receiver.__aexit__.return_value = None
     mock_receiver.session.session_id = "test_session_id"
     mock_receiver.__aiter__.return_value = [AsyncMock()]
-    mock_receiver.__aiter__.return_value[0] = json.dumps({"id": "test_id", "action": "install", "stepId": "test_step_id", "operationId": "test_operation_id"})
+    mock_receiver.__aiter__.return_value[0] = json.dumps({
+        "id": "test_id",
+        "action": "install",
+        "stepId": "test_step_id",
+        "operationId": "test_operation_id",
+        "name": "test_bundle",
+        "version": "1.0.0",
+        "parameters": {},
+    })
 
     mock_service_bus_client_instance.get_queue_receiver.return_value.__aenter__.return_value = mock_receiver
 
@@ -124,10 +132,95 @@ async def test_receive_message(mock_invoke_porter_action, mock_service_bus_clien
     config = {"resource_request_queue": "test_queue"}
 
     await receive_message(mock_service_bus_client_instance, config, keep_running=run_once)
-    mock_receiver.complete_message.assert_called_once()
+    mock_receiver.complete_message.assert_awaited_once()
     mock_service_bus_client_instance.get_queue_receiver.assert_called_once_with(queue_name="test_queue", max_wait_time=1, session_id=ServiceBusSessionFilter.NEXT_AVAILABLE)
 
 
+@pytest.mark.asyncio
+async def test_receive_message_bad_json(mock_service_bus_client, mock_auto_lock_renewer):
+    mock_service_bus_client_instance = mock_service_bus_client.return_value
+
+    # Set up the lock renewer mock correctly
+    mock_renewer = AsyncMock()
+    mock_renewer.register = Mock()
+    mock_auto_lock_renewer.return_value.__aenter__.return_value = mock_renewer
+
+    mock_receiver = AsyncMock()
+    mock_receiver.__aenter__.return_value = mock_receiver
+    mock_receiver.__aexit__.return_value = None
+    mock_receiver.session.session_id = "test_session_id"
+    mock_receiver.__aiter__.return_value = ["invalid_json_string"]
+
+    mock_service_bus_client_instance.get_queue_receiver.return_value.__aenter__.return_value = mock_receiver
+
+    run_once = Mock(side_effect=[True, False])
+
+    config = {"resource_request_queue": "test_queue"}
+
+    await receive_message(mock_service_bus_client_instance, config, keep_running=run_once)
+    mock_receiver.dead_letter_message.assert_awaited_once_with(
+        "invalid_json_string",
+        reason="InvalidJSON",
+        error_description="Expecting value: line 1 column 1 (char 0)"
+    )
+    mock_receiver.complete_message.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "message, error_description",
+    [
+        ("null", "Resource request message must be a JSON object"),
+        ("[]", "Resource request message must be a JSON object"),
+        ("\"text\"", "Resource request message must be a JSON object"),
+        (
+            '{"id": "test_id"}',
+            "Resource request message is missing fields: ['action', 'name', 'operationId', 'parameters', 'stepId', 'version']",
+        ),
+        (
+            '{"id": "test_id", "action": [], "stepId": "test_step_id", "operationId": "test_operation_id", "name": "test_bundle", "version": "1.0.0", "parameters": {}}',
+            "Resource request message has invalid field types: ['action']",
+        ),
+        (
+            '{"id": "test_id", "action": "install", "stepId": "test_step_id", "operationId": "test_operation_id", "name": "test_bundle", "version": "1.0.0", "parameters": null}',
+            "Resource request message has invalid field types: ['parameters']",
+        ),
+        (
+            '{"id": "test_id", "action": "install", "stepId": "test_step_id", "operationId": "test_operation_id", "name": "test_bundle", "version": "1.0.0", "parameters": {}, "user": []}',
+            "Resource request message has invalid field types: ['user']",
+        ),
+    ],
+)
+async def test_receive_message_invalid_json_structure(
+    message, error_description, mock_service_bus_client, mock_auto_lock_renewer
+):
+    mock_service_bus_client_instance = mock_service_bus_client.return_value
+
+    mock_renewer = AsyncMock()
+    mock_renewer.register = Mock()
+    mock_auto_lock_renewer.return_value.__aenter__.return_value = mock_renewer
+
+    mock_receiver = AsyncMock()
+    mock_receiver.__aenter__.return_value = mock_receiver
+    mock_receiver.__aexit__.return_value = None
+    mock_receiver.session.session_id = "test_session_id"
+    mock_receiver.__aiter__.return_value = [message]
+
+    mock_service_bus_client_instance.get_queue_receiver.return_value.__aenter__.return_value = mock_receiver
+
+    run_once = Mock(side_effect=[True, False])
+    config = {"resource_request_queue": "test_queue"}
+
+    await receive_message(mock_service_bus_client_instance, config, keep_running=run_once)
+
+    mock_receiver.dead_letter_message.assert_awaited_once_with(
+        message,
+        reason="InvalidJSON",
+        error_description=error_description,
+    )
+    mock_receiver.complete_message.assert_not_awaited()
+
+
 @pytest.mark.asyncio
 async def test_receive_message_unknown_exception(mock_auto_lock_renewer, mock_service_bus_client, mock_logger):
     """Test receiving a message with an unknown exception."""
@@ -143,7 +236,15 @@ async def test_receive_message_unknown_exception(mock_auto_lock_renewer, mock_se
     mock_receiver.__aexit__.return_value = None
     mock_receiver.session.session_id = "test_session_id"
     mock_receiver.__aiter__.return_value = [AsyncMock()]
-    mock_receiver.__aiter__.return_value[0] = json.dumps({"id": "test_id", "action": "install", "stepId": "test_step_id", "operationId": "test_operation_id"})
+    mock_receiver.__aiter__.return_value[0] = json.dumps({
+        "id": "test_id",
+        "action": "install",
+        "stepId": "test_step_id",
+        "operationId": "test_operation_id",
+        "name": "test_bundle",
+        "version": "1.0.0",
+        "parameters": {},
+    })
 
     mock_service_bus_client_instance.get_queue_receiver.return_value.__aenter__.return_value = mock_receiver
 

resource_processor/vmss_porter/runner.py

@@ -18,6 +18,43 @@
 from azure.identity.aio import DefaultAzureCredential
 
 
+RESOURCE_REQUEST_FIELD_TYPES = {
+    "id": str,
+    "action": str,
+    "stepId": str,
+    "operationId": str,
+    "name": str,
+    "version": str,
+    "parameters": dict,
+}
+OPTIONAL_RESOURCE_REQUEST_FIELD_TYPES = {"user": dict}
+
+
+def validate_resource_request(message: object) -> None:
+    if not isinstance(message, dict):
+        raise ValueError("Resource request message must be a JSON object")
+
+    missing_fields = set(RESOURCE_REQUEST_FIELD_TYPES) - message.keys()
+    if missing_fields:
+        raise ValueError(f"Resource request message is missing fields: {sorted(missing_fields)}")
+
+    invalid_fields = [
+        field_name
+        for field_name, field_type in RESOURCE_REQUEST_FIELD_TYPES.items()
+        if not isinstance(message[field_name], field_type)
+    ]
+    if invalid_fields:
+        raise ValueError(f"Resource request message has invalid field types: {sorted(invalid_fields)}")
+
+    invalid_optional_fields = [
+        field_name
+        for field_name, field_type in OPTIONAL_RESOURCE_REQUEST_FIELD_TYPES.items()
+        if field_name in message and not isinstance(message[field_name], field_type)
+    ]
+    if invalid_optional_fields:
+        raise ValueError(f"Resource request message has invalid field types: {sorted(invalid_optional_fields)}")
+
+
 def set_up_config() -> Optional[dict]:
     try:
         config = get_config()
@@ -71,8 +108,11 @@ async def receive_message(service_bus_client, config: dict, keep_running=lambda:
 
                         try:
                             message = json.loads(str(msg))
-                        except (json.JSONDecodeError) as e:
+                            validate_resource_request(message)
+                        except (json.JSONDecodeError, ValueError) as e:
                             logger.error(f"Received bad service bus resource request message: {e}")
+                            await receiver.dead_letter_message(msg, reason="InvalidJSON", error_description=str(e))
+                            continue
 
                         with tracer.start_as_current_span("receive_message") as current_span:
                             current_span.set_attribute("resource_id", message["id"])