brain_edge/state/
mod.rs

1//! Shared state: the [`BrainPool`] connection cache to Brain plus the injected
2//! [`CredentialResolver`] + [`MeteringSink`] ports.
3//!
4//! The resolver turns an incoming request into a Brain wire credential; the
5//! [`BrainPool`] reuses one handshake per credential so we don't re-connect on
6//! every request; the metering sink observes each op's outcome. Self-host uses
7//! the bearer/no-op defaults ([`EdgeState::new`]); the gateway injects its own
8//! ports ([`EdgeState::with_ports`]).
9
10use std::sync::{Arc, OnceLock};
11use std::time::Duration;
12
13use axum::http::HeaderMap;
14use brain_db_sdk::BrainClient;
15
16use crate::config::EdgeConfig;
17use crate::crypto::{KeyRing, SealedSecret, SecretBinding};
18use crate::error::ApiError;
19use crate::pool::{BrainPool, BrainPoolConfig};
20use crate::port::{
21    BearerResolver, CredentialResolver, LlmCredentialStore, MeterEvent, MeteringSink, NoopMeter,
22    Outcome, ResolvedCredential,
23};
24use zeroize::Zeroizing;
25
26/// Cheap-to-clone application state.
27#[derive(Clone)]
28pub struct EdgeState {
29    inner: Arc<Inner>,
30}
31
32impl std::fmt::Debug for EdgeState {
33    /// The ports are trait objects and the pool holds live connections, so this
34    /// prints the shape rather than the contents.
35    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
36        f.debug_struct("EdgeState")
37            .field("pool", &self.inner.pool)
38            .finish_non_exhaustive()
39    }
40}
41
42struct Inner {
43    pool: BrainPool,
44    resolver: Arc<dyn CredentialResolver>,
45    meter: Arc<dyn MeteringSink>,
46    /// BYOK: where sealed credentials live and the keys that seal them.
47    ///
48    /// Unset by default (self-host) and settable once at startup. A
49    /// `OnceLock` rather than a field because the pair is all-or-nothing — a
50    /// store without a keyring could only hold plaintext, which is the exact
51    /// outcome this module exists to prevent — and because installing it must
52    /// not require `Inner: Clone`.
53    byok: OnceLock<Byok>,
54}
55
56/// The BYOK pair. Both or neither.
57struct Byok {
58    store: Arc<dyn LlmCredentialStore>,
59    keyring: Arc<KeyRing>,
60}
61
62impl EdgeState {
63    /// The self-host default: bearer passthrough, no metering, no BYOK.
64    #[must_use]
65    pub fn new(config: EdgeConfig) -> Self {
66        Self::with_ports(config, Arc::new(BearerResolver), Arc::new(NoopMeter))
67    }
68
69    /// Build state with an injected resolver + metering sink (the hosted gateway
70    /// path). The gateway passes an API-key resolver and its analytics sink.
71    #[must_use]
72    pub fn with_ports(
73        config: EdgeConfig,
74        resolver: Arc<dyn CredentialResolver>,
75        meter: Arc<dyn MeteringSink>,
76    ) -> Self {
77        let pool = BrainPool::new(BrainPoolConfig {
78            brain_addr: config.brain_addr,
79            pool_size: config.pool_size,
80            max_credentials: config.max_credentials,
81            idle_ttl: Duration::from_secs(config.idle_ttl_secs),
82        });
83        Self {
84            inner: Arc::new(Inner {
85                pool,
86                resolver,
87                meter,
88                byok: OnceLock::new(),
89            }),
90        }
91    }
92
93    /// Install the BYOK pieces: where sealed credentials live, and the keys
94    /// that seal them.
95    ///
96    /// Separate from [`Self::with_ports`] so the existing two-port constructor
97    /// keeps working unchanged and BYOK is opt-in. A deployment that never
98    /// calls this has no KEK, so [`Self::seal_llm_key`] refuses — the edge
99    /// will not fall back to storing a provider key in the clear.
100    #[must_use]
101    pub fn with_llm_credentials(
102        self,
103        store: Arc<dyn LlmCredentialStore>,
104        keyring: Arc<KeyRing>,
105    ) -> Self {
106        // Set-once: a second call is a wiring bug, not a reconfiguration, and
107        // silently swapping the keyring under live ciphertext would strand it.
108        if self.inner.byok.set(Byok { store, keyring }).is_err() {
109            tracing::warn!("BYOK already installed on this EdgeState; ignoring repeat install");
110        }
111        self
112    }
113
114    /// True when a KEK is configured, i.e. BYOK is usable on this edge.
115    #[must_use]
116    pub fn byok_enabled(&self) -> bool {
117        self.inner.byok.get().is_some()
118    }
119
120    /// Seal a plaintext provider key for `namespace`/`provider`.
121    ///
122    /// The plaintext is borrowed and never stored; only the returned
123    /// [`SealedSecret`] crosses the store boundary.
124    ///
125    /// # Errors
126    /// `501` when no KEK is configured — storing the key would mean writing it
127    /// in the clear, so the edge declines instead. `500` if sealing fails.
128    pub fn seal_llm_key(
129        &self,
130        plaintext: &str,
131        binding: SecretBinding<'_>,
132    ) -> Result<SealedSecret, ApiError> {
133        let keyring = self.inner.byok.get().map(|b| &b.keyring).ok_or_else(|| {
134            ApiError::not_implemented(
135                "BYOK is not configured on this edge: no key-encryption key is set, \
136                 so a provider key cannot be stored",
137            )
138        })?;
139        keyring
140            .seal(plaintext.as_bytes(), binding)
141            .map_err(|e| ApiError::internal(format!("seal provider key: {e}")))
142    }
143
144    /// Open a sealed provider key. The plaintext is wiped when dropped.
145    ///
146    /// # Errors
147    /// `501` when no KEK is configured; `500` when the ciphertext does not open
148    /// (wrong key, tampering, or a binding that does not match).
149    pub fn open_llm_key(
150        &self,
151        sealed: &SealedSecret,
152        binding: SecretBinding<'_>,
153    ) -> Result<Zeroizing<Vec<u8>>, ApiError> {
154        let keyring = self
155            .inner
156            .byok
157            .get()
158            .map(|b| &b.keyring)
159            .ok_or_else(|| ApiError::not_implemented("BYOK is not configured on this edge"))?;
160        keyring
161            .open(sealed, binding)
162            .map_err(|e| ApiError::internal(format!("open provider key: {e}")))
163    }
164
165    /// The sealed-credential store, when BYOK is configured.
166    ///
167    /// `None` without a keyring: a store with nothing to seal with could only
168    /// ever hold plaintext, so there is deliberately no way to reach one.
169    #[must_use]
170    pub fn llm_credentials(&self) -> Option<&Arc<dyn LlmCredentialStore>> {
171        self.inner.byok.get().map(|b| &b.store)
172    }
173
174    /// Resolve a request to a Brain credential via the injected resolver.
175    ///
176    /// # Errors
177    /// Propagates the resolver's `401`/`403` when no valid credential is present.
178    pub async fn resolve(&self, headers: &HeaderMap) -> Result<ResolvedCredential, ApiError> {
179        self.inner.resolver.resolve(headers).await
180    }
181
182    /// Record one op's outcome via the injected sink (fire-and-forget).
183    pub fn record(&self, op: &str, ident: &ResolvedCredential, outcome: Outcome) {
184        self.inner.meter.record(&MeterEvent {
185            tenant: ident.tenant.as_deref(),
186            op,
187            outcome,
188        });
189    }
190
191    /// A Brain client authenticated as `credential`, from the [`BrainPool`]'s
192    /// cached (or freshly opened) per-credential pool.
193    ///
194    /// # Errors
195    /// Returns an engine error if the pool fails to open a connection.
196    pub async fn client_for(&self, credential: &str) -> Result<Arc<BrainClient>, ApiError> {
197        self.inner
198            .pool
199            .client_for(credential)
200            .await
201            .map_err(|e| ApiError::from_brain(&e))
202    }
203
204    /// The shared Brain connection pool (for the background idle sweeper and
205    /// readiness probing).
206    #[must_use]
207    pub fn pool(&self) -> &BrainPool {
208        &self.inner.pool
209    }
210
211    /// The injected metering sink, shared. The wire proxy reuses it so wire-path
212    /// ops meter through the same sink as the HTTP data plane.
213    #[must_use]
214    pub fn meter(&self) -> Arc<dyn MeteringSink> {
215        Arc::clone(&self.inner.meter)
216    }
217}
218
219#[cfg(test)]
220mod byok_state_tests {
221    use super::*;
222    use crate::port::{InMemoryLlmCredentialStore, StoredLlmCredential};
223
224    const KEK: &str = "1111111111111111111111111111111111111111111111111111111111111111";
225
226    fn config() -> EdgeConfig {
227        // Explicit rather than `..Default::default()`: EdgeConfig has no
228        // Default, and these tests never open a socket or a Brain connection.
229        EdgeConfig {
230            listen_addr: "127.0.0.1:0".parse().unwrap(),
231            brain_addr: "127.0.0.1:1".parse().unwrap(),
232            pool_size: 1,
233            max_credentials: 1,
234            idle_ttl_secs: 1,
235            request_timeout_secs: 1,
236            max_body_bytes: 1024,
237            wire_listen_addr: None,
238            wire_rate_capacity: 0,
239            wire_rate_refill_per_sec: 0,
240        }
241    }
242
243    fn binding<'a>() -> SecretBinding<'a> {
244        SecretBinding {
245            namespace: "acme",
246            provider: "openai",
247        }
248    }
249
250    #[test]
251    fn byok_is_off_by_default() {
252        let st = EdgeState::new(config());
253        assert!(!st.byok_enabled());
254        assert!(
255            st.llm_credentials().is_none(),
256            "a store with no keyring could only hold plaintext, so it must be unreachable"
257        );
258    }
259
260    #[test]
261    fn without_a_kek_sealing_is_refused_rather_than_stored_in_the_clear() {
262        // The central safety property of this whole design: a misconfigured
263        // edge declines to hold the key at all. There is deliberately no
264        // plaintext fallback path to reach.
265        let st = EdgeState::new(config());
266        let err = st.seal_llm_key("sk-live-secret", binding()).unwrap_err();
267        assert_eq!(err.status(), axum::http::StatusCode::NOT_IMPLEMENTED);
268        assert!(
269            !format!("{err:?}").contains("sk-live-secret"),
270            "the rejected key must not appear in the error"
271        );
272    }
273
274    #[test]
275    fn installing_byok_enables_seal_and_open() {
276        let keyring = Arc::new(KeyRing::from_spec(&format!("k1:{KEK}"), "k1").unwrap());
277        let st = EdgeState::new(config())
278            .with_llm_credentials(Arc::new(InMemoryLlmCredentialStore::default()), keyring);
279
280        assert!(st.byok_enabled());
281        assert!(st.llm_credentials().is_some());
282
283        let sealed = st.seal_llm_key("sk-live-secret", binding()).unwrap();
284        assert_eq!(
285            &*st.open_llm_key(&sealed, binding()).unwrap(),
286            b"sk-live-secret"
287        );
288    }
289
290    #[test]
291    fn a_sealed_key_does_not_open_under_another_tenant_via_state() {
292        // End-to-end at the state layer, not just the crypto layer: the
293        // binding must survive the wiring.
294        let keyring = Arc::new(KeyRing::from_spec(&format!("k1:{KEK}"), "k1").unwrap());
295        let st = EdgeState::new(config())
296            .with_llm_credentials(Arc::new(InMemoryLlmCredentialStore::default()), keyring);
297
298        let sealed = st.seal_llm_key("acme-key", binding()).unwrap();
299        let other = SecretBinding {
300            namespace: "globex",
301            provider: "openai",
302        };
303        assert!(st.open_llm_key(&sealed, other).is_err());
304    }
305
306    #[tokio::test]
307    async fn a_stored_record_carries_ciphertext_not_the_key() {
308        let keyring = Arc::new(KeyRing::from_spec(&format!("k1:{KEK}"), "k1").unwrap());
309        let st = EdgeState::new(config())
310            .with_llm_credentials(Arc::new(InMemoryLlmCredentialStore::default()), keyring);
311
312        let secret = "sk-live-0123456789abcdef";
313        let sealed = st.seal_llm_key(secret, binding()).unwrap();
314        st.llm_credentials()
315            .unwrap()
316            .put(
317                "acme",
318                StoredLlmCredential {
319                    provider: "openai".into(),
320                    model: "gpt-4o-mini".into(),
321                    sealed,
322                    fingerprint: crate::port::key_fingerprint(secret),
323                    created_at_unix_secs: 0,
324                },
325            )
326            .await
327            .unwrap();
328
329        let got = st
330            .llm_credentials()
331            .unwrap()
332            .get("acme")
333            .await
334            .unwrap()
335            .unwrap();
336        assert!(
337            !got.sealed
338                .ciphertext
339                .windows(secret.len())
340                .any(|w| w == secret.as_bytes()),
341            "the stored record must not contain the provider key"
342        );
343        assert_eq!(got.fingerprint, "…cdef");
344        assert_eq!(
345            &*st.open_llm_key(&got.sealed, binding()).unwrap(),
346            secret.as_bytes()
347        );
348    }
349}
brain_edge/state/mod.rs

brain_edge/state/
mod.rs
