use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use base64::{engine::general_purpose::STANDARD as BASE64, Engine};
use chrono::{DateTime, Utc};
use ed25519_dalek::VerifyingKey;
use tokio::sync::RwLock;
use crate::error::{RegistryError, RegistryResult};
use crate::types::{KeysManifest, TrustedKey};
use crate::verify::compute_key_id;
const DEFAULT_KEYS_TTL_SECS: i64 = 24 * 60 * 60;
const PRODUCTION_TRUST_ROOTS_JSON: &str = include_str!("../assets/production-trust-roots.json");
#[derive(Debug, Clone)]
pub struct TrustStore {
inner: Arc<RwLock<TrustStoreInner>>,
}
#[derive(Debug)]
struct TrustStoreInner {
keys: HashMap<String, VerifyingKey>,
metadata: HashMap<String, KeyMetadata>,
pinned_roots: Vec<String>,
manifest_fetched_at: Option<DateTime<Utc>>,
manifest_expires_at: Option<DateTime<Utc>>,
}
#[derive(Debug, Clone)]
pub struct KeyMetadata {
pub description: Option<String>,
pub added_at: Option<DateTime<Utc>>,
pub expires_at: Option<DateTime<Utc>>,
pub revoked: bool,
pub is_pinned: bool,
}
impl TrustStore {
pub fn new() -> Self {
Self {
inner: Arc::new(RwLock::new(Self::empty_inner())),
}
}
pub fn from_pinned_roots(roots: Vec<TrustedKey>) -> RegistryResult<Self> {
let mut inner = Self::empty_inner();
for root in &roots {
insert_pinned_key(&mut inner, root)?;
}
Ok(Self {
inner: Arc::new(RwLock::new(inner)),
})
}
pub async fn with_pinned_roots(roots: Vec<TrustedKey>) -> RegistryResult<Self> {
Self::from_pinned_roots(roots)
}
pub fn from_production_roots() -> RegistryResult<Self> {
load_production_roots_impl(PRODUCTION_TRUST_ROOTS_JSON)
}
pub async fn with_production_roots() -> RegistryResult<Self> {
Self::from_production_roots()
}
pub async fn add_pinned_key(&self, key: &TrustedKey) -> RegistryResult<()> {
let verifying_key = decode_verifying_key(&key.public_key)?;
let computed_id = compute_key_id(&decode_public_key_bytes(&key.public_key)?);
if computed_id != key.key_id {
return Err(RegistryError::SignatureInvalid {
reason: format!(
"key_id mismatch: claimed {}, computed {}",
key.key_id, computed_id
),
});
}
let mut inner = self.inner.write().await;
inner.keys.insert(key.key_id.clone(), verifying_key);
inner.metadata.insert(
key.key_id.clone(),
KeyMetadata {
description: key.description.clone(),
added_at: key.added_at,
expires_at: key.expires_at,
revoked: false, is_pinned: true,
},
);
if !inner.pinned_roots.contains(&key.key_id) {
inner.pinned_roots.push(key.key_id.clone());
}
Ok(())
}
pub async fn add_from_manifest(&self, manifest: &KeysManifest) -> RegistryResult<()> {
let now = Utc::now();
let mut inner = self.inner.write().await;
for key in &manifest.keys {
if key.revoked {
if !inner.pinned_roots.contains(&key.key_id) {
inner.keys.remove(&key.key_id);
if let Some(meta) = inner.metadata.get_mut(&key.key_id) {
meta.revoked = true;
}
}
continue;
}
if let Some(expires_at) = key.expires_at {
if expires_at < now {
continue;
}
}
if inner.pinned_roots.contains(&key.key_id) {
continue;
}
match decode_verifying_key(&key.public_key) {
Ok(verifying_key) => {
let computed_id = match decode_public_key_bytes(&key.public_key) {
Ok(bytes) => compute_key_id(&bytes),
Err(_) => continue,
};
if computed_id != key.key_id {
tracing::warn!(
claimed = %key.key_id,
computed = %computed_id,
"key_id mismatch, skipping"
);
continue;
}
inner.keys.insert(key.key_id.clone(), verifying_key);
inner.metadata.insert(
key.key_id.clone(),
KeyMetadata {
description: key.description.clone(),
added_at: key.added_at,
expires_at: key.expires_at,
revoked: false,
is_pinned: false,
},
);
}
Err(e) => {
tracing::warn!(key_id = %key.key_id, error = %e, "failed to decode key");
}
}
}
inner.manifest_fetched_at = Some(now);
inner.manifest_expires_at = manifest
.expires_at
.or(Some(now + chrono::Duration::seconds(DEFAULT_KEYS_TTL_SECS)));
Ok(())
}
pub async fn get_key_async(&self, key_id: &str) -> RegistryResult<VerifyingKey> {
let inner = self.inner.read().await;
self.get_key_inner(&inner, key_id)
}
pub fn get_key(&self, key_id: &str) -> RegistryResult<VerifyingKey> {
match self.inner.try_read() {
Ok(inner) => self.get_key_inner(&inner, key_id),
Err(_) => Err(RegistryError::KeyNotTrusted {
key_id: key_id.to_string(),
}),
}
}
fn get_key_inner(&self, inner: &TrustStoreInner, key_id: &str) -> RegistryResult<VerifyingKey> {
let key = inner
.keys
.get(key_id)
.ok_or_else(|| RegistryError::KeyNotTrusted {
key_id: key_id.to_string(),
})?;
if let Some(meta) = inner.metadata.get(key_id) {
if meta.revoked {
return Err(RegistryError::KeyNotTrusted {
key_id: key_id.to_string(),
});
}
if let Some(expires_at) = meta.expires_at {
if expires_at < Utc::now() {
return Err(RegistryError::KeyNotTrusted {
key_id: key_id.to_string(),
});
}
}
}
Ok(*key)
}
pub async fn needs_refresh(&self) -> bool {
let inner = self.inner.read().await;
match inner.manifest_expires_at {
Some(expires_at) => Utc::now() >= expires_at,
None => inner.manifest_fetched_at.is_none(),
}
}
pub async fn is_trusted(&self, key_id: &str) -> bool {
self.get_key_async(key_id).await.is_ok()
}
pub async fn list_keys(&self) -> Vec<String> {
let inner = self.inner.read().await;
inner.keys.keys().cloned().collect()
}
pub async fn get_metadata(&self, key_id: &str) -> Option<KeyMetadata> {
let inner = self.inner.read().await;
inner.metadata.get(key_id).cloned()
}
pub async fn clear_cached_keys(&self) {
let mut inner = self.inner.write().await;
let pinned_roots: std::collections::HashSet<_> =
inner.pinned_roots.iter().cloned().collect();
inner.keys.retain(|k, _| pinned_roots.contains(k));
inner.metadata.retain(|k, _| pinned_roots.contains(k));
inner.manifest_fetched_at = None;
inner.manifest_expires_at = None;
}
fn empty_inner() -> TrustStoreInner {
TrustStoreInner {
keys: HashMap::new(),
metadata: HashMap::new(),
pinned_roots: Vec::new(),
manifest_fetched_at: None,
manifest_expires_at: None,
}
}
}
impl Default for TrustStore {
fn default() -> Self {
Self::new()
}
}
fn decode_verifying_key(b64: &str) -> RegistryResult<VerifyingKey> {
use pkcs8::DecodePublicKey;
let bytes = BASE64.decode(b64).map_err(|e| RegistryError::Config {
message: format!("invalid base64 public key: {}", e),
})?;
VerifyingKey::from_public_key_der(&bytes).map_err(|e| RegistryError::Config {
message: format!("invalid SPKI public key: {}", e),
})
}
fn decode_public_key_bytes(b64: &str) -> RegistryResult<Vec<u8>> {
BASE64.decode(b64).map_err(|e| RegistryError::Config {
message: format!("invalid base64 public key: {}", e),
})
}
fn parse_pinned_roots_json_impl(raw: &str) -> RegistryResult<Vec<TrustedKey>> {
let roots: Vec<TrustedKey> = serde_json::from_str(raw).map_err(|e| RegistryError::Config {
message: format!("invalid production trust roots: {}", e),
})?;
if roots.is_empty() {
return Err(RegistryError::Config {
message: "production trust roots are empty".to_string(),
});
}
let mut seen = HashSet::new();
for root in &roots {
if root.algorithm != "Ed25519" {
return Err(RegistryError::Config {
message: format!(
"production trust root {} uses unsupported algorithm {}",
root.key_id, root.algorithm
),
});
}
if root.revoked {
return Err(RegistryError::Config {
message: format!("production trust root {} is revoked", root.key_id),
});
}
if !seen.insert(root.key_id.clone()) {
return Err(RegistryError::Config {
message: format!("duplicate production trust root {}", root.key_id),
});
}
}
Ok(roots)
}
fn load_production_roots_impl(raw: &str) -> RegistryResult<TrustStore> {
let roots = parse_pinned_roots_json_impl(raw)?;
let mut inner = TrustStore::empty_inner();
for root in &roots {
insert_pinned_key(&mut inner, root).map_err(|err| RegistryError::Config {
message: format!("invalid production trust root {}: {}", root.key_id, err),
})?;
}
Ok(TrustStore {
inner: Arc::new(RwLock::new(inner)),
})
}
fn insert_pinned_key(inner: &mut TrustStoreInner, key: &TrustedKey) -> RegistryResult<()> {
let verifying_key = decode_verifying_key(&key.public_key)?;
let computed_id = compute_key_id(&decode_public_key_bytes(&key.public_key)?);
if computed_id != key.key_id {
return Err(RegistryError::SignatureInvalid {
reason: format!(
"key_id mismatch: claimed {}, computed {}",
key.key_id, computed_id
),
});
}
inner.keys.insert(key.key_id.clone(), verifying_key);
inner.metadata.insert(
key.key_id.clone(),
KeyMetadata {
description: key.description.clone(),
added_at: key.added_at,
expires_at: key.expires_at,
revoked: false,
is_pinned: true,
},
);
if !inner.pinned_roots.contains(&key.key_id) {
inner.pinned_roots.push(key.key_id.clone());
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use ed25519_dalek::SigningKey;
use pkcs8::EncodePublicKey;
fn generate_trusted_key() -> (SigningKey, TrustedKey) {
let signing_key = SigningKey::generate(&mut rand::thread_rng());
let verifying_key = signing_key.verifying_key();
let spki_der = verifying_key.to_public_key_der().unwrap();
let public_key_b64 = BASE64.encode(spki_der.as_bytes());
let key_id = compute_key_id(spki_der.as_bytes());
let trusted = TrustedKey {
key_id,
algorithm: "Ed25519".to_string(),
public_key: public_key_b64,
description: Some("Test key".to_string()),
added_at: Some(Utc::now()),
expires_at: None,
revoked: false,
};
(signing_key, trusted)
}
#[tokio::test]
async fn test_empty_trust_store() {
let store = TrustStore::new();
let result = store.get_key_async("sha256:unknown").await;
assert!(matches!(result, Err(RegistryError::KeyNotTrusted { .. })));
}
#[tokio::test]
async fn test_with_production_roots_loads_embedded_roots() -> RegistryResult<()> {
let store = TrustStore::with_production_roots().await?;
let keys = store.list_keys().await;
assert_eq!(keys.len(), 1);
assert_eq!(
keys[0],
"sha256:3a64307d5655ba86fa3c95118ed8fe9665ef6bd37c752ca93f3bbe8f16e83a7f"
);
let meta = store
.get_metadata(&keys[0])
.await
.ok_or_else(|| RegistryError::Config {
message: "embedded production root metadata missing".to_string(),
})?;
assert!(meta.is_pinned);
assert!(!meta.revoked);
Ok(())
}
#[test]
fn test_parse_pinned_roots_json_rejects_empty_rootset() {
assert!(matches!(
parse_pinned_roots_json_impl("[]"),
Err(RegistryError::Config { .. })
));
}
#[test]
fn test_parse_pinned_roots_json_rejects_duplicate_key_ids() {
let duplicate = r#"[
{
"key_id": "sha256:dup",
"algorithm": "Ed25519",
"public_key": "MCowBQYDK2VwAyEAykCN7Cf9EQAB4UPonG5AtKfTVny0H4xaKpPI6wIGBwE=",
"revoked": false
},
{
"key_id": "sha256:dup",
"algorithm": "Ed25519",
"public_key": "MCowBQYDK2VwAyEAykCN7Cf9EQAB4UPonG5AtKfTVny0H4xaKpPI6wIGBwE=",
"revoked": false
}
]"#;
assert!(matches!(
parse_pinned_roots_json_impl(duplicate),
Err(RegistryError::Config { .. })
));
}
#[test]
fn test_load_production_roots_maps_key_mismatch_to_config() {
let mismatched = r#"[
{
"key_id": "sha256:not-the-real-key-id",
"algorithm": "Ed25519",
"public_key": "MCowBQYDK2VwAyEAykCN7Cf9EQAB4UPonG5AtKfTVny0H4xaKpPI6wIGBwE=",
"revoked": false
}
]"#;
let err = load_production_roots_impl(mismatched).unwrap_err();
assert!(matches!(err, RegistryError::Config { .. }));
assert!(err
.to_string()
.contains("invalid production trust root sha256:not-the-real-key-id"));
}
#[tokio::test]
async fn test_add_pinned_key() {
let store = TrustStore::new();
let (_signing_key, trusted) = generate_trusted_key();
store.add_pinned_key(&trusted).await.unwrap();
let key = store.get_key_async(&trusted.key_id).await.unwrap();
assert_eq!(key.as_bytes().len(), 32);
let meta = store.get_metadata(&trusted.key_id).await.unwrap();
assert!(meta.is_pinned);
assert!(!meta.revoked);
}
#[tokio::test]
async fn test_add_from_manifest() {
let store = TrustStore::new();
let (_, trusted1) = generate_trusted_key();
let (_, trusted2) = generate_trusted_key();
let manifest = KeysManifest {
version: 1,
keys: vec![trusted1.clone(), trusted2.clone()],
expires_at: Some(Utc::now() + chrono::Duration::hours(24)),
};
store.add_from_manifest(&manifest).await.unwrap();
assert!(store.is_trusted(&trusted1.key_id).await);
assert!(store.is_trusted(&trusted2.key_id).await);
}
#[tokio::test]
async fn test_revoked_key_in_manifest() {
let store = TrustStore::new();
let (_, mut trusted) = generate_trusted_key();
trusted.revoked = true;
let manifest = KeysManifest {
version: 1,
keys: vec![trusted.clone()],
expires_at: None,
};
store.add_from_manifest(&manifest).await.unwrap();
assert!(!store.is_trusted(&trusted.key_id).await);
}
#[tokio::test]
async fn test_expired_key_in_manifest() {
let store = TrustStore::new();
let (_, mut trusted) = generate_trusted_key();
trusted.expires_at = Some(Utc::now() - chrono::Duration::hours(1));
let manifest = KeysManifest {
version: 1,
keys: vec![trusted.clone()],
expires_at: None,
};
store.add_from_manifest(&manifest).await.unwrap();
assert!(!store.is_trusted(&trusted.key_id).await);
}
#[tokio::test]
async fn test_pinned_key_not_overwritten() {
let store = TrustStore::new();
let (_, trusted) = generate_trusted_key();
store.add_pinned_key(&trusted).await.unwrap();
let mut revoked = trusted.clone();
revoked.revoked = true;
let manifest = KeysManifest {
version: 1,
keys: vec![revoked],
expires_at: None,
};
store.add_from_manifest(&manifest).await.unwrap();
assert!(store.is_trusted(&trusted.key_id).await);
let meta = store.get_metadata(&trusted.key_id).await.unwrap();
assert!(meta.is_pinned);
assert!(!meta.revoked);
}
#[tokio::test]
async fn test_needs_refresh() {
let store = TrustStore::new();
assert!(store.needs_refresh().await);
let manifest = KeysManifest {
version: 1,
keys: vec![],
expires_at: Some(Utc::now() + chrono::Duration::hours(24)),
};
store.add_from_manifest(&manifest).await.unwrap();
assert!(!store.needs_refresh().await);
}
#[tokio::test]
async fn test_clear_cached_keys() {
let store = TrustStore::new();
let (_, pinned) = generate_trusted_key();
let (_, cached) = generate_trusted_key();
store.add_pinned_key(&pinned).await.unwrap();
let manifest = KeysManifest {
version: 1,
keys: vec![cached.clone()],
expires_at: None,
};
store.add_from_manifest(&manifest).await.unwrap();
assert!(store.is_trusted(&pinned.key_id).await);
assert!(store.is_trusted(&cached.key_id).await);
store.clear_cached_keys().await;
assert!(store.is_trusted(&pinned.key_id).await);
assert!(!store.is_trusted(&cached.key_id).await);
}
#[tokio::test]
async fn test_list_keys() {
let store = TrustStore::new();
let (_, key1) = generate_trusted_key();
let (_, key2) = generate_trusted_key();
store.add_pinned_key(&key1).await.unwrap();
store.add_pinned_key(&key2).await.unwrap();
let keys = store.list_keys().await;
assert_eq!(keys.len(), 2);
assert!(keys.contains(&key1.key_id));
assert!(keys.contains(&key2.key_id));
}
#[tokio::test]
async fn test_key_id_mismatch_rejected() {
let store = TrustStore::new();
let (_, mut trusted) = generate_trusted_key();
trusted.key_id =
"sha256:0000000000000000000000000000000000000000000000000000000000000000".to_string();
let result = store.add_pinned_key(&trusted).await;
assert!(matches!(
result,
Err(RegistryError::SignatureInvalid { .. })
));
}
#[tokio::test]
async fn test_trust_rotation_new_key_via_manifest() {
let store = TrustStore::new();
let (_, key_a) = generate_trusted_key();
let manifest_a = KeysManifest {
version: 1,
keys: vec![key_a.clone()],
expires_at: Some(Utc::now() + chrono::Duration::hours(24)),
};
store.add_from_manifest(&manifest_a).await.unwrap();
assert!(store.is_trusted(&key_a.key_id).await);
let (_, key_b) = generate_trusted_key();
let manifest_b = KeysManifest {
version: 1,
keys: vec![key_a.clone(), key_b.clone()], expires_at: Some(Utc::now() + chrono::Duration::hours(24)),
};
store.add_from_manifest(&manifest_b).await.unwrap();
assert!(store.is_trusted(&key_a.key_id).await);
assert!(store.is_trusted(&key_b.key_id).await);
}
#[tokio::test]
async fn test_trust_rotation_revoke_old_key() {
let store = TrustStore::new();
let (_, key_a) = generate_trusted_key();
let manifest_v1 = KeysManifest {
version: 1,
keys: vec![key_a.clone()],
expires_at: Some(Utc::now() + chrono::Duration::hours(24)),
};
store.add_from_manifest(&manifest_v1).await.unwrap();
assert!(store.is_trusted(&key_a.key_id).await);
let mut key_a_revoked = key_a.clone();
key_a_revoked.revoked = true;
let (_, key_b) = generate_trusted_key();
let manifest_v2 = KeysManifest {
version: 1,
keys: vec![key_a_revoked, key_b.clone()],
expires_at: Some(Utc::now() + chrono::Duration::hours(24)),
};
store.add_from_manifest(&manifest_v2).await.unwrap();
assert!(!store.is_trusted(&key_a.key_id).await);
assert!(store.is_trusted(&key_b.key_id).await);
}
#[tokio::test]
async fn test_trust_rotation_pinned_root_survives_revocation() {
let store = TrustStore::new();
let (_, pinned_root) = generate_trusted_key();
store.add_pinned_key(&pinned_root).await.unwrap();
assert!(store.is_trusted(&pinned_root.key_id).await);
let mut revoked_root = pinned_root.clone();
revoked_root.revoked = true;
let manifest = KeysManifest {
version: 1,
keys: vec![revoked_root],
expires_at: Some(Utc::now() + chrono::Duration::hours(24)),
};
store.add_from_manifest(&manifest).await.unwrap();
assert!(store.is_trusted(&pinned_root.key_id).await);
let meta = store.get_metadata(&pinned_root.key_id).await.unwrap();
assert!(meta.is_pinned);
assert!(!meta.revoked);
}
#[tokio::test]
async fn test_trust_rotation_expired_key_not_added() {
let store = TrustStore::new();
let (_, mut expired_key) = generate_trusted_key();
expired_key.expires_at = Some(Utc::now() - chrono::Duration::hours(1));
let manifest = KeysManifest {
version: 1,
keys: vec![expired_key.clone()],
expires_at: Some(Utc::now() + chrono::Duration::hours(24)),
};
store.add_from_manifest(&manifest).await.unwrap();
assert!(!store.is_trusted(&expired_key.key_id).await);
}
#[tokio::test]
async fn test_trust_rotation_key_expires_after_added() {
let store = TrustStore::new();
let (_, mut soon_to_expire) = generate_trusted_key();
soon_to_expire.expires_at = Some(Utc::now() - chrono::Duration::seconds(1));
let manifest = KeysManifest {
version: 1,
keys: vec![soon_to_expire.clone()],
expires_at: None,
};
store.add_from_manifest(&manifest).await.unwrap();
assert!(!store.is_trusted(&soon_to_expire.key_id).await);
}
}