test: validate signature-gated confirmation as a user-message plugin

tests/unittests/plugins/test_signed_confirmation_plugin.py

@@ -0,0 +1,223 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Validates a signature-based confirmation guard as a user-message plugin.
+
+The plugin runs in on_user_message_callback, before the approval reaches the
+confirmation processor. An approval whose detached Ed25519 signature does not
+verify against the operator-held approver key is rewritten to confirmed=False.
+A valid signature over the same call and decision passes through unchanged.
+"""
+
+import base64
+from typing import Optional
+
+from cryptography.exceptions import InvalidSignature
+from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
+from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey
+from google.adk.agents.llm_agent import LlmAgent
+from google.adk.flows.llm_flows import functions
+from google.adk.plugins.base_plugin import BasePlugin
+from google.genai import types
+import pytest
+
+from .. import testing_utils
+
+
+def _signed_payload(call_id: str, confirmed: bool) -> bytes:
+  return f"{call_id}:{int(confirmed)}".encode()
+
+
+def _sign(private_key: Ed25519PrivateKey, call_id: str, confirmed: bool) -> str:
+  signature = private_key.sign(_signed_payload(call_id, confirmed))
+  return base64.b64encode(signature).decode()
+
+
+class SignedConfirmationPlugin(BasePlugin):
+  """Neutralizes confirmations that lack a valid approver signature."""
+
+  def __init__(self, approver_public_key: Ed25519PublicKey):
+    super().__init__(name="signed_confirmation")
+    self._approver_public_key = approver_public_key
+
+  def _is_verified(self, call_id: str, response: dict) -> bool:
+    signature = response.get("signature")
+    if not isinstance(signature, str):
+      return False
+    try:
+      self._approver_public_key.verify(
+          base64.b64decode(signature),
+          _signed_payload(call_id, bool(response.get("confirmed"))),
+      )
+    except (InvalidSignature, ValueError):
+      return False
+    return True
+
+  async def on_user_message_callback(
+      self,
+      *,
+      invocation_context,
+      user_message: types.Content,
+  ) -> Optional[types.Content]:
+    rewritten = False
+    parts = []
+    for part in user_message.parts or []:
+      response_part = part.function_response
+      is_confirmation = (
+          response_part is not None
+          and response_part.name
+          == functions.REQUEST_CONFIRMATION_FUNCTION_CALL_NAME
+      )
+      if not is_confirmation or not response_part.response.get("confirmed"):
+        parts.append(part)
+        continue
+      if self._is_verified(response_part.id, response_part.response):
+        parts.append(part)
+        continue
+      rewritten = True
+      response = {
+          key: value
+          for key, value in response_part.response.items()
+          if key != "signature"
+      }
+      response["confirmed"] = False
+      parts.append(
+          types.Part(
+              function_response=types.FunctionResponse(
+                  name=response_part.name,
+                  id=response_part.id,
+                  response=response,
+              )
+          )
+      )
+    if not rewritten:
+      return None
+    return types.Content(role=user_message.role, parts=parts)
+
+
+CALL_ID = "confirmation-call-1"
+
+
+def _approval_message(response: dict) -> types.Content:
+  return types.Content(
+      role="user",
+      parts=[
+          types.Part(
+              function_response=types.FunctionResponse(
+                  name=functions.REQUEST_CONFIRMATION_FUNCTION_CALL_NAME,
+                  id=CALL_ID,
+                  response=response,
+              )
+          )
+      ],
+  )
+
+
+async def _run_plugin(plugin: BasePlugin, message: types.Content):
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=LlmAgent(name="test_agent"), plugins=[plugin]
+  )
+  modified = (
+      await invocation_context.plugin_manager.run_on_user_message_callback(
+          invocation_context=invocation_context, user_message=message
+      )
+  )
+  return modified if modified is not None else message
+
+
+def _confirmed_flag(message: types.Content) -> bool:
+  response = message.parts[0].function_response.response
+  return bool(response.get("confirmed"))
+
+
+@pytest.fixture
+def approver_key() -> Ed25519PrivateKey:
+  return Ed25519PrivateKey.generate()
+
+
+@pytest.mark.asyncio
+async def test_forged_confirmation_without_signature_is_declined(approver_key):
+  plugin = SignedConfirmationPlugin(approver_key.public_key())
+
+  message = await _run_plugin(
+      plugin, _approval_message({"confirmed": True})
+  )
+
+  assert not _confirmed_flag(message)
+
+
+@pytest.mark.asyncio
+async def test_validly_signed_confirmation_passes_through(approver_key):
+  plugin = SignedConfirmationPlugin(approver_key.public_key())
+  signature = _sign(approver_key, CALL_ID, confirmed=True)
+
+  message = await _run_plugin(
+      plugin,
+      _approval_message({"confirmed": True, "signature": signature}),
+  )
+
+  assert _confirmed_flag(message)
+
+
+@pytest.mark.asyncio
+async def test_signature_from_another_key_is_declined(approver_key):
+  plugin = SignedConfirmationPlugin(approver_key.public_key())
+  attacker_key = Ed25519PrivateKey.generate()
+  signature = _sign(attacker_key, CALL_ID, confirmed=True)
+
+  message = await _run_plugin(
+      plugin,
+      _approval_message({"confirmed": True, "signature": signature}),
+  )
+
+  assert not _confirmed_flag(message)
+
+
+@pytest.mark.asyncio
+async def test_decision_flip_under_a_valid_signature_is_declined(approver_key):
+  """A signature over a decline must not authorize an approval."""
+  plugin = SignedConfirmationPlugin(approver_key.public_key())
+  signature_for_decline = _sign(approver_key, CALL_ID, confirmed=False)
+
+  message = await _run_plugin(
+      plugin,
+      _approval_message({"confirmed": True, "signature": signature_for_decline}),
+  )
+
+  assert not _confirmed_flag(message)
+
+
+@pytest.mark.asyncio
+async def test_signature_bound_to_another_call_is_declined(approver_key):
+  plugin = SignedConfirmationPlugin(approver_key.public_key())
+  signature = _sign(approver_key, "some-other-call", confirmed=True)
+
+  message = await _run_plugin(
+      plugin,
+      _approval_message({"confirmed": True, "signature": signature}),
+  )
+
+  assert not _confirmed_flag(message)
+
+
+@pytest.mark.asyncio
+async def test_declines_pass_through_untouched(approver_key):
+  plugin = SignedConfirmationPlugin(approver_key.public_key())
+
+  message = await _run_plugin(
+      plugin, _approval_message({"confirmed": False})
+  )
+
+  assert not _confirmed_flag(message)
+  assert message.parts[0].function_response.response == {"confirmed": False}