superfluid_daemon/
keypolicy.rs

1//! Per-API-key serving policy (`--key-policy <file.json>`).
2
3use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
4use std::sync::Arc;
5
6use serde::Deserialize;
7
8use crate::qos;
9use crate::ratelimit::TokenBucket;
10
11#[derive(Deserialize)]
12#[serde(deny_unknown_fields)]
13struct FileTable {
14    keys: Vec<FileKey>,
15}
16
17#[derive(Deserialize)]
18#[serde(deny_unknown_fields)]
19struct FileKey {
20    name: String,
21    #[serde(default)]
22    key: Option<String>,
23    #[serde(default)]
24    key_env: Option<String>,
25    #[serde(default)]
26    class: Option<String>,
27    #[serde(default)]
28    max_class: Option<String>,
29    #[serde(default)]
30    batch_invariant: bool,
31    #[serde(default)]
32    rate_limit_rpm: u32,
33    #[serde(default)]
34    max_concurrent: u32,
35    #[serde(default)]
36    admin: bool,
37}
38
39#[derive(Debug, Clone, Copy, PartialEq, Eq)]
40pub enum Reject {
41    RateLimit,
42    Concurrency,
43    Qos,
44}
45
46#[derive(Debug, Default)]
47pub struct KeyStats {
48    pub requests: AtomicU64,
49    pub rejected_rate: AtomicU64,
50    pub rejected_concurrency: AtomicU64,
51    pub rejected_qos: AtomicU64,
52    pub rejected_admin: AtomicU64,
53}
54
55pub struct KeyPolicy {
56    pub name: String,
57    secret: String,
58    digest: [u8; 32],
59    pub class: u8,
60    pub max_class: u8,
61    pub allow_batch_invariant: bool,
62    pub rate_limit_rpm: u32,
63    pub max_concurrent: u32,
64    pub admin: bool,
65    bucket: TokenBucket,
66    in_flight: Arc<AtomicU32>,
67    pub stats: KeyStats,
68}
69
70impl std::fmt::Debug for KeyPolicy {
71    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
72        f.debug_struct("KeyPolicy")
73            .field("name", &self.name)
74            .field("class", &self.class)
75            .field("max_class", &self.max_class)
76            .field("allow_batch_invariant", &self.allow_batch_invariant)
77            .field("rate_limit_rpm", &self.rate_limit_rpm)
78            .field("max_concurrent", &self.max_concurrent)
79            .finish_non_exhaustive()
80    }
81}
82
83#[derive(Debug)]
84pub struct InFlight(Option<Arc<AtomicU32>>);
85
86impl Drop for InFlight {
87    fn drop(&mut self) {
88        if let Some(n) = self.0.take() {
89            n.fetch_sub(1, Ordering::AcqRel);
90        }
91    }
92}
93
94impl KeyPolicy {
95    pub fn admit(&self, takes_slot: bool, spends_rate: bool) -> Result<InFlight, Reject> {
96        self.stats.requests.fetch_add(1, Ordering::Relaxed);
97        let slot = if self.max_concurrent == 0 || !takes_slot {
98            InFlight(None)
99        } else {
100            let cap = self.max_concurrent;
101            let got = self
102                .in_flight
103                .try_update(Ordering::AcqRel, Ordering::Acquire, |n| (n < cap).then_some(n + 1));
104            if got.is_err() {
105                self.stats.rejected_concurrency.fetch_add(1, Ordering::Relaxed);
106                return Err(Reject::Concurrency);
107            }
108            InFlight(Some(Arc::clone(&self.in_flight)))
109        };
110        if spends_rate && !self.bucket.allow() {
111            self.stats.rejected_rate.fetch_add(1, Ordering::Relaxed);
112            return Err(Reject::RateLimit);
113        }
114        Ok(slot)
115    }
116
117    pub fn pace(&self, stop: impl Fn() -> bool) -> bool {
118        loop {
119            if stop() {
120                return false;
121            }
122            if self.bucket.allow() {
123                self.stats.requests.fetch_add(1, Ordering::Relaxed);
124                return true;
125            }
126            std::thread::sleep(self.bucket.next_token_in().min(std::time::Duration::from_millis(500)));
127        }
128    }
129
130    pub fn charge_rate(&self) -> bool {
131        if self.bucket.allow() {
132            return true;
133        }
134        self.stats.rejected_rate.fetch_add(1, Ordering::Relaxed);
135        false
136    }
137
138    pub fn in_flight(&self) -> u32 {
139        self.in_flight.load(Ordering::Acquire)
140    }
141
142    pub fn note_qos_reject(&self) {
143        self.stats.rejected_qos.fetch_add(1, Ordering::Relaxed);
144    }
145
146    pub fn note_admin_reject(&self) {
147        self.stats.rejected_admin.fetch_add(1, Ordering::Relaxed);
148    }
149}
150
151pub(crate) fn sha256(s: &str) -> [u8; 32] {
152    use sha2::Digest;
153    sha2::Sha256::digest(s.as_bytes()).into()
154}
155
156#[derive(Clone, Debug)]
157pub struct ClientKey(pub Arc<KeyPolicy>, pub Option<Arc<InFlight>>);
158
159#[derive(Debug)]
160pub struct KeyTable {
161    keys: Vec<Arc<KeyPolicy>>,
162    by_digest: std::collections::HashMap<[u8; 32], usize>,
163}
164
165fn parse_class(field: &str, name: &str, v: &str, redact: &dyn Fn(&str) -> String) -> Result<u8, String> {
166    qos::parse(v).ok_or_else(|| {
167        format!(
168            "key {name:?}: {field} {} is not a class (expected one of {})",
169            redact(v),
170            qos::NAMES.join(", ")
171        )
172    })
173}
174
175impl KeyTable {
176    pub fn load(path: &std::path::Path, api_key: Option<&str>) -> Result<KeyTable, String> {
177        let text = std::fs::read_to_string(path).map_err(|e| format!("{}: {e}", path.display()))?;
178        Self::parse_with_global(&text, |var| std::env::var(var).ok(), api_key)
179            .and_then(|t| t.check_global_key(api_key).map(|()| t))
180            .map_err(|e| format!("{}: {e}", path.display()))
181    }
182
183    pub fn parse(text: &str, env: impl Fn(&str) -> Option<String>) -> Result<KeyTable, String> {
184        Self::parse_with_global(text, env, None)
185    }
186
187    pub fn parse_with_global(
188        text: &str,
189        env: impl Fn(&str) -> Option<String>,
190        api_key: Option<&str>,
191    ) -> Result<KeyTable, String> {
192        let file: FileTable = serde_json::from_str(text).map_err(|e| e.to_string())?;
193        if file.keys.is_empty() {
194            return Err("no keys: a policy file with no keys would lock every client out".into());
195        }
196        let secrets: Vec<String> = file
197            .keys
198            .iter()
199            .flat_map(|k| k.key.clone().into_iter().chain(k.key_env.as_deref().and_then(&env)))
200            .collect();
201        let redact = |v: &str| -> String {
202            if secrets.iter().any(|s| s == v) || api_key.is_some_and(|g| g == v) {
203                "(a value equal to a credential, not shown)".to_string()
204            } else {
205                format!("{v:?}")
206            }
207        };
208        let mut keys: Vec<Arc<KeyPolicy>> = Vec::with_capacity(file.keys.len());
209        for (i, k) in file.keys.into_iter().enumerate() {
210            let name = k.name.trim().to_string();
211            if secrets.iter().any(|s| *s == name || *s == k.name) {
212                return Err(format!(
213                    "keys[{i}]: the name equals a key's secret; names are published in metrics and logs"
214                ));
215            }
216            if api_key.is_some_and(|g| g == name || g == k.name) {
217                return Err(format!("keys[{i}]: the name equals --api-key; names are published in metrics and logs"));
218            }
219            if name.is_empty()
220                || !name.chars().all(|c| c.is_ascii_alphanumeric() || matches!(c, '-' | '_' | '.'))
221            {
222                return Err(format!("key name {:?} must be non-empty [A-Za-z0-9._-]", k.name));
223            }
224            let secret = match (k.key, k.key_env) {
225                (Some(s), None) => s,
226                (None, Some(var)) => {
227                    env(&var).ok_or_else(|| format!("key {name:?}: environment variable {} is not set", redact(&var)))?
228                }
229                _ => return Err(format!("key {name:?}: give exactly one of \"key\" or \"key_env\"")),
230            };
231            if secret.trim().is_empty() {
232                return Err(format!("key {name:?}: the key is empty"));
233            }
234            if !secret.bytes().all(|b| b.is_ascii_graphic()) {
235                return Err(format!(
236                    "key {name:?}: the key must be visible ASCII (no spaces, control or non-ASCII characters)"
237                ));
238            }
239            let class = match &k.class {
240                Some(c) => parse_class("class", &name, c, &redact)?,
241                None => qos::FOREGROUND_AGENT,
242            };
243            let max_class = match &k.max_class {
244                Some(c) => parse_class("max_class", &name, c, &redact)?,
245                None => class,
246            };
247            if max_class > class {
248                return Err(format!(
249                    "key {name:?}: max_class {} is a lower priority than its default class {}; \
250                     max_class is the BEST class the header may ask for",
251                    qos::NAMES[max_class as usize],
252                    qos::NAMES[class as usize]
253                ));
254            }
255            if keys.iter().any(|o| o.name == name) {
256                return Err(format!("key name {name:?} is listed twice"));
257            }
258            if let Some(o) = keys.iter().find(|o| o.secret == secret) {
259                return Err(format!("keys {:?} and {name:?} share a secret", o.name));
260            }
261            let digest = sha256(&secret);
262            keys.push(Arc::new(KeyPolicy {
263                name,
264                secret,
265                digest,
266                class,
267                max_class,
268                allow_batch_invariant: k.batch_invariant,
269                rate_limit_rpm: k.rate_limit_rpm,
270                max_concurrent: k.max_concurrent,
271                admin: k.admin,
272                bucket: TokenBucket::new(k.rate_limit_rpm),
273                in_flight: Arc::new(AtomicU32::new(0)),
274                stats: KeyStats::default(),
275            }));
276        }
277        let by_digest = keys.iter().enumerate().map(|(i, k)| (k.digest, i)).collect();
278        Ok(KeyTable { keys, by_digest })
279    }
280
281    pub fn check_global_key(&self, api_key: Option<&str>) -> Result<(), String> {
282        if let Some(k) = api_key.and_then(|g| self.keys.iter().find(|k| k.secret == g)) {
283            return Err(format!("key {:?} has the same secret as --api-key", k.name));
284        }
285        if let Some(i) = api_key.and_then(|g| self.keys.iter().position(|k| k.name == g)) {
286            return Err(format!("keys[{i}]: the name equals --api-key; names are published in metrics and logs"));
287        }
288        Ok(())
289    }
290
291    pub fn lookup(&self, presented: &str) -> Option<Arc<KeyPolicy>> {
292        self.by_digest.get(&sha256(presented)).map(|&i| Arc::clone(&self.keys[i]))
293    }
294
295    pub fn len(&self) -> usize {
296        self.keys.len()
297    }
298
299    pub fn is_empty(&self) -> bool {
300        self.keys.is_empty()
301    }
302
303    pub fn keys(&self) -> impl Iterator<Item = &Arc<KeyPolicy>> {
304        self.keys.iter()
305    }
306
307    pub fn metrics_text(&self) -> String {
308        self.metrics_text_for(None)
309    }
310
311    pub fn metrics_text_for(&self, only: Option<&KeyPolicy>) -> String {
312        let g = |a: &AtomicU64| a.load(Ordering::Relaxed);
313        let shown: Vec<&Arc<KeyPolicy>> =
314            self.keys.iter().filter(|k| only.is_none_or(|o| std::ptr::eq(&***k, o))).collect();
315        let mut out = String::new();
316        out.push_str("# HELP superfluid_key_requests_total HTTP requests authenticated by a --key-policy key.\n# TYPE superfluid_key_requests_total counter\n");
317        for k in &shown {
318            out.push_str(&format!("superfluid_key_requests_total{{key=\"{}\"}} {}\n", k.name, g(&k.stats.requests)));
319        }
320        out.push_str("# HELP superfluid_key_rejected_total Keyed requests refused at the boundary, by reason.\n# TYPE superfluid_key_rejected_total counter\n");
321        for k in &shown {
322            for (reason, n) in [
323                ("rate_limit", &k.stats.rejected_rate),
324                ("concurrency", &k.stats.rejected_concurrency),
325                ("qos", &k.stats.rejected_qos),
326                ("admin", &k.stats.rejected_admin),
327            ] {
328                out.push_str(&format!(
329                    "superfluid_key_rejected_total{{key=\"{}\",reason=\"{reason}\"}} {}\n",
330                    k.name,
331                    g(n)
332                ));
333            }
334        }
335        out.push_str("# HELP superfluid_key_in_flight Requests holding one of a key's max_concurrent slots.\n# TYPE superfluid_key_in_flight gauge\n");
336        for k in &shown {
337            out.push_str(&format!("superfluid_key_in_flight{{key=\"{}\"}} {}\n", k.name, k.in_flight()));
338        }
339        out
340    }
341}
342
343#[cfg(test)]
344mod tests {
345    use super::*;
346
347    fn no_env(_: &str) -> Option<String> {
348        None
349    }
350
351    #[test]
352    fn parses_defaults_and_env_keys() {
353        let t = KeyTable::parse(
354            r#"{"keys":[
355                {"name":"ui","key":"sk-ui","class":"interactive"},
356                {"name":"bg","key_env":"BG_KEY"}
357            ]}"#,
358            |v| (v == "BG_KEY").then(|| "sk-bg".to_string()),
359        )
360        .unwrap();
361        let ui = t.lookup("sk-ui").unwrap();
362        assert_eq!((ui.class, ui.max_class), (qos::INTERACTIVE_CHAT, qos::INTERACTIVE_CHAT));
363        let bg = t.lookup("sk-bg").unwrap();
364        assert_eq!(bg.name, "bg");
365        assert_eq!((bg.class, bg.max_class), (qos::FOREGROUND_AGENT, qos::FOREGROUND_AGENT), "defaults: agent, header may only demote");
366        assert!(!bg.allow_batch_invariant);
367        assert!(t.lookup("sk-nope").is_none());
368        assert!(t.lookup("sk-u").is_none(), "a prefix is not a match");
369    }
370
371    #[test]
372    fn refuses_bad_tables() {
373        let err = |s: &str| KeyTable::parse(s, no_env).unwrap_err();
374        assert!(err(r#"{"keys":[]}"#).contains("no keys"));
375        assert!(err(r#"{"keys":[{"name":"a","key":"k","clas":"agent"}]}"#).contains("unknown field"));
376        assert!(err(r#"{"keys":[{"name":"a"}]}"#).contains("exactly one"));
377        assert!(err(r#"{"keys":[{"name":"a","key":"k","key_env":"K"}]}"#).contains("exactly one"));
378        assert!(err(r#"{"keys":[{"name":"a","key_env":"K"}]}"#).contains("not set"));
379        assert!(err(r#"{"keys":[{"name":"a","key":"  "}]}"#).contains("empty"));
380        assert!(err(r#"{"keys":[{"name":"a","key":"sk-é"}]}"#).contains("visible ASCII"));
381        assert!(err(r#"{"keys":[{"name":"a","key":"sk a"}]}"#).contains("visible ASCII"));
382        assert!(err(r#"{"keys":[{"name":"a b","key":"k"}]}"#).contains("[A-Za-z0-9._-]"));
383        assert!(err(r#"{"keys":[{"name":"a","key":"k","class":"urgent"}]}"#).contains("not a class"));
384        assert!(err(r#"{"keys":[{"name":"a","key":"k","class":"interactive","max_class":"agent"}]}"#)
385            .contains("lower priority"));
386        assert!(err(r#"{"keys":[{"name":"a","key":"k"},{"name":"a","key":"j"}]}"#).contains("twice"));
387        assert!(err(r#"{"keys":[{"name":"a","key":"k"},{"name":"b","key":"k"}]}"#).contains("share a secret"));
388        for bad in [
389            r#"{"keys":[{"name":"sk-b","key":"sk-a"},{"name":"b","key":"sk-b"}]}"#,
390            r#"{"keys":[{"name":"sk-b","key":"sk-a","class":"urgent"},{"name":"b","key":"sk-b"}]}"#,
391            r#"{"keys":[{"name":"b","key":"sk-b"},{"name":"sk-b","key":"sk-a"},{"name":"sk-b","key":"sk-c"}]}"#,
392            r#"{"keys":[{"name":"sk-b","key":"sk-b","key_env":"K"}]}"#,
393        ] {
394            let e = err(bad);
395            assert!(e.contains("equals a key's secret") && !e.contains("sk-b"), "{e}");
396        }
397        let e = KeyTable::parse_with_global(r#"{"keys":[{"name":"sk-g","key":"k","class":"urgent"}]}"#, no_env, Some("sk-g"))
398            .unwrap_err();
399        assert!(e.contains("keys[0]") && e.contains("--api-key") && !e.contains("sk-g"), "{e}");
400        let t = KeyTable::parse(r#"{"keys":[{"name":"a","key":"k"}]}"#, no_env).unwrap();
401        assert!(t.check_global_key(Some("k")).unwrap_err().contains("--api-key"));
402        let e = t.check_global_key(Some("a")).unwrap_err();
403        assert!(e.contains("keys[0]") && !e.contains("\"a\""), "{e}");
404        assert!(t.check_global_key(Some("other")).is_ok());
405        assert!(t.check_global_key(None).is_ok());
406    }
407
408    #[test]
409    fn concurrency_slots_release_on_drop() {
410        let t = KeyTable::parse(r#"{"keys":[{"name":"a","key":"k","max_concurrent":2}]}"#, no_env).unwrap();
411        let k = t.lookup("k").unwrap();
412        let a = k.admit(true, true).unwrap();
413        let _b = k.admit(true, true).unwrap();
414        assert_eq!(k.admit(true, true).unwrap_err(), Reject::Concurrency);
415        assert_eq!(k.in_flight(), 2);
416        drop(a);
417        assert_eq!(k.in_flight(), 1);
418        let _c = k.admit(true, true).unwrap();
419        assert_eq!(k.stats.requests.load(Ordering::Relaxed), 4);
420        assert_eq!(k.stats.rejected_concurrency.load(Ordering::Relaxed), 1);
421    }
422
423    #[test]
424    fn rate_limit_is_per_key_and_does_not_leak_slots() {
425        let t = KeyTable::parse(
426            r#"{"keys":[{"name":"a","key":"k","rate_limit_rpm":2,"max_concurrent":8},{"name":"b","key":"j"}]}"#,
427            no_env,
428        )
429        .unwrap();
430        let a = t.lookup("k").unwrap();
431        let _s1 = a.admit(true, true).unwrap();
432        let _s2 = a.admit(true, true).unwrap();
433        assert_eq!(a.admit(true, true).unwrap_err(), Reject::RateLimit);
434        assert_eq!(a.in_flight(), 2, "a rate-limited request gives its slot back");
435        let b = t.lookup("j").unwrap();
436        for _ in 0..10 {
437            b.admit(true, true).unwrap();
438        }
439        let m = t.metrics_text();
440        assert!(m.contains("superfluid_key_rejected_total{key=\"a\",reason=\"rate_limit\"} 1"), "{m}");
441        assert!(m.contains("superfluid_key_requests_total{key=\"b\"} 10"), "{m}");
442        assert!(m.contains("superfluid_key_in_flight{key=\"a\"} 2"), "{m}");
443        assert!(!m.contains("sk-") && !m.contains("\"k\""), "no secret in metrics: {m}");
444    }
445}
superfluid_daemon/keypolicy.rs

superfluid_daemon/
keypolicy.rs
