agentguard_core/
auth_keys.rs

1//! Key registry for JWS (JWT/DPoP/delegation) signature verification.
2//!
3//! Supports rotation: a new key with the same `kid` can be added while the
4//! old one remains valid for a configurable grace period.
5//!
6//! Lives in `agentguard-core` (not `agentguard-auth`) because the registry
7//! itself has no auth-specific dependencies — `Algorithm` and `KeyMaterial`
8//! are the JOSE primitives, and the same registry is shared by the JWT
9//! validator (`agentguard-auth`), the DPoP verifier (which doesn't
10//! currently use a registry but could), and the delegation verifier
11//! (`agentguard-core::delegation`).
12
13use serde::{Deserialize, Serialize};
14use std::collections::HashMap;
15use std::time::{Duration, Instant};
16
17/// JWS signing algorithm (RFC 7518, RFC 8725 §3.1).
18///
19/// All four variants parse out of header strings, but only `EdDSA`
20/// is implemented by the signature verifier today. `JwtConfig::with_algorithms`
21/// rejects non-`EdDSA` values at construction time so a misconfigured
22/// operator sees the failure before the first token arrives. RS256
23/// and ES256 are reserved for a future release.
24#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
25#[serde(rename_all = "UPPERCASE")]
26pub enum Algorithm {
27    /// HMAC-SHA256 (symmetric). Avoid for asymmetric protocols.
28    HS256,
29    /// RSASSA-PKCS1-v1_5 with SHA-256.
30    RS256,
31    /// ECDSA with P-256 and SHA-256.
32    ES256,
33    /// Ed25519.
34    EdDSA,
35}
36
37impl std::str::FromStr for Algorithm {
38    type Err = ParseAlgorithmError;
39    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
40        parse_alg(s).ok_or_else(|| ParseAlgorithmError(s.to_string()))
41    }
42}
43
44impl std::fmt::Display for Algorithm {
45    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
46        f.write_str(match self {
47            Algorithm::HS256 => "HS256",
48            Algorithm::RS256 => "RS256",
49            Algorithm::ES256 => "ES256",
50            Algorithm::EdDSA => "EdDSA",
51        })
52    }
53}
54
55/// Error from parsing a [`Algorithm`] string.
56#[derive(Debug, thiserror::Error)]
57#[error("unknown JOSE algorithm: {0}")]
58pub struct ParseAlgorithmError(pub String);
59
60/// Key material — abstracted so the registry can hold heterogeneous keys.
61#[derive(Debug, Clone)]
62pub enum KeyMaterial {
63    /// Raw HMAC secret bytes (32+ bytes for HS256).
64    Hmac(Vec<u8>),
65    /// PEM-encoded RSA public key.
66    Rsa(Vec<u8>),
67    /// PEM-encoded EC P-256 public key.
68    Ecdsa(Vec<u8>),
69    /// Raw 32-byte Ed25519 public key.
70    Ed25519(Vec<u8>),
71}
72
73#[derive(Debug, Clone)]
74struct KeyEntry {
75    #[allow(dead_code)]
76    kid: String,
77    alg: Algorithm,
78    key: KeyMaterial,
79    grace_expires_at: Option<Instant>,
80}
81
82/// Thread-safe registry of trusted verification keys, keyed by `kid`.
83///
84/// The registry supports key rotation via [`Self::rotate`]: when a new
85/// key is added under the same `kid`, the old key is preserved for the
86/// grace period so tokens signed by the old key continue to verify
87/// during the cutover.
88///
89/// Bounded by [`DEFAULT_KEY_KID_CAP`] distinct kids. When the cap is
90/// reached, [`Self::add`] and [`Self::rotate`] evict the oldest kid before
91/// inserting a new one. This prevents a misbehaving IdP with thousands of
92/// distinct kids from growing the registry without bound.
93#[derive(Debug)]
94pub struct KeyRegistry {
95    inner: parking_lot::RwLock<HashMap<String, Vec<KeyEntry>>>,
96    /// Insertion order for FIFO eviction. New kids are pushed to the
97    /// back; both `add` and `rotate` evict the front at capacity. `VecDeque`
98    /// keeps eviction O(1).
99    order: parking_lot::Mutex<std::collections::VecDeque<String>>,
100    cap: usize,
101}
102
103/// Default cap on the number of distinct `kid`s a [`KeyRegistry`]
104/// will hold. Mirrors the cap enforced inside
105/// `JwtValidator::refresh_jwks` so the two cannot drift.
106pub const DEFAULT_KEY_KID_CAP: usize = 64;
107
108impl Default for KeyRegistry {
109    fn default() -> Self {
110        Self::with_cap(DEFAULT_KEY_KID_CAP)
111    }
112}
113
114impl KeyRegistry {
115    /// Construct an empty registry.
116    pub fn new() -> Self {
117        Self::default()
118    }
119
120    /// Construct a registry with a custom kid cap (used in tests).
121    pub fn with_cap(cap: usize) -> Self {
122        Self {
123            inner: parking_lot::RwLock::new(HashMap::new()),
124            order: parking_lot::Mutex::new(std::collections::VecDeque::new()),
125            cap: cap.max(1),
126        }
127    }
128
129    pub fn cap(&self) -> usize {
130        self.cap
131    }
132
133    fn reserve_kid_slot(&self, entries: &mut HashMap<String, Vec<KeyEntry>>, kid: &str) {
134        if entries.contains_key(kid) {
135            return;
136        }
137
138        let mut order = self.order.lock();
139        while entries.len() >= self.cap {
140            let Some(oldest) = order.pop_front() else {
141                break;
142            };
143            entries.remove(&oldest);
144            tracing::warn!(
145                kid = %oldest,
146                cap = self.cap,
147                "key registry at cap; evicted oldest kid"
148            );
149        }
150        order.push_back(kid.to_owned());
151    }
152
153    /// Register a key. Replaces any existing active key with the same
154    /// `kid` and `alg` (grace-window keys from a rotation are preserved).
155    /// If the registry is at its cap, the oldest kid is evicted first.
156    pub fn add(&self, kid: impl Into<String>, alg: Algorithm, key: KeyMaterial) {
157        let kid = kid.into();
158        let mut guard = self.inner.write();
159        self.reserve_kid_slot(&mut guard, &kid);
160        // Drop a stale entry (if any) so a re-add counts as an update,
161        // not a new insertion, when computing capacity.
162        let entries = guard.entry(kid.clone()).or_default();
163        entries.retain(|e| !(e.alg == alg && e.grace_expires_at.is_none()));
164        entries.push(KeyEntry {
165            kid,
166            alg,
167            key,
168            grace_expires_at: None,
169        });
170    }
171
172    /// Register a new key under the same `kid`, marking the previous key as
173    /// in grace for the given duration.
174    pub fn rotate(
175        &self,
176        kid: impl Into<String>,
177        alg: Algorithm,
178        key: KeyMaterial,
179        grace: Duration,
180    ) {
181        let kid = kid.into();
182        let mut guard = self.inner.write();
183        self.reserve_kid_slot(&mut guard, &kid);
184        let entries = guard.entry(kid.clone()).or_default();
185        for entry in entries.iter_mut() {
186            entry.grace_expires_at = Some(Instant::now() + grace);
187        }
188        entries.push(KeyEntry {
189            kid,
190            alg,
191            key,
192            grace_expires_at: None,
193        });
194    }
195
196    /// Look up a key by `kid` and `alg`. Returns all matching keys that have
197    /// not yet exhausted their grace period. Returns an empty `Vec` if no
198    /// matching key is found (caller decides the error mapping).
199    pub fn get(&self, kid: &str, alg: Algorithm) -> Vec<KeyMaterial> {
200        let guard = self.inner.read();
201        let now = Instant::now();
202        guard
203            .get(kid)
204            .map(|entries| {
205                entries
206                    .iter()
207                    .filter(|e| e.alg == alg)
208                    .filter(|e| e.grace_expires_at.map(|t| t > now).unwrap_or(true))
209                    .map(|e| e.key.clone())
210                    .collect()
211            })
212            .unwrap_or_default()
213    }
214
215    /// Returns true if a key with this `kid` is registered (any algorithm).
216    pub fn contains(&self, kid: &str) -> bool {
217        self.inner.read().contains_key(kid)
218    }
219
220    /// Number of distinct `kid`s registered.
221    pub fn kid_count(&self) -> usize {
222        self.inner.read().len()
223    }
224}
225
226/// Parse a JOSE algorithm name into our [`Algorithm`] enum.
227pub fn parse_alg(s: &str) -> Option<Algorithm> {
228    match s {
229        "HS256" => Some(Algorithm::HS256),
230        "RS256" => Some(Algorithm::RS256),
231        "ES256" => Some(Algorithm::ES256),
232        "EdDSA" => Some(Algorithm::EdDSA),
233        _ => None,
234    }
235}
236
237#[cfg(test)]
238mod tests {
239    use super::*;
240
241    #[test]
242    fn add_and_get_returns_key() {
243        let r = KeyRegistry::new();
244        r.add(
245            "kid1",
246            Algorithm::EdDSA,
247            KeyMaterial::Ed25519(vec![0u8; 32]),
248        );
249        let keys = r.get("kid1", Algorithm::EdDSA);
250        assert_eq!(keys.len(), 1);
251    }
252
253    #[test]
254    fn rotate_preserves_old_during_grace() {
255        let r = KeyRegistry::new();
256        r.add(
257            "kid1",
258            Algorithm::EdDSA,
259            KeyMaterial::Ed25519(vec![1u8; 32]),
260        );
261        r.rotate(
262            "kid1",
263            Algorithm::EdDSA,
264            KeyMaterial::Ed25519(vec![2u8; 32]),
265            Duration::from_secs(60),
266        );
267        let keys = r.get("kid1", Algorithm::EdDSA);
268        assert_eq!(keys.len(), 2, "both old and new should verify during grace");
269    }
270
271    #[test]
272    fn unknown_kid_returns_empty() {
273        let r = KeyRegistry::new();
274        assert!(r.get("missing", Algorithm::EdDSA).is_empty());
275        assert!(!r.contains("missing"));
276    }
277
278    #[test]
279    fn contains_returns_true_after_add() {
280        let r = KeyRegistry::new();
281        r.add(
282            "kid1",
283            Algorithm::EdDSA,
284            KeyMaterial::Ed25519(vec![0u8; 32]),
285        );
286        assert!(r.contains("kid1"));
287    }
288
289    #[test]
290    fn add_replaces_active_key() {
291        let r = KeyRegistry::new();
292        r.add(
293            "kid1",
294            Algorithm::EdDSA,
295            KeyMaterial::Ed25519(vec![1u8; 32]),
296        );
297        r.add(
298            "kid1",
299            Algorithm::EdDSA,
300            KeyMaterial::Ed25519(vec![2u8; 32]),
301        );
302        // Second add replaces the first; we expect only the new key
303        // (no grace window, so the old one is gone).
304        let keys = r.get("kid1", Algorithm::EdDSA);
305        assert_eq!(keys.len(), 1);
306    }
307
308    #[test]
309    fn cap_evicts_oldest_kid() {
310        let r = KeyRegistry::with_cap(3);
311        for i in 0..5 {
312            r.add(
313                format!("kid{i}"),
314                Algorithm::EdDSA,
315                KeyMaterial::Ed25519(vec![i as u8; 32]),
316            );
317        }
318        // Only the last 3 kids are retained.
319        assert_eq!(r.kid_count(), 3);
320        assert!(!r.contains("kid0"));
321        assert!(!r.contains("kid1"));
322        assert!(r.contains("kid2"));
323        assert!(r.contains("kid3"));
324        assert!(r.contains("kid4"));
325    }
326
327    #[test]
328    fn adding_existing_kid_does_not_evict() {
329        let r = KeyRegistry::with_cap(4);
330        r.add("kid1", Algorithm::EdDSA, KeyMaterial::Ed25519(vec![1; 32]));
331        // Re-adding the same kid updates in place; no eviction.
332        r.add("kid1", Algorithm::EdDSA, KeyMaterial::Ed25519(vec![2; 32]));
333        assert_eq!(r.kid_count(), 1);
334    }
335
336    #[test]
337    fn rotate_new_kids_respects_cap_and_fifo_eviction_order() {
338        let r = KeyRegistry::with_cap(2);
339        r.add("kid1", Algorithm::EdDSA, KeyMaterial::Ed25519(vec![1; 32]));
340        r.rotate(
341            "kid2",
342            Algorithm::EdDSA,
343            KeyMaterial::Ed25519(vec![2; 32]),
344            Duration::from_secs(60),
345        );
346
347        assert_eq!(r.kid_count(), 2);
348        assert!(r.contains("kid1"));
349        assert!(r.contains("kid2"));
350
351        r.rotate(
352            "kid3",
353            Algorithm::EdDSA,
354            KeyMaterial::Ed25519(vec![3; 32]),
355            Duration::from_secs(60),
356        );
357
358        assert_eq!(r.kid_count(), 2);
359        assert!(!r.contains("kid1"));
360        assert!(r.contains("kid2"));
361        assert!(r.contains("kid3"));
362    }
363
364    #[test]
365    fn rotating_an_existing_kid_preserves_grace_keys_without_consuming_capacity() {
366        let r = KeyRegistry::with_cap(1);
367        r.add("kid1", Algorithm::EdDSA, KeyMaterial::Ed25519(vec![1; 32]));
368        r.rotate(
369            "kid1",
370            Algorithm::EdDSA,
371            KeyMaterial::Ed25519(vec![2; 32]),
372            Duration::from_secs(60),
373        );
374
375        assert_eq!(r.kid_count(), 1);
376        assert_eq!(r.get("kid1", Algorithm::EdDSA).len(), 2);
377    }
378}
agentguard_core/auth_keys.rs

agentguard_core/
auth_keys.rs
