mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
fix(vertex_ai): key known Gemini caches by API key without hashing it
The known-cache scope hashed the Gemini API key with SHA-256, which CodeQL reports as weak hashing of sensitive data. The scope is only an in-memory dict key, never logged or persisted, so it now holds the key as-is and no hash is involved.
This commit is contained in:
parent
755efd48a2
commit
6b2fae2dfe
1 changed files with 12 additions and 10 deletions
|
|
@ -1,5 +1,4 @@
|
|||
import asyncio
|
||||
import hashlib
|
||||
import re
|
||||
import threading
|
||||
import time
|
||||
|
|
@ -107,10 +106,10 @@ class _KnownCachedContents:
|
|||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._entries: dict[tuple[str, str], tuple[str, float]] = {} # mutable-ok: private, guarded by _lock
|
||||
self._entries: dict[tuple[tuple[str, str], str], tuple[str, float]] = {} # mutable-ok: guarded by _lock
|
||||
self._lock: Final = threading.Lock()
|
||||
self._inflight: weakref.WeakKeyDictionary[ # mutable-ok: private, guarded by _lock
|
||||
asyncio.AbstractEventLoop, dict[tuple[str, str], asyncio.Future[str | None]]
|
||||
asyncio.AbstractEventLoop, dict[tuple[tuple[str, str], str], asyncio.Future[str | None]]
|
||||
] = weakref.WeakKeyDictionary()
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -118,18 +117,21 @@ class _KnownCachedContents:
|
|||
base_url: str,
|
||||
custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"],
|
||||
api_key: str | None,
|
||||
) -> str:
|
||||
"""Where a cache name is valid: the Vertex URL names project and location, a Gemini key names the project."""
|
||||
) -> tuple[str, str]:
|
||||
"""Where a cache name is valid: the Vertex URL names project and location, a Gemini key names the project.
|
||||
|
||||
The Gemini key is kept as-is in this in-memory scope and is never logged or persisted.
|
||||
"""
|
||||
if custom_llm_provider == "gemini":
|
||||
return f"{base_url}#{hashlib.sha256((api_key or '').encode()).hexdigest()}"
|
||||
return base_url
|
||||
return (base_url, api_key or "")
|
||||
return (base_url, "")
|
||||
|
||||
@staticmethod
|
||||
def is_usable(expire_time: object) -> bool:
|
||||
expires_at: Final = _parse_expire_time(expire_time)
|
||||
return expires_at is None or expires_at - KNOWN_CACHE_EXPIRY_MARGIN_SECONDS > time.time()
|
||||
|
||||
def get(self, scope: str, cache_key: str) -> str | None:
|
||||
def get(self, scope: tuple[str, str], cache_key: str) -> str | None:
|
||||
with self._lock:
|
||||
entry: Final = self._entries.get((scope, cache_key))
|
||||
if entry is None:
|
||||
|
|
@ -140,7 +142,7 @@ class _KnownCachedContents:
|
|||
del self._entries[(scope, cache_key)]
|
||||
return None
|
||||
|
||||
def remember(self, scope: str, cache_key: str, name: str | None, expire_time: object) -> None:
|
||||
def remember(self, scope: tuple[str, str], cache_key: str, name: str | None, expire_time: object) -> None:
|
||||
expires_at: Final = _parse_expire_time(expire_time)
|
||||
if not name or expires_at is None:
|
||||
return
|
||||
|
|
@ -156,7 +158,7 @@ class _KnownCachedContents:
|
|||
self._entries[(scope, cache_key)] = (name, reuse_until)
|
||||
|
||||
async def single_flight(
|
||||
self, scope: str, cache_key: str, resolve: Callable[[], Awaitable[str | None]]
|
||||
self, scope: tuple[str, str], cache_key: str, resolve: Callable[[], Awaitable[str | None]]
|
||||
) -> str | None:
|
||||
"""Run ``resolve`` once per (scope, cache_key) at a time within this event loop.
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue