use std::collections::BTreeMap;
use std::fmt;
use async_trait::async_trait;
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use url::Url;
use crate::error::{Error, Result};
pub const CREDENTIAL_SCHEMA_VERSION: u32 = 2;
const LEGACY_SCHEMA_VERSION: u32 = 1;
const REDACTED: &str = "<redacted>";
const MISSING_ISSUER_REASON: &str =
"entry records no issuer; it cannot be re-keyed without guessing which \
authorization server issued it, so it was dropped rather than misattributed";
type AccountMap = BTreeMap<String, StoredCredentials>;
type ServerMap = BTreeMap<String, AccountMap>;
type IssuerMap = BTreeMap<String, ServerMap>;
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct CredentialKey {
issuer: String,
account: String,
server: String,
}
impl CredentialKey {
pub fn new<I, A, S>(issuer: I, account: A, server: S) -> Self
where
I: Into<String>,
A: Into<String>,
S: Into<String>,
{
Self {
issuer: issuer.into(),
account: account.into(),
server: server.into(),
}
}
pub fn issuer(&self) -> &str {
&self.issuer
}
pub fn account(&self) -> &str {
&self.account
}
pub fn server(&self) -> &str {
&self.server
}
}
#[derive(Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct StoredCredentials {
access_token: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
refresh_token: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
expires_at: Option<u64>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
scopes: Vec<String>,
client_id: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
registered_application_type: Option<String>,
}
impl StoredCredentials {
pub fn new(access_token: impl Into<String>, client_id: impl Into<String>) -> Self {
Self {
access_token: access_token.into(),
refresh_token: None,
expires_at: None,
scopes: Vec::new(),
client_id: client_id.into(),
registered_application_type: None,
}
}
pub fn with_refresh_token(mut self, refresh_token: impl Into<String>) -> Self {
self.refresh_token = Some(refresh_token.into());
self
}
pub fn with_expires_at(mut self, expires_at: u64) -> Self {
self.expires_at = Some(expires_at);
self
}
pub fn with_granted_scopes<S, I>(mut self, scopes: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
self.scopes = scopes.into_iter().map(Into::into).collect();
self
}
pub fn with_registered_application_type(mut self, application_type: impl Into<String>) -> Self {
self.registered_application_type = Some(application_type.into());
self
}
pub fn access_token(&self) -> &str {
&self.access_token
}
pub fn refresh_token(&self) -> Option<&str> {
self.refresh_token.as_deref()
}
pub fn expires_at(&self) -> Option<u64> {
self.expires_at
}
pub fn granted_scopes(&self) -> &[String] {
&self.scopes
}
pub fn client_id(&self) -> &str {
&self.client_id
}
pub fn registered_application_type(&self) -> Option<&str> {
self.registered_application_type.as_deref()
}
}
impl fmt::Debug for StoredCredentials {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("StoredCredentials")
.field("access_token", &REDACTED)
.field(
"refresh_token",
&self.refresh_token.as_ref().map(|_| REDACTED),
)
.field("expires_at", &self.expires_at)
.field("scopes", &self.scopes)
.field("client_id", &self.client_id)
.field(
"registered_application_type",
&self.registered_application_type,
)
.finish()
}
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct CredentialSnapshot {
credentials: IssuerMap,
issuers: BTreeMap<String, String>,
}
impl CredentialSnapshot {
pub fn new() -> Self {
Self::default()
}
pub fn get(&self, key: &CredentialKey) -> Option<&StoredCredentials> {
self.credentials
.get(&key.issuer)?
.get(&key.server)?
.get(&key.account)
}
pub fn insert(&mut self, key: CredentialKey, credentials: StoredCredentials) {
self.credentials
.entry(key.issuer)
.or_default()
.entry(key.server)
.or_default()
.insert(key.account, credentials);
}
pub fn remove(&mut self, key: &CredentialKey) -> bool {
let Some(by_server) = self.credentials.get_mut(&key.issuer) else {
return false;
};
let Some(by_account) = by_server.get_mut(&key.server) else {
return false;
};
let removed = by_account.remove(&key.account).is_some();
if by_account.is_empty() {
by_server.remove(&key.server);
}
if by_server.is_empty() {
self.credentials.remove(&key.issuer);
}
removed
}
pub fn keys(&self) -> Vec<CredentialKey> {
let mut out = Vec::new();
for (issuer, by_server) in &self.credentials {
for (server, by_account) in by_server {
for account in by_account.keys() {
out.push(CredentialKey::new(issuer, account, server));
}
}
}
out
}
pub fn keys_for_server(&self, server_key: &str) -> Vec<CredentialKey> {
let mut out = Vec::new();
for (issuer, by_server) in &self.credentials {
if let Some(by_account) = by_server.get(server_key) {
for account in by_account.keys() {
out.push(CredentialKey::new(issuer, account, server_key));
}
}
}
out
}
pub fn clear(&mut self) -> usize {
let removed = self.credential_count();
self.credentials.clear();
self.issuers.clear();
removed
}
pub fn last_issuer(&self, server_key: &str) -> Option<&str> {
self.issuers.get(server_key).map(String::as_str)
}
pub fn record_issuer(&mut self, server_key: &str, issuer: &str) {
self.issuers
.insert(server_key.to_owned(), issuer.to_owned());
}
pub fn to_bytes(&self) -> Result<Vec<u8>> {
let document = DocumentRef {
schema_version: CREDENTIAL_SCHEMA_VERSION,
credentials: &self.credentials,
issuers: &self.issuers,
};
serde_json::to_vec_pretty(&document)
.map_err(|e| Error::internal(format!("failed to serialize credentials: {e}")))
}
pub(crate) fn forget_issuer(&mut self, server_key: &str) {
self.issuers.remove(server_key);
}
fn credential_count(&self) -> usize {
self.credentials
.values()
.map(|by_server| by_server.values().map(BTreeMap::len).sum::<usize>())
.sum()
}
}
#[derive(Serialize)]
struct DocumentRef<'a> {
schema_version: u32,
credentials: &'a IssuerMap,
issuers: &'a BTreeMap<String, String>,
}
#[derive(Deserialize)]
struct Document {
#[serde(default)]
credentials: IssuerMap,
#[serde(default)]
issuers: BTreeMap<String, String>,
}
#[derive(Deserialize)]
struct VersionProbe {
schema_version: u32,
}
#[derive(Deserialize)]
struct LegacyCache {
#[serde(default)]
entries: BTreeMap<String, LegacyEntry>,
}
#[derive(Deserialize)]
struct LegacyEntry {
access_token: String,
#[serde(default)]
refresh_token: Option<String>,
#[serde(default)]
expires_at: Option<u64>,
#[serde(default)]
scopes: Vec<String>,
#[serde(default)]
issuer: Option<String>,
#[serde(default)]
client_id: String,
}
impl LegacyEntry {
fn into_stored(self) -> StoredCredentials {
StoredCredentials {
access_token: self.access_token,
refresh_token: self.refresh_token,
expires_at: self.expires_at,
scopes: self.scopes,
client_id: self.client_id,
registered_application_type: None,
}
}
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
#[non_exhaustive]
pub struct MigrationReport {
migrated: usize,
dropped: Vec<DroppedEntry>,
}
impl MigrationReport {
pub fn migrated(&self) -> usize {
self.migrated
}
pub fn dropped(&self) -> &[DroppedEntry] {
&self.dropped
}
pub fn is_noop(&self) -> bool {
self.migrated == 0 && self.dropped.is_empty()
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct DroppedEntry {
server_key: String,
reason: String,
}
impl DroppedEntry {
pub fn server_key(&self) -> &str {
&self.server_key
}
pub fn reason(&self) -> &str {
&self.reason
}
}
pub fn parse_credential_snapshot(bytes: &[u8]) -> Result<(CredentialSnapshot, MigrationReport)> {
let probe: VersionProbe = serde_json::from_slice(bytes).map_err(|e| malformed_document(&e))?;
match probe.schema_version {
CREDENTIAL_SCHEMA_VERSION => parse_current(bytes),
LEGACY_SCHEMA_VERSION => migrate_legacy(bytes),
observed => Err(unsupported_schema_version(observed)),
}
}
fn parse_current(bytes: &[u8]) -> Result<(CredentialSnapshot, MigrationReport)> {
let document: Document = serde_json::from_slice(bytes).map_err(|e| malformed_document(&e))?;
let snapshot = CredentialSnapshot {
credentials: document.credentials,
issuers: document.issuers,
};
Ok((snapshot, MigrationReport::default()))
}
fn migrate_legacy(bytes: &[u8]) -> Result<(CredentialSnapshot, MigrationReport)> {
let cache: LegacyCache = serde_json::from_slice(bytes).map_err(|e| malformed_document(&e))?;
let mut snapshot = CredentialSnapshot::new();
let mut migrated = 0usize;
let mut dropped = Vec::new();
for (server_key, mut entry) in cache.entries {
let recorded = entry.issuer.take().filter(|value| !value.is_empty());
let Some(issuer) = recorded else {
dropped.push(DroppedEntry {
server_key,
reason: MISSING_ISSUER_REASON.to_owned(),
});
continue;
};
snapshot.record_issuer(&server_key, &issuer);
snapshot.insert(
CredentialKey::new(issuer, "", server_key),
entry.into_stored(),
);
migrated += 1;
}
Ok((snapshot, MigrationReport { migrated, dropped }))
}
fn malformed_document(err: &serde_json::Error) -> Error {
Error::validation(format!(
"credential document is malformed: {:?} error at line {}, column {}",
err.classify(),
err.line(),
err.column()
))
}
fn unsupported_schema_version(observed: u32) -> Error {
Error::validation(format!(
"credential document schema_version {observed} is not supported by this build, \
which reads version {CREDENTIAL_SCHEMA_VERSION}; upgrade pmcp to read it"
))
}
pub fn normalize_server_key(server_url: &str) -> Result<String> {
let parsed = Url::parse(server_url)
.map_err(|e| Error::validation(format!("invalid MCP server URL ({e})")))?;
let host = parsed
.host_str()
.ok_or_else(|| Error::validation("MCP server URL has no host"))?
.to_ascii_lowercase();
let mut key = format!("{}://{}", parsed.scheme(), host);
if let Some(port) = parsed.port() {
let is_default = (parsed.scheme() == "https" && port == 443)
|| (parsed.scheme() == "http" && port == 80);
if !is_default {
key.push_str(&format!(":{port}"));
}
}
Ok(key)
}
#[async_trait]
pub trait CredentialStore: Send + Sync + fmt::Debug {
async fn load(&self, key: &CredentialKey) -> Result<Option<StoredCredentials>>;
async fn save(&self, key: &CredentialKey, credentials: &StoredCredentials) -> Result<()>;
async fn delete(&self, key: &CredentialKey) -> Result<()>;
async fn save_with_issuer(
&self,
key: &CredentialKey,
credentials: &StoredCredentials,
server_key: &str,
issuer: &str,
) -> Result<()> {
self.save(key, credentials).await?;
self.record_issuer(server_key, issuer).await
}
async fn last_issuer(&self, _server_key: &str) -> Result<Option<String>> {
Ok(None)
}
async fn record_issuer(&self, _server_key: &str, _issuer: &str) -> Result<()> {
Ok(())
}
}
#[async_trait]
pub trait CredentialStoreAdmin: CredentialStore {
async fn list_keys(&self) -> Result<Vec<CredentialKey>>;
async fn delete_by_server(&self, server_key: &str) -> Result<usize>;
async fn clear_all(&self) -> Result<usize>;
async fn take_migration_report(&self) -> Result<Option<MigrationReport>>;
}
#[derive(Debug, Default)]
pub struct InMemoryCredentialStore {
snapshot: RwLock<CredentialSnapshot>,
migration_report: RwLock<Option<MigrationReport>>,
}
impl InMemoryCredentialStore {
pub fn new() -> Self {
Self::default()
}
pub fn from_bytes(bytes: &[u8]) -> Result<Self> {
let (snapshot, report) = parse_credential_snapshot(bytes)?;
let retained = if report.is_noop() { None } else { Some(report) };
Ok(Self {
snapshot: RwLock::new(snapshot),
migration_report: RwLock::new(retained),
})
}
}
#[async_trait]
impl CredentialStore for InMemoryCredentialStore {
async fn load(&self, key: &CredentialKey) -> Result<Option<StoredCredentials>> {
Ok(self.snapshot.read().get(key).cloned())
}
async fn save(&self, key: &CredentialKey, credentials: &StoredCredentials) -> Result<()> {
self.snapshot
.write()
.insert(key.clone(), credentials.clone());
Ok(())
}
async fn delete(&self, key: &CredentialKey) -> Result<()> {
self.snapshot.write().remove(key);
Ok(())
}
async fn save_with_issuer(
&self,
key: &CredentialKey,
credentials: &StoredCredentials,
server_key: &str,
issuer: &str,
) -> Result<()> {
let mut snapshot = self.snapshot.write();
snapshot.insert(key.clone(), credentials.clone());
snapshot.record_issuer(server_key, issuer);
Ok(())
}
async fn last_issuer(&self, server_key: &str) -> Result<Option<String>> {
Ok(self
.snapshot
.read()
.last_issuer(server_key)
.map(str::to_owned))
}
async fn record_issuer(&self, server_key: &str, issuer: &str) -> Result<()> {
self.snapshot.write().record_issuer(server_key, issuer);
Ok(())
}
}
#[async_trait]
impl CredentialStoreAdmin for InMemoryCredentialStore {
async fn list_keys(&self) -> Result<Vec<CredentialKey>> {
Ok(self.snapshot.read().keys())
}
async fn delete_by_server(&self, server_key: &str) -> Result<usize> {
let mut snapshot = self.snapshot.write();
let mut removed = 0usize;
for key in snapshot.keys_for_server(server_key) {
if snapshot.remove(&key) {
removed += 1;
}
}
snapshot.forget_issuer(server_key);
Ok(removed)
}
async fn clear_all(&self) -> Result<usize> {
Ok(self.snapshot.write().clear())
}
async fn take_migration_report(&self) -> Result<Option<MigrationReport>> {
Ok(self.migration_report.write().take())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn the_redaction_marker_is_not_a_field_name_substring() {
for field in [
"access_token",
"refresh_token",
"expires_at",
"scopes",
"client_id",
"registered_application_type",
] {
assert!(!field.contains(REDACTED));
assert!(!REDACTED.contains(field));
}
}
#[test]
fn a_legacy_entry_carries_every_field_across() {
let entry = LegacyEntry {
access_token: "at".to_owned(),
refresh_token: Some("rt".to_owned()),
expires_at: Some(42),
scopes: vec!["a".to_owned()],
issuer: Some("https://as.example".to_owned()),
client_id: "cid".to_owned(),
};
let stored = entry.into_stored();
assert_eq!(stored.access_token(), "at");
assert_eq!(stored.refresh_token(), Some("rt"));
assert_eq!(stored.expires_at(), Some(42));
assert_eq!(stored.granted_scopes(), ["a"]);
assert_eq!(stored.client_id(), "cid");
assert!(stored.registered_application_type().is_none());
}
#[test]
fn an_emptied_issuer_map_is_pruned_so_keys_stays_accurate() {
let key = CredentialKey::new("https://as.example", "", "https://mcp.example");
let mut snapshot = CredentialSnapshot::new();
snapshot.insert(key.clone(), StoredCredentials::new("at", "cid"));
assert!(snapshot.remove(&key));
assert!(snapshot.credentials.is_empty(), "empty maps must be pruned");
assert_eq!(snapshot.credential_count(), 0);
}
#[test]
fn forget_issuer_touches_only_the_named_server() {
let mut snapshot = CredentialSnapshot::new();
snapshot.record_issuer("https://a.example", "https://as.example");
snapshot.record_issuer("https://b.example", "https://as.example");
snapshot.forget_issuer("https://a.example");
assert!(snapshot.last_issuer("https://a.example").is_none());
assert!(snapshot.last_issuer("https://b.example").is_some());
}
#[test]
fn an_unsupported_version_refusal_names_both_versions() {
let message = unsupported_schema_version(7).to_string();
assert!(message.contains('7'), "{message}");
assert!(message.contains('2'), "{message}");
}
}