Skip to main content

sim_lib_openai_server/runtime/
keys.rs

1use std::sync::{Arc, Mutex, OnceLock};
2
3use serde_json::Value as JsonValue;
4use sha2::{Digest, Sha256};
5use sim_citizen_derive::non_citizen;
6use sim_kernel::{
7    CapabilityName, CapabilitySet, Cx, DefaultFactory, Error, Expr, GrantSeat, NoopEvalPolicy,
8    Object, ObjectCompat, Result, Symbol, Table, Value,
9};
10use sim_lib_net_core::hex_encode;
11
12use crate::objects::GatewayRequest;
13
14macro_rules! grant_into_result {
15    ($grant:expr) => {{
16        #[allow(clippy::let_unit_value)]
17        let grant_result = $grant;
18        #[allow(clippy::unit_arg)]
19        grant_result.into_result()
20    }};
21}
22
23/// Object tag identifying a serialized [`OpenAiGatewayKey`] descriptor.
24pub const OPENAI_GATEWAY_KEY_OBJECT: &str = "openai-gateway/key";
25const REDACTED_HEADER_VALUE: &str = "[redacted]";
26
27/// An API key registered with the gateway, holding its hashed secret and the
28/// capability ceiling it grants.
29///
30/// The plaintext secret is never stored; the key is identified by the SHA-256
31/// hash of the secret and carries the [`CapabilitySet`] that bounds what a
32/// request authenticated with it may do.
33#[derive(Clone, Debug, PartialEq, Eq)]
34#[non_citizen(
35    reason = "gateway key object stores redacted credential policy; serializable projection is openai/GatewayKey descriptor",
36    kind = "handle",
37    descriptor = "openai/GatewayKey"
38)]
39pub struct OpenAiGatewayKey {
40    id: String,
41    key_hash: String,
42    capabilities: CapabilitySet,
43    default_policy: Expr,
44}
45
46impl OpenAiGatewayKey {
47    /// Returns a key from a precomputed secret hash and its capability ceiling.
48    pub fn new(key_hash: impl Into<String>, capabilities: CapabilitySet) -> Self {
49        let key_hash = key_hash.into();
50        let id = key_id(&key_hash);
51        let default_policy = key_default_policy_expr(&capabilities);
52        Self {
53            id,
54            key_hash,
55            capabilities,
56            default_policy,
57        }
58    }
59
60    /// Returns a key built by hashing `secret`, with the given capability ceiling.
61    pub fn from_secret(secret: &str, capabilities: CapabilitySet) -> Self {
62        Self::new(key_hash(secret), capabilities)
63    }
64
65    /// Returns the public key id (a `key_`-prefixed hash fragment).
66    pub fn id(&self) -> &str {
67        &self.id
68    }
69
70    /// Returns the hex SHA-256 hash of the key secret.
71    pub fn key_hash(&self) -> &str {
72        &self.key_hash
73    }
74
75    /// Returns a short, non-verifier fingerprint for reports.
76    pub fn fingerprint(&self) -> String {
77        key_fingerprint(&self.key_hash)
78    }
79
80    /// Returns the capability ceiling granted by this key.
81    pub fn capabilities(&self) -> &CapabilitySet {
82        &self.capabilities
83    }
84
85    /// Returns the default policy expression derived from the capability ceiling.
86    pub fn default_policy(&self) -> &Expr {
87        &self.default_policy
88    }
89
90    /// Returns the key as a serializable map descriptor.
91    pub fn to_expr(&self) -> Expr {
92        Expr::Map(vec![
93            field("object", Expr::String(OPENAI_GATEWAY_KEY_OBJECT.to_owned())),
94            field("id", Expr::String(self.id.clone())),
95            field("fingerprint", Expr::String(self.fingerprint())),
96            field("capabilities", capabilities_expr(&self.capabilities)),
97            field("default-policy", self.default_policy.clone()),
98        ])
99    }
100}
101
102impl Object for OpenAiGatewayKey {
103    fn display(&self, _cx: &mut Cx) -> Result<String> {
104        Ok(format!("#<openai-gateway-key {}>", self.id))
105    }
106
107    fn as_any(&self) -> &dyn std::any::Any {
108        self
109    }
110}
111
112impl ObjectCompat for OpenAiGatewayKey {
113    fn as_expr(&self, _cx: &mut Cx) -> Result<Expr> {
114        Ok(self.to_expr())
115    }
116}
117
118/// Thread-safe table of gateway API keys indexed by secret hash.
119///
120/// Keys are stored in a SIM table behind a shared mutex so the same table can
121/// be cloned across route state; an anonymous capability ceiling applies to
122/// requests that present no recognized key.
123#[derive(Clone)]
124pub struct OpenAiKeyTable {
125    inner: Arc<OpenAiKeyTableInner>,
126}
127
128struct OpenAiKeyTableInner {
129    cx: Mutex<Cx>,
130    keys: Value,
131    anonymous: CapabilitySet,
132}
133
134impl OpenAiKeyTable {
135    /// Returns an empty key table with no anonymous capabilities.
136    pub fn new() -> Result<Self> {
137        Self::with_anonymous(CapabilitySet::new())
138    }
139
140    /// Returns an empty key table whose unauthenticated requests get `anonymous`.
141    pub fn with_anonymous(anonymous: CapabilitySet) -> Result<Self> {
142        let cx = Cx::new(Arc::new(NoopEvalPolicy), Arc::new(DefaultFactory));
143        let keys = cx.factory().table(Vec::new())?;
144        Ok(Self {
145            inner: Arc::new(OpenAiKeyTableInner {
146                cx: Mutex::new(cx),
147                keys,
148                anonymous,
149            }),
150        })
151    }
152
153    /// Hashes `secret`, registers it with `capabilities`, and returns the new key.
154    pub fn add_secret(
155        &self,
156        secret: &str,
157        capabilities: CapabilitySet,
158    ) -> Result<OpenAiGatewayKey> {
159        let key = OpenAiGatewayKey::from_secret(secret, capabilities);
160        self.add_key(key.clone())?;
161        Ok(key)
162    }
163
164    /// Inserts an already-constructed key, indexing it by its secret hash.
165    pub fn add_key(&self, key: OpenAiGatewayKey) -> Result<()> {
166        let mut cx = self.cx()?;
167        let value = cx.factory().opaque(Arc::new(key.clone()))?;
168        table_impl(&self.inner.keys)?.set(&mut cx, Symbol::new(key.key_hash()), value)
169    }
170
171    /// Returns the registered key matching `secret`, or `None` if unknown.
172    pub fn key_for_secret(&self, secret: &str) -> Result<Option<OpenAiGatewayKey>> {
173        self.key_for_hash(&key_hash(secret))
174    }
175
176    /// Returns the key authenticating `request`, reading its bearer/api-key header.
177    pub fn key_for_request(&self, request: &GatewayRequest) -> Result<Option<OpenAiGatewayKey>> {
178        presented_key(request)
179            .map(|secret| self.key_for_secret(secret))
180            .unwrap_or(Ok(None))
181    }
182
183    /// Returns all registered keys.
184    pub fn list_keys(&self) -> Result<Vec<OpenAiGatewayKey>> {
185        let mut cx = self.cx()?;
186        Ok(table_impl(&self.inner.keys)?
187            .entries(&mut cx)?
188            .into_iter()
189            .filter_map(|(_, value)| value.object().downcast_ref::<OpenAiGatewayKey>().cloned())
190            .collect())
191    }
192
193    /// Returns the capabilities `request` may exercise.
194    ///
195    /// The result is the key's ceiling (or the anonymous ceiling when no key is
196    /// presented), intersected with any capabilities the request body explicitly
197    /// requests, so a request can only narrow -- never widen -- its key's grant.
198    pub fn effective_capabilities(&self, request: &GatewayRequest) -> Result<CapabilitySet> {
199        let ceiling = self
200            .key_for_request(request)?
201            .map(|key| key.capabilities().clone())
202            .unwrap_or_else(|| self.inner.anonymous.clone());
203        Ok(requested_capabilities(request)
204            .map(|requested| intersect_capabilities(&requested, &ceiling))
205            .unwrap_or(ceiling))
206    }
207
208    /// Runs `f` with `cx` scoped to the request's effective capabilities.
209    pub fn with_effective_capabilities<T>(
210        &self,
211        cx: &mut Cx,
212        request: &GatewayRequest,
213        f: impl FnOnce(&mut Cx) -> Result<T>,
214    ) -> Result<T> {
215        cx.with_capabilities(self.effective_capabilities(request)?, f)
216    }
217
218    fn key_for_hash(&self, hash: &str) -> Result<Option<OpenAiGatewayKey>> {
219        let mut cx = self.cx()?;
220        let value = table_impl(&self.inner.keys)?.get(&mut cx, Symbol::new(hash))?;
221        Ok(value.object().downcast_ref::<OpenAiGatewayKey>().cloned())
222    }
223
224    fn cx(&self) -> Result<std::sync::MutexGuard<'_, Cx>> {
225        self.inner
226            .cx
227            .lock()
228            .map_err(|_| Error::PoisonedLock("openai gateway key table"))
229    }
230}
231
232impl Default for OpenAiKeyTable {
233    fn default() -> Self {
234        Self::new().expect("in-memory SIM key table creation is infallible")
235    }
236}
237
238/// Returns the process-wide shared gateway key table, initializing it on first use.
239pub fn global_openai_key_table() -> &'static OpenAiKeyTable {
240    static TABLE: OnceLock<OpenAiKeyTable> = OnceLock::new();
241    TABLE.get_or_init(OpenAiKeyTable::default)
242}
243
244/// Returns the lowercase hex SHA-256 hash of a key secret.
245pub fn key_hash(secret: &str) -> String {
246    hex_encode(&Sha256::digest(secret.as_bytes()))
247}
248
249/// Returns a copy of `request` with sensitive auth headers replaced by `[redacted]`.
250pub fn redacted_gateway_request(request: &GatewayRequest) -> GatewayRequest {
251    GatewayRequest::new(
252        request.method().to_owned(),
253        request.path().to_owned(),
254        redact_headers(request.headers()),
255        request.body().to_vec(),
256    )
257}
258
259/// Grants every capability in `capabilities` to `cx`, through the host-held
260/// `seat` minted when `cx` was constructed.
261pub fn grant_capability_set(
262    seat: &GrantSeat,
263    cx: &mut Cx,
264    capabilities: &CapabilitySet,
265) -> Result<()> {
266    for capability in capabilities.iter().cloned() {
267        grant_into_result!(seat.grant(cx, capability))?;
268    }
269    Ok(())
270}
271
272trait GrantOutcome {
273    fn into_result(self) -> Result<()>;
274}
275
276impl GrantOutcome for () {
277    fn into_result(self) -> Result<()> {
278        Ok(())
279    }
280}
281
282impl GrantOutcome for Result<()> {
283    fn into_result(self) -> Result<()> {
284        self
285    }
286}
287
288fn key_id(hash: &str) -> String {
289    let prefix_len = hash.len().min(12);
290    format!("key_{}", &hash[..prefix_len])
291}
292
293fn key_fingerprint(hash: &str) -> String {
294    let prefix_len = hash.len().min(8);
295    format!("sha256:{}...", &hash[..prefix_len])
296}
297
298fn key_default_policy_expr(capabilities: &CapabilitySet) -> Expr {
299    Expr::Map(vec![field(
300        "capability-ceiling",
301        capabilities_expr(capabilities),
302    )])
303}
304
305fn capabilities_expr(capabilities: &CapabilitySet) -> Expr {
306    Expr::Vector(
307        capabilities
308            .iter()
309            .map(|capability| Expr::String(capability.as_str().to_owned()))
310            .collect(),
311    )
312}
313
314fn requested_capabilities(request: &GatewayRequest) -> Option<CapabilitySet> {
315    let body = serde_json::from_slice::<JsonValue>(request.body()).ok()?;
316    let object = body.as_object()?;
317    object
318        .get("capabilities")
319        .or_else(|| {
320            object
321                .get("sim")
322                .and_then(JsonValue::as_object)
323                .and_then(|sim| sim.get("capabilities"))
324        })
325        .and_then(capability_set_from_json)
326}
327
328fn capability_set_from_json(value: &JsonValue) -> Option<CapabilitySet> {
329    let mut capabilities = CapabilitySet::new();
330    match value {
331        JsonValue::String(name) => capabilities.insert(CapabilityName::new(name.clone())),
332        JsonValue::Array(items) => {
333            for item in items {
334                let name = item.as_str()?;
335                capabilities.insert(CapabilityName::new(name.to_owned()));
336            }
337        }
338        _ => return None,
339    }
340    Some(capabilities)
341}
342
343fn intersect_capabilities(left: &CapabilitySet, right: &CapabilitySet) -> CapabilitySet {
344    let mut capabilities = CapabilitySet::new();
345    for capability in left.iter() {
346        if right.contains(capability) {
347            capabilities.insert(capability.clone());
348        }
349    }
350    capabilities
351}
352
353fn presented_key(request: &GatewayRequest) -> Option<&str> {
354    for (name, value) in request.headers() {
355        if name.eq_ignore_ascii_case("authorization") {
356            let value = value.trim();
357            if let Some((scheme, token)) = value.split_once(' ')
358                && scheme.eq_ignore_ascii_case("bearer")
359            {
360                let token = token.trim();
361                if !token.is_empty() {
362                    return Some(token);
363                }
364            }
365        }
366        if is_api_key_header(name) {
367            let value = value.trim();
368            if !value.is_empty() {
369                return Some(value);
370            }
371        }
372    }
373    None
374}
375
376fn redact_headers(headers: &[(String, String)]) -> Vec<(String, String)> {
377    headers
378        .iter()
379        .map(|(name, value)| {
380            if is_sensitive_header(name) {
381                (name.clone(), REDACTED_HEADER_VALUE.to_owned())
382            } else {
383                (name.clone(), value.clone())
384            }
385        })
386        .collect()
387}
388
389fn is_api_key_header(name: &str) -> bool {
390    name.eq_ignore_ascii_case("x-api-key")
391        || name.eq_ignore_ascii_case("api-key")
392        || name.eq_ignore_ascii_case("openai-api-key")
393        || name.eq_ignore_ascii_case("x-openai-api-key")
394}
395
396fn is_sensitive_header(name: &str) -> bool {
397    name.eq_ignore_ascii_case("authorization")
398        || name.eq_ignore_ascii_case("proxy-authorization")
399        || is_api_key_header(name)
400}
401
402fn table_impl(value: &Value) -> Result<&dyn Table> {
403    value.object().as_table_impl().ok_or(Error::TypeMismatch {
404        expected: "SIM table",
405        found: "non-table",
406    })
407}
408
409use sim_value::build::entry as field;