swtich to using hmac message signing

appsettings.json

@@ -1,4 +1,5 @@
 {
+  "GradLlmMessageHmacKey": ".env",
   "Logging": {
     "LogLevel": {
       "Default": "Debug",

gradllm_hmac.py

@@ -0,0 +1,60 @@
+import base64
+import hashlib
+import hmac
+import json
+from typing import Any, Dict
+
+from config import settings
+
+PROTOCOL = "gradllm-hmac-sha256-v1"
+
+
+def _key() -> bytes:
+    value = getattr(settings, "GradLlmMessageHmacKey", None)
+    if not isinstance(value, str) or len(value.encode("utf-8")) < 32:
+        raise RuntimeError("GradLlmMessageHmacKey must contain at least 32 UTF-8 bytes")
+    return value.encode("utf-8")
+
+
+def _signature(direction: str, exchange: str, routing_key: str, envelope: Dict[str, Any]) -> str:
+    fields = (
+        PROTOCOL,
+        direction,
+        exchange,
+        routing_key,
+        envelope["reply_key"],
+        str(envelope["sequence"]),
+        "1" if envelope["final"] else "0",
+        envelope["payload_base64"],
+    )
+    digest = hmac.new(_key(), "\n".join(fields).encode("utf-8"), hashlib.sha256).digest()
+    return base64.b64encode(digest).decode("ascii")
+
+
+def verify_request(envelope: Dict[str, Any], exchange: str, routing_key: str) -> Dict[str, Any]:
+    if not isinstance(envelope, dict) or envelope.get("protocol") != PROTOCOL:
+        raise ValueError("missing GradLLM HMAC envelope")
+    if envelope.get("sequence") != 0 or envelope.get("final") is not False:
+        raise ValueError("invalid GradLLM request envelope")
+    expected = _signature("request", exchange, routing_key, envelope)
+    if not hmac.compare_digest(str(envelope.get("signature", "")), expected):
+        raise ValueError("invalid GradLLM request HMAC")
+    payload = base64.b64decode(envelope["payload_base64"], validate=True)
+    decoded = json.loads(payload.decode("utf-8"))
+    if not isinstance(decoded, dict):
+        raise ValueError("GradLLM request payload must be an object")
+    decoded["reply_key"] = envelope["reply_key"]
+    return decoded
+
+
+def create_response(exchange: str, reply_key: str, sequence: int, payload: Any, final: bool = False) -> Dict[str, Any]:
+    payload_bytes = json.dumps(payload, separators=(",", ":"), ensure_ascii=False).encode("utf-8")
+    envelope = {
+        "protocol": PROTOCOL,
+        "reply_key": reply_key,
+        "sequence": sequence,
+        "final": final,
+        "payload_base64": base64.b64encode(payload_bytes).decode("ascii"),
+    }
+    envelope["signature"] = _signature("response", exchange, reply_key, envelope)
+    return envelope

oa_server.py

@@ -5,6 +5,8 @@ from typing import Any, Dict, List, AsyncIterable, Optional
 from backends_base import ChatBackend, ImagesBackend
 
 from rabbit_repo import RabbitRepo
+from gradllm_hmac import create_response, verify_request
+from config import settings
 
 logger = logging.getLogger(__name__)
 
@@ -72,20 +74,28 @@ class OpenAIServers:
             logger.warning("oaChatCreate: data is not a dict")
             return
 
+        try:
+            data = verify_request(data, "oa.chat.create", settings.RABBIT_ROUTING_KEY)
+        except Exception:
+            logger.warning("oaChatCreate: rejected unauthenticated request", exc_info=True)
+            return
+
         reply_key = data.get("reply_key")
         if not reply_key:
             logger.error("oaChatCreate: missing reply_key")
             return
 
         user_id = _extract_user_id(data)
+        sequence = 0
         try:
             async for chunk in self._chat.stream(data):
                 try:
                     await self._pub.publish(
                         "oa.chat.reply",
-                        _with_user_id(chunk, user_id),
+                        create_response("oa.chat.reply", reply_key, sequence, _with_user_id(chunk, user_id)),
                         routing_key=reply_key,
                     )
+                    sequence += 1
                 except Exception:
                     logger.exception("oaChatCreate: publish failed")
                     break  # stop streaming on publish failure
@@ -94,7 +104,7 @@ class OpenAIServers:
             try:
                 await self._pub.publish(
                     "oa.chat.reply",
-                    _with_user_id({"object": "stream.end"}, user_id),
+                    create_response("oa.chat.reply", reply_key, sequence, {}, final=True),
                     routing_key=reply_key,
                 )
             except Exception:
@@ -110,6 +120,11 @@ class OpenAIServers:
         if not isinstance(data, dict):
             logger.warning("oaImagesGenerate: data is not a dict")
             return
+        try:
+            data = verify_request(data, "oa.images.generate", settings.RABBIT_ROUTING_KEY)
+        except Exception:
+            logger.warning("oaImagesGenerate: rejected unauthenticated request", exc_info=True)
+            return
         reply_key = data.get("reply_key")
         if not reply_key:
             logger.error("oaImagesGenerate: missing reply_key")
@@ -122,7 +137,7 @@ class OpenAIServers:
             try:
                 await self._pub.publish(
                     "oa.images.reply",
-                    _with_user_id(resp, user_id),
+                    create_response("oa.images.reply", reply_key, 0, _with_user_id(resp, user_id)),
                     routing_key=reply_key,
                 )
             except Exception:

test_gradllm_hmac.py

@@ -0,0 +1,41 @@
+import base64
+import importlib
+import json
+import os
+import unittest
+
+
+class GradLlmHmacTests(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        os.environ["GradLlmMessageHmacKey"] = "test-gradllm-hmac-key-that-is-at-least-32-bytes"
+        cls.protocol = importlib.import_module("gradllm_hmac")
+
+    def test_response_round_trip_and_tamper_rejection(self):
+        envelope = self.protocol.create_response("oa.chat.reply", "rk.test", 3, {"value": 42})
+        expected = self.protocol._signature("response", "oa.chat.reply", "rk.test", envelope)
+        self.assertEqual(envelope["signature"], expected)
+
+        envelope["payload_base64"] = base64.b64encode(json.dumps({"value": 43}).encode()).decode()
+        self.assertNotEqual(envelope["signature"], self.protocol._signature("response", "oa.chat.reply", "rk.test", envelope))
+
+    def test_request_verification_returns_original_payload_and_bound_reply_key(self):
+        payload = base64.b64encode(b'{"model":"timesfm"}').decode()
+        envelope = {
+            "protocol": self.protocol.PROTOCOL,
+            "reply_key": "rk.test",
+            "sequence": 0,
+            "final": False,
+            "payload_base64": payload,
+        }
+        envelope["signature"] = self.protocol._signature("request", "oa.chat.create", "execute.api", envelope)
+        decoded = self.protocol.verify_request(envelope, "oa.chat.create", "execute.api")
+        self.assertEqual(decoded, {"model": "timesfm", "reply_key": "rk.test"})
+
+        envelope["reply_key"] = "rk.attacker"
+        with self.assertRaises(ValueError):
+            self.protocol.verify_request(envelope, "oa.chat.create", "execute.api")
+
+
+if __name__ == "__main__":
+    unittest.main()