From bc771fbc1e3e8cde07ce2064e691db0ae2130e02 Mon Sep 17 00:00:00 2001
From: jagadish-zentiti <jagadish.l@zentiti.com>
Date: Sat, 11 Jul 2026 09:56:23 +0530
Subject: [PATCH] fix(mcp): guard DbTokenStorage against non-dict oauth_tokens
 JSON (#5107)

_load() returned whatever json.loads() produced without checking it was a
dict; _update() did the same before assigning data[key] = value. If the
oauth_tokens column ever held a JSON array or primitive (DB corruption,
manual edit, migration drift), _load()'s callers crashed with
AttributeError on .get(), and _update() crashed with TypeError trying to
item-assign into a list/string/int.

Validate the parsed value is a dict in both methods, falling back to {}
otherwise - same recovery behavior already used elsewhere in the codebase
for this exact JSON-blob-is-not-a-dict shape (_parse_tool_args,
_is_sensitive_path's siblings).

Adds 3 regression tests for _load, get_tokens, and _update against a
non-dict oauth_tokens value.

Fixes #5082
---
 src/mcp_oauth.py        |  6 ++++-
 tests/test_mcp_oauth.py | 51 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/src/mcp_oauth.py b/src/mcp_oauth.py
index 9f3b2ad4d..27a30383e 100644
--- a/src/mcp_oauth.py
+++ b/src/mcp_oauth.py
@@ -96,7 +96,9 @@ class DbTokenStorage:
         try:
             srv = db.query(McpServer).filter(McpServer.id == self.server_id).first()
             if srv and srv.oauth_tokens:
-                return json.loads(srv.oauth_tokens)
+                parsed = json.loads(srv.oauth_tokens)
+                if isinstance(parsed, dict):
+                    return parsed
         finally:
             db.close()
         return {}
@@ -111,6 +113,8 @@ class DbTokenStorage:
             if srv is None:
                 return
             data = json.loads(srv.oauth_tokens) if srv.oauth_tokens else {}
+            if not isinstance(data, dict):
+                data = {}
             data[key] = value
             srv.oauth_tokens = json.dumps(data)
             db.commit()
diff --git a/tests/test_mcp_oauth.py b/tests/test_mcp_oauth.py
index a9f5fdf6b..6fb6f43b9 100644
--- a/tests/test_mcp_oauth.py
+++ b/tests/test_mcp_oauth.py
@@ -1,4 +1,5 @@
 import asyncio
+import json
 from src import mcp_oauth
 
 
@@ -79,3 +80,53 @@ def test_db_token_storage_round_trip():
     t = asyncio.run(go())
     assert t.access_token == "abc"
     assert srv.oauth_tokens is not None  # persisted as JSON
+
+
+def _fake_storage(oauth_tokens):
+    class FakeSrv:
+        pass
+
+    srv = FakeSrv()
+    srv.oauth_tokens = oauth_tokens
+
+    class FakeQuery:
+        def filter(self, *a):
+            return self
+
+        def first(self):
+            return srv
+
+    class FakeSession:
+        def query(self, *a):
+            return FakeQuery()
+
+        def commit(self):
+            pass
+
+        def close(self):
+            pass
+
+    return srv, mcp_oauth.DbTokenStorage("srv-1", session_factory=lambda: FakeSession())
+
+
+def test_load_falls_back_to_empty_dict_for_non_dict_json():
+    # A corrupted/migrated oauth_tokens column holding a JSON array, not an
+    # object, must not crash _load()'s callers with AttributeError.
+    _srv, storage = _fake_storage('["stale", "data"]')
+    assert storage._load() == {}
+
+
+def test_get_tokens_returns_none_for_non_dict_oauth_tokens():
+    _srv, storage = _fake_storage("42")
+
+    async def go():
+        return await storage.get_tokens()
+
+    assert asyncio.run(go()) is None
+
+
+def test_update_recovers_from_non_dict_oauth_tokens():
+    # _update() must not raise TypeError trying to item-assign into a list.
+    srv, storage = _fake_storage('["stale", "data"]')
+    storage._update("tokens", {"access_token": "new"})
+    assert json.loads(srv.oauth_tokens) == {"tokens": {"access_token": "new"}}