mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
fix(mcp): address Greptile P2s on credential encoding helpers
Three minor fixes from Greptile review: 1. _decode_user_credential now also catches TypeError so a null credential_b64 value returns None instead of propagating, matching the documented "returns None when neither path yields a valid string" contract. 2. The OAuth2 BYOK guard error no longer claims the existing row is a BYOK credential — after a salt-key rotation, an OAuth2 row can fail to decrypt and reach the same guard. Reword to "could not be verified as an OAuth2 token", which is accurate for both cases. 3. Drop the no-op sys.path.insert in the new test file (other tests in the directory don't need it; pytest picks up the package via the installed editable wheel). Adds a regression test for the None-input case.
This commit is contained in:
parent
f3000bda36
commit
c76c300392
2 changed files with 16 additions and 10 deletions
|
|
@ -517,7 +517,7 @@ def _decode_user_credential(stored: str) -> Optional[str]:
|
|||
return decrypted
|
||||
try:
|
||||
return base64.urlsafe_b64decode(stored).decode()
|
||||
except (binascii.Error, UnicodeDecodeError, ValueError):
|
||||
except (binascii.Error, UnicodeDecodeError, ValueError, TypeError):
|
||||
return None
|
||||
|
||||
|
||||
|
|
@ -649,9 +649,14 @@ async def store_user_oauth_credential(
|
|||
existing is not None
|
||||
and _decode_oauth_payload(existing.credential_b64) is None
|
||||
):
|
||||
# Existing row is either a BYOK secret or an OAuth2 row that no
|
||||
# longer decrypts (e.g. after a salt-key rotation). In either
|
||||
# case, refuse to overwrite — the caller would clobber data
|
||||
# that may still be recoverable.
|
||||
raise ValueError(
|
||||
f"A non-OAuth2 credential already exists for user {user_id} "
|
||||
f"and server {server_id}. Refusing to overwrite."
|
||||
f"Existing credential for user {user_id} and server "
|
||||
f"{server_id} could not be verified as an OAuth2 token. "
|
||||
f"Refusing to overwrite."
|
||||
)
|
||||
|
||||
encoded = encrypt_value_helper(json.dumps(payload))
|
||||
|
|
|
|||
|
|
@ -10,15 +10,11 @@ keeps a plain-base64 fallback on read so existing rows continue to work.
|
|||
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../../../.."))
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.db import ( # noqa: E402
|
||||
from litellm.proxy._experimental.mcp_server.db import (
|
||||
_decode_user_credential,
|
||||
get_user_credential,
|
||||
get_user_oauth_credential,
|
||||
|
|
@ -201,7 +197,7 @@ async def test_oauth_get_returns_none_for_byok_row():
|
|||
@pytest.mark.asyncio
|
||||
async def test_byok_guard_rejects_overwriting_legacy_byok():
|
||||
prisma = _make_prisma_with_existing(row=_legacy_row("plain-byok-key"))
|
||||
with pytest.raises(ValueError, match="non-OAuth2 credential"):
|
||||
with pytest.raises(ValueError, match="could not be verified as an OAuth2"):
|
||||
await store_user_oauth_credential(prisma, "alice", "srv-1", "tok")
|
||||
|
||||
|
||||
|
|
@ -218,7 +214,7 @@ async def test_byok_guard_rejects_overwriting_encrypted_byok():
|
|||
return_value=encrypted_row
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="non-OAuth2 credential"):
|
||||
with pytest.raises(ValueError, match="could not be verified as an OAuth2"):
|
||||
await store_user_oauth_credential(prisma, "alice", "srv-1", "tok")
|
||||
|
||||
|
||||
|
|
@ -287,6 +283,11 @@ def test_decode_user_credential_handles_garbage():
|
|||
assert _decode_user_credential("not-base64-and-not-encrypted!!!") is None
|
||||
|
||||
|
||||
def test_decode_user_credential_handles_none():
|
||||
# Defensive: a null DB value must return None, not propagate TypeError.
|
||||
assert _decode_user_credential(None) is None
|
||||
|
||||
|
||||
def test_decode_user_credential_legacy_path():
|
||||
plain = "legacy-secret"
|
||||
stored = base64.urlsafe_b64encode(plain.encode()).decode()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue