feat(rust): add virtual key storage contracts (#43468)

Co-authored-by: Yujong Lee <yujong@berri.ai>

litellm-rust/Cargo.lock

@@ -3537,6 +3537,19 @@ dependencies = [
  "wiremock",
 ]
 
+[[package]]
+name = "litellm-gateway-management"
+version = "0.1.0"
+dependencies = [
+ "base64 0.22.1",
+ "litellm-auth-types",
+ "litellm-gateway-auth",
+ "rand 0.8.7",
+ "rstest",
+ "thiserror 2.0.19",
+ "tokio",
+]
+
 [[package]]
 name = "litellm-host"
 version = "0.1.0"

litellm-rust/Cargo.toml

@@ -16,6 +16,7 @@ litellm-core = { path = "crates/core" }
 litellm-gateway = { path = "crates/gateway" }
 litellm-gateway-inference = { path = "crates/gateway-inference" }
 litellm-gateway-auth = { path = "crates/gateway-auth" }
+litellm-gateway-management = { path = "crates/gateway-management" }
 litellm-coroutine = { path = "crates/coroutine" }
 litellm-host = { path = "crates/host" }
 litellm-host-http = { path = "crates/host-http" }

litellm-rust/crates/gateway-auth/src/error.rs

@@ -3,6 +3,16 @@ use axum::{
     response::{IntoResponse, Response},
 };
 
+#[derive(Debug, thiserror::Error)]
+pub enum KeyError {
+    #[error("missing or invalid virtual key")]
+    Invalid,
+    #[error("virtual key has expired")]
+    Expired,
+    #[error("virtual key lookup unavailable")]
+    Lookup(#[source] Box<dyn std::error::Error + Send + Sync>),
+}
+
 #[derive(Debug, thiserror::Error)]
 pub enum Error {
     #[error("operation is not permitted")]

litellm-rust/crates/gateway-auth/src/keys.rs

@@ -0,0 +1,48 @@
+use std::{future::Future, pin::Pin, time::SystemTime};
+
+use litellm_auth_types::SecretValue;
+
+use crate::{KeyError, hash_token};
+
+#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
+pub struct KeyHash(String);
+
+impl KeyHash {
+    pub fn from_token(token: &SecretValue) -> Self {
+        Self(hash_token(token.expose()))
+    }
+
+    pub fn as_str(&self) -> &str {
+        &self.0
+    }
+}
+
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub enum KeyStatus {
+    Active { expires_at: Option<SystemTime> },
+    Revoked,
+}
+
+pub type LookupFuture<'a> =
+    Pin<Box<dyn Future<Output = Result<Option<KeyStatus>, KeyError>> + Send + 'a>>;
+
+pub trait KeyLookup: Send + Sync {
+    fn lookup<'a>(&'a self, hash: &'a KeyHash) -> LookupFuture<'a>;
+}
+
+pub async fn verify_key(
+    lookup: &dyn KeyLookup,
+    token: &SecretValue,
+    now: impl FnOnce() -> SystemTime,
+) -> Result<KeyHash, KeyError> {
+    let hash = KeyHash::from_token(token);
+    match lookup.lookup(&hash).await? {
+        Some(KeyStatus::Active { expires_at }) => {
+            if expires_at.is_some_and(|expiry| expiry <= now()) {
+                return Err(KeyError::Expired);
+            }
+            Ok(hash)
+        }
+        Some(KeyStatus::Revoked) | None => Err(KeyError::Invalid),
+    }
+}

litellm-rust/crates/gateway-auth/src/lib.rs

@@ -3,6 +3,7 @@ mod authorization;
 mod error;
 mod http;
 mod identity;
+pub mod keys;
 
 use sha2::{Digest, Sha256};
 
@@ -14,7 +15,7 @@ pub use authorization::{
     AccessRequest, AuthenticatedRequest, AuthorizedOperation, Authorizer, McpAction,
     NoAdditionalPolicy, Permissions, UiAction,
 };
-pub use error::Error;
+pub use error::{Error, KeyError};
 pub use http::{Bearer, CredentialExtractor, RequireMasterKey, authenticate};
 pub use identity::{
     AuthenticatedCaller, Authentication, AuthenticationMethod, Principal, PrincipalKind,

litellm-rust/crates/gateway-auth/tests/keys.rs

@@ -0,0 +1,123 @@
+use std::time::{Duration, SystemTime};
+
+use litellm_auth_types::SecretValue;
+use litellm_gateway_auth::{
+    KeyError,
+    keys::{KeyHash, KeyLookup, KeyStatus, LookupFuture, verify_key},
+};
+use rstest::{fixture, rstest};
+
+struct Lookup {
+    hash: KeyHash,
+    status: Option<KeyStatus>,
+}
+
+impl KeyLookup for Lookup {
+    fn lookup<'a>(&'a self, hash: &'a KeyHash) -> LookupFuture<'a> {
+        Box::pin(async move {
+            Ok(if hash == &self.hash {
+                self.status.clone()
+            } else {
+                None
+            })
+        })
+    }
+}
+
+#[fixture]
+fn token() -> SecretValue {
+    SecretValue::new("sk-test-credential")
+}
+
+#[fixture]
+fn now() -> SystemTime {
+    SystemTime::UNIX_EPOCH + Duration::from_secs(100)
+}
+
+#[rstest]
+#[case::no_expiration(None, true)]
+#[case::future(Some(101), true)]
+#[case::boundary(Some(100), false)]
+#[case::past(Some(99), false)]
+#[tokio::test]
+async fn checks_expiration(
+    token: SecretValue,
+    now: SystemTime,
+    #[case] expiration: Option<u64>,
+    #[case] valid: bool,
+) {
+    let hash = KeyHash::from_token(&token);
+    let lookup = Lookup {
+        hash: hash.clone(),
+        status: Some(KeyStatus::Active {
+            expires_at: expiration
+                .map(|seconds| SystemTime::UNIX_EPOCH + Duration::from_secs(seconds)),
+        }),
+    };
+
+    let result = verify_key(&lookup, &token, || now).await;
+
+    match valid {
+        true => assert_eq!(result.unwrap(), hash),
+        false => assert!(matches!(result, Err(KeyError::Expired))),
+    }
+}
+
+#[rstest]
+#[case::missing(None)]
+#[case::revoked(Some(KeyStatus::Revoked))]
+#[tokio::test]
+async fn rejects_unusable_keys(
+    token: SecretValue,
+    now: SystemTime,
+    #[case] status: Option<KeyStatus>,
+) {
+    let lookup = Lookup {
+        hash: KeyHash::from_token(&token),
+        status,
+    };
+
+    assert!(matches!(
+        verify_key(&lookup, &token, || now).await,
+        Err(KeyError::Invalid)
+    ));
+}
+
+#[rstest]
+#[tokio::test]
+async fn rejects_a_different_token(token: SecretValue, now: SystemTime) {
+    let lookup = Lookup {
+        hash: KeyHash::from_token(&token),
+        status: Some(KeyStatus::Active { expires_at: None }),
+    };
+
+    assert!(matches!(
+        verify_key(&lookup, &SecretValue::new("sk-other-credential"), || now).await,
+        Err(KeyError::Invalid)
+    ));
+}
+
+struct Unavailable;
+
+impl KeyLookup for Unavailable {
+    fn lookup<'a>(&'a self, _: &'a KeyHash) -> LookupFuture<'a> {
+        Box::pin(async {
+            Err(KeyError::Lookup(Box::new(std::io::Error::other(
+                "backend detail",
+            ))))
+        })
+    }
+}
+
+#[rstest]
+#[tokio::test]
+async fn lookup_failure_preserves_source_without_exposing_it(token: SecretValue, now: SystemTime) {
+    let error = verify_key(&Unavailable, &token, || now).await.unwrap_err();
+
+    assert!(matches!(&error, KeyError::Lookup(_)));
+    assert_eq!(
+        std::error::Error::source(&error).unwrap().to_string(),
+        "backend detail"
+    );
+    assert!(!error.to_string().contains("backend detail"));
+}

litellm-rust/crates/gateway-management/Cargo.toml

@@ -0,0 +1,17 @@
+[package]
+name = "litellm-gateway-management"
+version = "0.1.0"
+edition.workspace = true
+license.workspace = true
+repository.workspace = true
+
+[dependencies]
+base64.workspace = true
+litellm-auth-types.workspace = true
+litellm-gateway-auth.workspace = true
+rand.workspace = true
+thiserror.workspace = true
+
+[dev-dependencies]
+rstest.workspace = true
+tokio.workspace = true

litellm-rust/crates/gateway-management/README.md

@@ -0,0 +1,19 @@
+# Gateway management
+
+This crate scaffolds virtual key creation, retrieval, and revocation with an injected `KeyStore`. `gateway-auth::keys` owns the separate read-only `KeyLookup` interface and credential verification. Neither crate imports a database implementation
+
+The future `db` key repository should implement both interfaces against the existing LiteLLM database. The gateway composition layer can inject the same repository into management and authentication, sharing its connection pool
+
+## Storage contract
+
+`create` inserts the supplied hash and expiration atomically, returning `AlreadyExists` without changing an existing record on collision. Management generates a token from 32 bytes of operating-system randomness and returns its plaintext only after storage succeeds. `KeyHash` uses the existing SHA-256 token hashing function. Storage never receives the plaintext token
+
+`get` returns the stored status, including revoked keys. `revoke` returns `Revoked` for an existing key, including one already revoked, and `NotFound` for a missing key. Once revocation succeeds, every subsequent lookup must return revoked or missing, including lookups from other gateway instances. Database writes and any cache invalidation must preserve this contract
+
+`KeyLookup` returns active status with optional expiration, revoked status, or no record. Backend failures remain errors. `verify_key` checks the current lookup result and expiration on every call, reads its supplied clock after lookup completes, and returns the verified hash. Pass `SystemTime::now` in production or a fixed clock in tests
+
+## Scope
+
+These are backend operations, with no HTTP routes, database adapter, production in-memory store, or gateway CLI wiring. The in-memory repository under `tests` exercises the shared storage contract. Verification establishes only that a credential is valid; it does not grant permissions. Caller identity, management authorization, and inference policy enforcement belong to the gateway auth integration before these operations are exposed over HTTP
+
+The existing master-key authentication path is unchanged. User, team, model, and budget policies are not represented in this scaffold, so existing database rows carrying those policies require that integration before they can authorize traffic

litellm-rust/crates/gateway-management/src/error.rs

@@ -0,0 +1,11 @@
+#[derive(Debug, thiserror::Error)]
+pub enum Error {
+    #[error("key expiration must be in the future")]
+    InvalidExpiration,
+    #[error("key already exists")]
+    AlreadyExists,
+    #[error("key storage unavailable")]
+    Storage(#[source] Box<dyn std::error::Error + Send + Sync>),
+    #[error("key generation unavailable")]
+    Entropy(#[source] rand::Error),
+}

litellm-rust/crates/gateway-management/src/keys.rs

@@ -0,0 +1,85 @@
+use std::{future::Future, pin::Pin, sync::Arc, time::SystemTime};
+
+use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD};
+use litellm_auth_types::SecretValue;
+use litellm_gateway_auth::keys::{KeyHash, KeyStatus};
+use rand::{RngCore, rngs::OsRng};
+
+use crate::Error;
+
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub struct NewKey {
+    pub hash: KeyHash,
+    pub expires_at: Option<SystemTime>,
+}
+
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub struct KeyRecord {
+    pub hash: KeyHash,
+    pub status: KeyStatus,
+}
+
+#[derive(Debug)]
+pub struct GeneratedKey {
+    pub token: SecretValue,
+    pub record: KeyRecord,
+}
+
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub enum Revocation {
+    Revoked,
+    NotFound,
+}
+
+pub type StoreFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
+
+pub trait KeyStore: Send + Sync {
+    fn create(&self, key: NewKey) -> StoreFuture<'_, ()>;
+    fn get<'a>(&'a self, hash: &'a KeyHash) -> StoreFuture<'a, Option<KeyRecord>>;
+    fn revoke<'a>(&'a self, hash: &'a KeyHash) -> StoreFuture<'a, Revocation>;
+}
+
+pub struct Keys {
+    store: Arc<dyn KeyStore>,
+}
+
+impl Keys {
+    pub fn new(store: Arc<dyn KeyStore>) -> Self {
+        Self { store }
+    }
+
+    pub async fn generate(
+        &self,
+        expires_at: Option<SystemTime>,
+        now: SystemTime,
+    ) -> Result<GeneratedKey, Error> {
+        if expires_at.is_some_and(|expiry| expiry <= now) {
+            return Err(Error::InvalidExpiration);
+        }
+        let mut entropy = [0_u8; 32];
+        OsRng.try_fill_bytes(&mut entropy).map_err(Error::Entropy)?;
+        let token = SecretValue::new(format!("sk-{}", URL_SAFE_NO_PAD.encode(entropy)));
+        let hash = KeyHash::from_token(&token);
+        self.store
+            .create(NewKey {
+                hash: hash.clone(),
+                expires_at,
+            })
+            .await?;
+        Ok(GeneratedKey {
+            token,
+            record: KeyRecord {
+                hash,
+                status: KeyStatus::Active { expires_at },
+            },
+        })
+    }
+
+    pub async fn get(&self, hash: &KeyHash) -> Result<Option<KeyRecord>, Error> {
+        self.store.get(hash).await
+    }
+
+    pub async fn revoke(&self, hash: &KeyHash) -> Result<Revocation, Error> {
+        self.store.revoke(hash).await
+    }
+}

litellm-rust/crates/gateway-management/src/lib.rs

@@ -0,0 +1,4 @@
+mod error;
+pub mod keys;
+
+pub use error::Error;

litellm-rust/crates/gateway-management/tests/keys.rs

@@ -0,0 +1,145 @@
+#[path = "keys/support.rs"]
+mod support;
+
+use std::{
+    sync::Arc,
+    time::{Duration, SystemTime},
+};
+
+use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD};
+use litellm_auth_types::SecretValue;
+use litellm_gateway_auth::{
+    KeyError, hash_token,
+    keys::{KeyHash, KeyLookup, KeyStatus, verify_key},
+};
+use litellm_gateway_management::{
+    Error,
+    keys::{Keys, Revocation},
+};
+use rstest::{fixture, rstest};
+use support::{MemoryStore, Unavailable};
+
+#[fixture]
+fn store() -> Arc<MemoryStore> {
+    Arc::new(MemoryStore::default())
+}
+
+#[fixture]
+fn now() -> SystemTime {
+    SystemTime::UNIX_EPOCH + Duration::from_secs(100)
+}
+
+#[rstest]
+#[case::permanent(false)]
+#[case::expiring(true)]
+#[tokio::test]
+async fn generated_key_can_be_retrieved_verified_and_revoked(
+    store: Arc<MemoryStore>,
+    now: SystemTime,
+    #[case] expiring: bool,
+) {
+    let keys = Keys::new(store.clone());
+    let lookup: Arc<dyn KeyLookup> = store;
+    let expires_at = expiring.then_some(now + Duration::from_secs(60));
+
+    let generated = keys.generate(expires_at, now).await.unwrap();
+    let hash = &generated.record.hash;
+
+    assert_eq!(hash.as_str(), hash_token(generated.token.expose()));
+    assert_eq!(generated.record.status, KeyStatus::Active { expires_at });
+    assert_eq!(
+        keys.get(hash).await.unwrap(),
+        Some(generated.record.clone())
+    );
+    assert_eq!(
+        verify_key(lookup.as_ref(), &generated.token, || now)
+            .await
+            .unwrap(),
+        *hash
+    );
+    assert!(!format!("{generated:?}").contains(generated.token.expose()));
+    assert_eq!(
+        URL_SAFE_NO_PAD
+            .decode(generated.token.expose().strip_prefix("sk-").unwrap())
+            .unwrap()
+            .len(),
+        32
+    );
+
+    assert_eq!(keys.revoke(hash).await.unwrap(), Revocation::Revoked);
+    assert_eq!(
+        keys.get(hash).await.unwrap().unwrap().status,
+        KeyStatus::Revoked
+    );
+    assert!(matches!(
+        verify_key(lookup.as_ref(), &generated.token, || now).await,
+        Err(KeyError::Invalid)
+    ));
+    assert_eq!(keys.revoke(hash).await.unwrap(), Revocation::Revoked);
+}
+
+#[rstest]
+#[tokio::test]
+async fn generates_independent_credentials(store: Arc<MemoryStore>, now: SystemTime) {
+    let keys = Keys::new(store.clone());
+    let first = keys.generate(None, now).await.unwrap();
+    let second = keys.generate(None, now).await.unwrap();
+
+    assert_ne!(first.token, second.token);
+    assert_ne!(first.record.hash, second.record.hash);
+    keys.revoke(&first.record.hash).await.unwrap();
+    assert_eq!(
+        verify_key(store.as_ref(), &second.token, || now)
+            .await
+            .unwrap(),
+        second.record.hash
+    );
+}
+
+#[rstest]
+#[case::past(99)]
+#[case::boundary(100)]
+#[tokio::test]
+async fn invalid_expiration_is_rejected_before_storage(now: SystemTime, #[case] seconds: u64) {
+    let keys = Keys::new(Arc::new(Unavailable));
+
+    let result = keys
+        .generate(
+            Some(SystemTime::UNIX_EPOCH + Duration::from_secs(seconds)),
+            now,
+        )
+        .await;
+
+    assert!(matches!(result, Err(Error::InvalidExpiration)));
+}
+
+#[rstest]
+#[tokio::test]
+async fn missing_key_is_not_created_by_get_or_revoke(store: Arc<MemoryStore>) {
+    let keys = Keys::new(store);
+    let hash = KeyHash::from_token(&SecretValue::new("sk-missing"));
+
+    assert_eq!(keys.get(&hash).await.unwrap(), None);
+    assert_eq!(keys.revoke(&hash).await.unwrap(), Revocation::NotFound);
+    assert_eq!(keys.get(&hash).await.unwrap(), None);
+}
+
+#[rstest]
+#[tokio::test]
+async fn storage_failures_are_returned_for_all_operations(now: SystemTime) {
+    let keys = Keys::new(Arc::new(Unavailable));
+    let hash = KeyHash::from_token(&SecretValue::new("sk-unavailable"));
+
+    assert!(matches!(
+        keys.generate(None, now).await,
+        Err(Error::Storage(_))
+    ));
+    assert!(matches!(keys.get(&hash).await, Err(Error::Storage(_))));
+    let error = keys.revoke(&hash).await.unwrap_err();
+    assert!(matches!(&error, Error::Storage(_)));
+    assert_eq!(
+        std::error::Error::source(&error).unwrap().to_string(),
+        "backend detail"
+    );
+    assert!(!error.to_string().contains("backend detail"));
+}

litellm-rust/crates/gateway-management/tests/keys/support.rs

@@ -0,0 +1,84 @@
+use std::{collections::BTreeMap, sync::Mutex};
+
+use litellm_gateway_auth::keys::KeyStatus;
+use litellm_gateway_auth::{
+    KeyError,
+    keys::{KeyHash, KeyLookup, LookupFuture},
+};
+use litellm_gateway_management::{
+    Error,
+    keys::{KeyRecord, KeyStore, NewKey, Revocation, StoreFuture},
+};
+
+#[derive(Default)]
+pub struct MemoryStore(Mutex<BTreeMap<KeyHash, KeyRecord>>);
+
+impl KeyStore for MemoryStore {
+    fn create(&self, key: NewKey) -> StoreFuture<'_, ()> {
+        Box::pin(async move {
+            let mut records = self.0.lock().unwrap();
+            match records.entry(key.hash.clone()) {
+                std::collections::btree_map::Entry::Occupied(_) => Err(Error::AlreadyExists),
+                std::collections::btree_map::Entry::Vacant(entry) => {
+                    entry.insert(KeyRecord {
+                        hash: key.hash,
+                        status: KeyStatus::Active {
+                            expires_at: key.expires_at,
+                        },
+                    });
+                    Ok(())
+                }
+            }
+        })
+    }
+
+    fn get<'a>(&'a self, hash: &'a KeyHash) -> StoreFuture<'a, Option<KeyRecord>> {
+        Box::pin(async move { Ok(self.0.lock().unwrap().get(hash).cloned()) })
+    }
+
+    fn revoke<'a>(&'a self, hash: &'a KeyHash) -> StoreFuture<'a, Revocation> {
+        Box::pin(async move {
+            let mut records = self.0.lock().unwrap();
+            match records.get_mut(hash) {
+                Some(record) => {
+                    record.status = KeyStatus::Revoked;
+                    Ok(Revocation::Revoked)
+                }
+                None => Ok(Revocation::NotFound),
+            }
+        })
+    }
+}
+
+impl KeyLookup for MemoryStore {
+    fn lookup<'a>(&'a self, hash: &'a KeyHash) -> LookupFuture<'a> {
+        Box::pin(async move {
+            self.get(hash)
+                .await
+                .map(|record| record.map(|record| record.status))
+                .map_err(|error| KeyError::Lookup(Box::new(error)))
+        })
+    }
+}
+
+pub struct Unavailable;
+
+fn unavailable<T>() -> Result<T, Error> {
+    Err(Error::Storage(Box::new(std::io::Error::other(
+        "backend detail",
+    ))))
+}
+
+impl KeyStore for Unavailable {
+    fn create(&self, _: NewKey) -> StoreFuture<'_, ()> {
+        Box::pin(async { unavailable() })
+    }
+
+    fn get<'a>(&'a self, _: &'a KeyHash) -> StoreFuture<'a, Option<KeyRecord>> {
+        Box::pin(async { unavailable() })
+    }
+
+    fn revoke<'a>(&'a self, _: &'a KeyHash) -> StoreFuture<'a, Revocation> {
+        Box::pin(async { unavailable() })
+    }
+}