use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use async_trait::async_trait;
use base64::Engine as _;
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use serde::{Deserialize, Serialize};
use cheers_core::{
CodecError, Principal, PrincipalError, PrincipalId, PrincipalKind, PrincipalStatus, StoreError,
};
use crate::codec::PasetoV4SecretMinter;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[non_exhaustive]
#[serde(rename_all = "snake_case")]
pub enum SigningKeyStatus {
Active,
Retiring,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[non_exhaustive]
pub struct SigningKey {
pub kid: String,
pub principal_id: PrincipalId,
#[serde(with = "public_key_serde")]
pub public_key: [u8; 32],
pub status: SigningKeyStatus,
pub created_at: i64,
pub retire_at: Option<i64>,
}
impl SigningKey {
pub fn new(
kid: impl Into<String>,
principal_id: PrincipalId,
public_key: [u8; 32],
status: SigningKeyStatus,
created_at: i64,
retire_at: Option<i64>,
) -> Self {
Self {
kid: kid.into(),
principal_id,
public_key,
status,
created_at,
retire_at,
}
}
}
mod public_key_serde {
use super::*;
use serde::de::Error as DeError;
pub fn serialize<S: serde::Serializer>(bytes: &[u8; 32], ser: S) -> Result<S::Ok, S::Error> {
ser.serialize_str(&URL_SAFE_NO_PAD.encode(bytes))
}
pub fn deserialize<'de, D: serde::Deserializer<'de>>(de: D) -> Result<[u8; 32], D::Error> {
let s = String::deserialize(de)?;
let decoded = URL_SAFE_NO_PAD
.decode(s.as_bytes())
.map_err(|e| D::Error::custom(format!("invalid base64url public_key: {e}")))?;
decoded
.try_into()
.map_err(|v: Vec<u8>| D::Error::custom(format!("expected 32 public_key bytes, got {}", v.len())))
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[non_exhaustive]
pub struct NewServicePrincipal {
pub desired_id: String,
}
impl NewServicePrincipal {
pub fn new(desired_id: impl Into<String>) -> Self {
Self {
desired_id: desired_id.into(),
}
}
}
#[derive(Clone)]
#[non_exhaustive]
pub struct ProvisionedKey {
pub principal: Principal,
pub signing_key: SigningKey,
pub secret_key: [u8; 64],
}
impl std::fmt::Debug for ProvisionedKey {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ProvisionedKey")
.field("principal", &self.principal)
.field("signing_key", &self.signing_key)
.field("secret_key", &"***OMITTED***")
.finish()
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub struct OverlapPolicy {
pub overlap_seconds: i64,
}
impl OverlapPolicy {
pub const DEFAULT_OVERLAP_SECONDS: i64 = 24 * 60 * 60;
pub fn new(overlap_seconds: i64) -> Self {
Self { overlap_seconds }
}
pub fn with_overlap_seconds(mut self, seconds: i64) -> Self {
self.overlap_seconds = seconds;
self
}
}
impl Default for OverlapPolicy {
fn default() -> Self {
Self {
overlap_seconds: Self::DEFAULT_OVERLAP_SECONDS,
}
}
}
#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum ServicePrincipalError {
#[error("service principal '{0}' already exists")]
AlreadyExists(PrincipalId),
#[error("unknown service principal '{0}'")]
UnknownPrincipal(PrincipalId),
#[error(transparent)]
Principal(#[from] PrincipalError),
#[error("service-principal authority requires kind=service; got {0}")]
WrongPrincipalKind(PrincipalKind),
#[error("service principal '{0}' has no active signing key")]
NoActiveKey(PrincipalId),
#[error(transparent)]
Codec(#[from] CodecError),
#[error(transparent)]
Store(#[from] StoreError),
}
#[async_trait]
pub trait ServicePrincipalStore: Send + Sync {
async fn insert_principal(&self, principal: &Principal) -> Result<(), StoreError>;
async fn get_principal(
&self,
id: &PrincipalId,
) -> Result<Option<Principal>, StoreError>;
async fn insert_signing_key(&self, key: &SigningKey) -> Result<(), StoreError>;
async fn list_signing_keys(
&self,
principal: &PrincipalId,
) -> Result<Vec<SigningKey>, StoreError>;
async fn list_all_signing_keys(&self) -> Result<Vec<SigningKey>, StoreError>;
async fn retire_signing_key(
&self,
kid: &str,
retire_at: i64,
) -> Result<(), StoreError>;
async fn prune_retired_keys(&self, now: i64) -> Result<u64, StoreError>;
}
pub struct ServicePrincipalAuthority<S> {
store: S,
policy: OverlapPolicy,
}
impl<S> std::fmt::Debug for ServicePrincipalAuthority<S> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ServicePrincipalAuthority")
.field("policy", &self.policy)
.finish_non_exhaustive()
}
}
impl<S> ServicePrincipalAuthority<S>
where
S: ServicePrincipalStore,
{
pub fn new(store: S) -> Self {
Self {
store,
policy: OverlapPolicy::default(),
}
}
pub fn with_policy(mut self, policy: OverlapPolicy) -> Self {
self.policy = policy;
self
}
pub fn policy(&self) -> &OverlapPolicy {
&self.policy
}
pub fn store(&self) -> &S {
&self.store
}
pub async fn provision(
&self,
input: NewServicePrincipal,
now: i64,
) -> Result<ProvisionedKey, ServicePrincipalError> {
let id = PrincipalId::service(input.desired_id);
let principal = Principal::try_new(id.clone(), None, PrincipalStatus::Active, now)?;
match self.store.insert_principal(&principal).await {
Ok(()) => {}
Err(StoreError::Conflict) => return Err(ServicePrincipalError::AlreadyExists(id)),
Err(other) => return Err(other.into()),
}
let (signing_key, secret_key) = self.mint_signing_key(&id, now).await?;
Ok(ProvisionedKey {
principal,
signing_key,
secret_key,
})
}
pub async fn rotate(
&self,
id: &PrincipalId,
now: i64,
) -> Result<ProvisionedKey, ServicePrincipalError> {
if id.kind != PrincipalKind::Service {
return Err(ServicePrincipalError::WrongPrincipalKind(id.kind));
}
let principal = self
.store
.get_principal(id)
.await?
.ok_or_else(|| ServicePrincipalError::UnknownPrincipal(id.clone()))?;
let retire_at = now + self.policy.overlap_seconds;
let existing = self.store.list_signing_keys(id).await?;
let active = existing
.iter()
.find(|k| matches!(k.status, SigningKeyStatus::Active))
.ok_or_else(|| ServicePrincipalError::NoActiveKey(id.clone()))?;
self.store
.retire_signing_key(&active.kid, retire_at)
.await?;
let (signing_key, secret_key) = self.mint_signing_key(id, now).await?;
Ok(ProvisionedKey {
principal,
signing_key,
secret_key,
})
}
pub async fn prune_retired_keys(
&self,
now: i64,
) -> Result<u64, ServicePrincipalError> {
Ok(self.store.prune_retired_keys(now).await?)
}
pub async fn published_signing_keys(
&self,
now: i64,
) -> Result<Vec<SigningKey>, ServicePrincipalError> {
let all = self.store.list_all_signing_keys().await?;
Ok(all
.into_iter()
.filter(|k| match k.status {
SigningKeyStatus::Active => true,
SigningKeyStatus::Retiring => k.retire_at.map(|t| t > now).unwrap_or(false),
})
.collect())
}
async fn mint_signing_key(
&self,
principal_id: &PrincipalId,
now: i64,
) -> Result<(SigningKey, [u8; 64]), ServicePrincipalError> {
let (minter, verifier) = PasetoV4SecretMinter::generate()?;
let public_bytes = verifier
.public_key()
.as_bytes()
.try_into()
.expect("v4.public key is 32 bytes");
let secret_slice = minter.secret_key_bytes();
let mut secret_key = [0u8; 64];
secret_key.copy_from_slice(secret_slice);
let signing_key = SigningKey {
kid: mint_kid(),
principal_id: principal_id.clone(),
public_key: public_bytes,
status: SigningKeyStatus::Active,
created_at: now,
retire_at: None,
};
self.store.insert_signing_key(&signing_key).await?;
Ok((signing_key, secret_key))
}
}
fn mint_kid() -> String {
let mut bytes = [0u8; 16];
getrandom::fill(&mut bytes).expect("OS CSPRNG must be available");
URL_SAFE_NO_PAD.encode(bytes)
}
#[derive(Default, Clone, Debug)]
pub struct MemoryServicePrincipalStore {
inner: Arc<Mutex<MemoryServicePrincipalInner>>,
}
#[derive(Default, Debug)]
struct MemoryServicePrincipalInner {
principals: HashMap<PrincipalId, Principal>,
keys: Vec<SigningKey>,
}
impl MemoryServicePrincipalStore {
pub fn new() -> Self {
Self::default()
}
}
#[async_trait]
impl ServicePrincipalStore for MemoryServicePrincipalStore {
async fn insert_principal(&self, principal: &Principal) -> Result<(), StoreError> {
let mut g = self.inner.lock().expect("service-principal store mutex poisoned");
if g.principals.contains_key(&principal.id) {
return Err(StoreError::Conflict);
}
g.principals.insert(principal.id.clone(), principal.clone());
Ok(())
}
async fn get_principal(
&self,
id: &PrincipalId,
) -> Result<Option<Principal>, StoreError> {
let g = self.inner.lock().expect("service-principal store mutex poisoned");
Ok(g.principals.get(id).cloned())
}
async fn insert_signing_key(&self, key: &SigningKey) -> Result<(), StoreError> {
let mut g = self.inner.lock().expect("service-principal store mutex poisoned");
if g.keys.iter().any(|k| k.kid == key.kid) {
return Err(StoreError::Conflict);
}
g.keys.push(key.clone());
Ok(())
}
async fn list_signing_keys(
&self,
principal: &PrincipalId,
) -> Result<Vec<SigningKey>, StoreError> {
let g = self.inner.lock().expect("service-principal store mutex poisoned");
Ok(g.keys
.iter()
.filter(|k| &k.principal_id == principal)
.cloned()
.collect())
}
async fn list_all_signing_keys(&self) -> Result<Vec<SigningKey>, StoreError> {
let g = self.inner.lock().expect("service-principal store mutex poisoned");
Ok(g.keys.clone())
}
async fn retire_signing_key(
&self,
kid: &str,
retire_at: i64,
) -> Result<(), StoreError> {
let mut g = self.inner.lock().expect("service-principal store mutex poisoned");
let key = g
.keys
.iter_mut()
.find(|k| k.kid == kid)
.ok_or(StoreError::NotFound)?;
key.status = SigningKeyStatus::Retiring;
key.retire_at = Some(retire_at);
Ok(())
}
async fn prune_retired_keys(&self, now: i64) -> Result<u64, StoreError> {
let mut g = self.inner.lock().expect("service-principal store mutex poisoned");
let before = g.keys.len();
g.keys.retain(|k| match k.status {
SigningKeyStatus::Active => true,
SigningKeyStatus::Retiring => k.retire_at.map(|t| t > now).unwrap_or(true),
});
Ok((before - g.keys.len()) as u64)
}
}
#[cfg(test)]
mod tests {
use super::*;
use cheers_verify::PasetoV4PublicVerifier;
use pollster::block_on;
fn rig() -> ServicePrincipalAuthority<MemoryServicePrincipalStore> {
ServicePrincipalAuthority::new(MemoryServicePrincipalStore::new())
}
#[test]
fn overlap_policy_default_is_24_hours() {
let p = OverlapPolicy::default();
assert_eq!(p.overlap_seconds, 24 * 60 * 60);
let custom = OverlapPolicy::new(3_600).with_overlap_seconds(7_200);
assert_eq!(custom.overlap_seconds, 7_200);
}
#[test]
fn provision_allocates_principal_and_returns_secret_once() {
let authority = rig();
block_on(async {
let provisioned = authority
.provision(NewServicePrincipal::new("yubaba-1"), 1_000)
.await
.unwrap();
assert_eq!(
provisioned.principal.id,
PrincipalId::service("yubaba-1")
);
assert_eq!(provisioned.principal.bound_to, None);
assert_eq!(provisioned.principal.status, PrincipalStatus::Active);
assert_eq!(provisioned.principal.created_at, 1_000);
assert_eq!(provisioned.signing_key.principal_id, PrincipalId::service("yubaba-1"));
assert_eq!(provisioned.signing_key.status, SigningKeyStatus::Active);
assert!(provisioned.signing_key.retire_at.is_none());
assert_eq!(provisioned.signing_key.created_at, 1_000);
assert!(!provisioned.signing_key.kid.is_empty());
assert_eq!(
&provisioned.secret_key[32..],
&provisioned.signing_key.public_key,
"PASETO V4 secret layout is seed||pubkey",
);
});
}
#[test]
fn provision_persisted_pubkey_verifies_tokens_minted_by_returned_secret() {
use cheers_core::{
Actor, AuthStrength, McpClaims, Owns, PrincipalId, Scope,
};
let authority = rig();
block_on(async {
let provisioned = authority
.provision(NewServicePrincipal::new("yubaba-2"), 1_000)
.await
.unwrap();
let off_cheers_minter =
PasetoV4SecretMinter::from_secret_key(&provisioned.secret_key).unwrap();
let verifier =
PasetoV4PublicVerifier::from_public_key(&provisioned.signing_key.public_key)
.unwrap();
let mut owns = Owns::default();
owns.service = vec!["svc-prod".into()];
let claims = McpClaims::new(
"https://cheers.example",
"https://kamaji.example",
PrincipalId::service("yubaba-2"),
1_000,
1_600,
"jti-rt",
vec![Scope::OwnershipWrite],
)
.with_act(Actor::new(PrincipalId::service("yubaba-2")))
.with_owns(owns)
.with_auth_strength(AuthStrength::Bootstrap);
let token = off_cheers_minter
.mint_mcp(&claims, &provisioned.signing_key.kid)
.unwrap();
let back = verifier
.verify_mcp_at(&token, 1_100, &provisioned.signing_key.kid)
.unwrap();
assert_eq!(back, claims);
});
}
#[test]
fn provision_rejects_duplicate_id() {
let authority = rig();
block_on(async {
authority
.provision(NewServicePrincipal::new("dup"), 1_000)
.await
.unwrap();
let err = authority
.provision(NewServicePrincipal::new("dup"), 1_001)
.await
.unwrap_err();
match err {
ServicePrincipalError::AlreadyExists(id) => {
assert_eq!(id, PrincipalId::service("dup"));
}
other => panic!("expected AlreadyExists, got {other:?}"),
}
});
}
#[test]
fn rotate_retires_old_key_and_returns_fresh_secret() {
let authority = rig();
block_on(async {
let first = authority
.provision(NewServicePrincipal::new("yubaba-r"), 1_000)
.await
.unwrap();
let id = first.principal.id.clone();
let second = authority.rotate(&id, 5_000).await.unwrap();
assert_ne!(second.signing_key.kid, first.signing_key.kid);
assert_ne!(second.signing_key.public_key, first.signing_key.public_key);
assert_eq!(second.signing_key.status, SigningKeyStatus::Active);
assert!(second.signing_key.retire_at.is_none());
let stored = authority
.store
.list_signing_keys(&id)
.await
.unwrap();
assert_eq!(stored.len(), 2);
let old = stored.iter().find(|k| k.kid == first.signing_key.kid).unwrap();
assert_eq!(old.status, SigningKeyStatus::Retiring);
assert_eq!(
old.retire_at,
Some(5_000 + OverlapPolicy::DEFAULT_OVERLAP_SECONDS)
);
});
}
#[test]
fn rotate_respects_custom_overlap_window() {
let authority = rig().with_policy(OverlapPolicy::new(60 * 60)); block_on(async {
let first = authority
.provision(NewServicePrincipal::new("yubaba-c"), 1_000)
.await
.unwrap();
authority.rotate(&first.principal.id, 2_000).await.unwrap();
let keys = authority
.store
.list_signing_keys(&first.principal.id)
.await
.unwrap();
let old = keys.iter().find(|k| k.kid == first.signing_key.kid).unwrap();
assert_eq!(old.retire_at, Some(2_000 + 3_600));
});
}
#[test]
fn rotate_rejects_unknown_principal() {
let authority = rig();
block_on(async {
let err = authority
.rotate(&PrincipalId::service("ghost"), 1_000)
.await
.unwrap_err();
match err {
ServicePrincipalError::UnknownPrincipal(id) => {
assert_eq!(id, PrincipalId::service("ghost"));
}
other => panic!("expected UnknownPrincipal, got {other:?}"),
}
});
}
#[test]
fn rotate_rejects_non_service_principal_kind() {
let authority = rig();
block_on(async {
let err = authority
.rotate(&PrincipalId::user("alice"), 1_000)
.await
.unwrap_err();
match err {
ServicePrincipalError::WrongPrincipalKind(k) => {
assert_eq!(k, PrincipalKind::User);
}
other => panic!("expected WrongPrincipalKind, got {other:?}"),
}
});
}
#[test]
fn published_signing_keys_includes_active_and_retiring_within_window() {
let authority = rig().with_policy(OverlapPolicy::new(100));
block_on(async {
let first = authority
.provision(NewServicePrincipal::new("yubaba-p"), 1_000)
.await
.unwrap();
authority.rotate(&first.principal.id, 2_000).await.unwrap();
let live = authority.published_signing_keys(2_050).await.unwrap();
assert_eq!(live.len(), 2, "got {live:#?}");
let live = authority.published_signing_keys(2_099).await.unwrap();
assert_eq!(live.len(), 2);
let live = authority.published_signing_keys(2_100).await.unwrap();
assert_eq!(live.len(), 1);
assert_eq!(live[0].status, SigningKeyStatus::Active);
});
}
#[test]
fn prune_retired_keys_removes_expired_only() {
let authority = rig().with_policy(OverlapPolicy::new(100));
block_on(async {
let first = authority
.provision(NewServicePrincipal::new("yubaba-x"), 1_000)
.await
.unwrap();
authority.rotate(&first.principal.id, 2_000).await.unwrap();
let dropped = authority.prune_retired_keys(2_099).await.unwrap();
assert_eq!(dropped, 0);
assert_eq!(
authority
.store
.list_signing_keys(&first.principal.id)
.await
.unwrap()
.len(),
2
);
let dropped = authority.prune_retired_keys(2_100).await.unwrap();
assert_eq!(dropped, 1);
let remaining = authority
.store
.list_signing_keys(&first.principal.id)
.await
.unwrap();
assert_eq!(remaining.len(), 1);
assert_eq!(remaining[0].status, SigningKeyStatus::Active);
let dropped = authority.prune_retired_keys(2_200).await.unwrap();
assert_eq!(dropped, 0);
});
}
#[test]
fn published_signing_keys_omits_keys_with_no_retire_at_in_retiring_state() {
let store = MemoryServicePrincipalStore::new();
let key = SigningKey {
kid: "k".into(),
principal_id: PrincipalId::service("zombie"),
public_key: [0u8; 32],
status: SigningKeyStatus::Retiring,
created_at: 100,
retire_at: None,
};
block_on(async {
store.insert_signing_key(&key).await.unwrap();
let authority = ServicePrincipalAuthority::new(store);
let live = authority.published_signing_keys(1_000).await.unwrap();
assert!(live.is_empty());
});
}
#[test]
fn signing_key_roundtrips_through_json_with_base64_pubkey() {
let key = SigningKey {
kid: "kid-abc".into(),
principal_id: PrincipalId::service("yubaba"),
public_key: [1u8; 32],
status: SigningKeyStatus::Active,
created_at: 100,
retire_at: None,
};
let json = serde_json::to_string(&key).unwrap();
assert!(json.contains("\"public_key\":\""), "got {json}");
assert!(!json.contains("[1,1,1"), "must NOT be an array: {json}");
let back: SigningKey = serde_json::from_str(&json).unwrap();
assert_eq!(back, key);
}
#[test]
fn signing_key_deserialize_rejects_wrong_length_pubkey() {
let json = r#"{
"kid": "k",
"principal_id": "svc:x",
"public_key": "AAAA",
"status": "active",
"created_at": 0,
"retire_at": null
}"#;
let err = serde_json::from_str::<SigningKey>(json).unwrap_err();
assert!(
err.to_string().contains("32 public_key bytes"),
"got {err}"
);
}
#[test]
fn service_principal_store_is_dyn_compatible() {
fn _f(_: &dyn ServicePrincipalStore) {}
}
}