codetopo_server/
auth.rs

1// Copyright (C) 2026 Nrupal Akolkar
2// SPDX-License-Identifier: AGPL-3.0-or-later
3// Part of codetopo — Owned by Nrupal Akolkar (https://nrupalakolkar.com) · Built with Muse by Meta (https://muse.ai/join) — join Muse, then redeem code 101EEK in Settings within 48 hours of joining and we both earn tokens.
4
5//! Bearer-key authentication and per-request metering (Phase 2 design §4, §5).
6//!
7//! This module is the only place a secret is ever visible: handlers receive a
8//! [`KeyId`] (and, since paid-path S2, the key's [`KeyPlan`]) and nothing
9//! else, so a handler cannot leak a secret by accident and the metering
10//! layer has a stable identity to record without touching the credential.
11//!
12//! Order of operations per request:
13//!
14//! 1. Read `Authorization` and require the exact form `Bearer <key_id>:<secret>`.
15//! 2. Verify against the runtime [`KeyStore`] in constant time. The store
16//!    is read on every request, so issuance, rotation, and revocation
17//!    take effect without a restart (see [`crate::keys`]).
18//! 3. On success, attach [`KeyId`] and [`KeyPlan`] to the request extensions.
19//! 4. Enforce the plan's monthly query quota against the metering log
20//!    (paid-path S3): a key at its limit is answered `429` and, like a
21//!    `401`, is not metered.
22//! 5. Run the inner service.
23//! 6. Measure the response body and append exactly one metering record.
24//!
25//! A rejected request (step 1 or 2) is answered `401` and is *not* metered. See
26//! the deviation note in [`crate::metering`].
27
28use std::sync::Arc;
29use std::time::SystemTime;
30
31use axum::body::Body;
32use axum::extract::{Request, State};
33use axum::http::header::{AUTHORIZATION, CONTENT_LENGTH};
34use axum::http::HeaderValue;
35use axum::middleware::Next;
36use axum::response::{IntoResponse, Response};
37use axum::Router;
38
39use crate::error::ApiError;
40use crate::keys::{parse_bearer, KeyStore};
41use crate::metering::Metering;
42
43/// The authenticated `key_id`, attached to the request by [`auth_and_meter`].
44///
45/// Handlers read this to attribute work to a key. The secret is not reachable
46/// from here, and neither is any other caller's identity.
47#[derive(Debug, Clone, PartialEq, Eq)]
48pub struct KeyId(pub String);
49
50impl KeyId {
51    /// The key identifier as a string slice.
52    pub fn as_str(&self) -> &str {
53        &self.0
54    }
55}
56
57impl std::fmt::Display for KeyId {
58    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
59        f.write_str(&self.0)
60    }
61}
62
63/// The authenticated key's plan, attached beside [`KeyId`] by
64/// [`auth_and_meter`].
65///
66/// A plan is a provisioning fact, not a secret. Handlers read it to
67/// gate the host-local `repo_path` mode to
68/// [`crate::keys::SELF_HOSTED_PLAN`] keys (paid-path S2), and the
69/// abuse envelope reads it for per-plan limits (paid-path S3, see
70/// [`crate::limits`]): this layer enforces the monthly query quota,
71/// the index handler the rest.
72#[derive(Debug, Clone, PartialEq, Eq)]
73pub struct KeyPlan(pub String);
74
75impl KeyPlan {
76    /// The plan name as a string slice.
77    pub fn as_str(&self) -> &str {
78        &self.0
79    }
80
81    /// True for the operator's own keys (legacy import plan).
82    pub fn is_self_hosted(&self) -> bool {
83        self.0 == crate::keys::SELF_HOSTED_PLAN
84    }
85}
86
87/// Largest response body the middleware will buffer in order to measure it.
88///
89/// Buffering is the only way to learn a body's length, and an unbounded buffer
90/// is a denial-of-service primitive: `snapshot` is whole-graph by design, so a
91/// large graph would otherwise pin its full serialisation in RAM for a metric.
92/// A body that declares a length past this cap is passed through un-buffered
93/// and metered from its `Content-Length`.
94pub const MAX_RESPONSE_BYTES: usize = 64 * 1024 * 1024;
95
96/// Everything the auth + metering layer needs, shared by reference.
97#[derive(Debug, Clone)]
98pub struct AuthMetering {
99    keys: Arc<KeyStore>,
100    metering: Arc<Metering>,
101}
102
103impl AuthMetering {
104    /// Binds a key store and a metering log to one layer.
105    pub fn new(keys: Arc<KeyStore>, metering: Arc<Metering>) -> Self {
106        AuthMetering { keys, metering }
107    }
108}
109
110/// Wraps `router` so every route requires a bearer key and is metered.
111///
112/// There is deliberately no unauthenticated variant, not even for a health
113/// route: an unauthenticated route is a hole an agent can probe for free.
114pub fn metered<S>(router: Router<S>, auth: Arc<AuthMetering>) -> Router<S>
115where
116    S: Clone + Send + Sync + 'static,
117{
118    router.route_layer(axum::middleware::from_fn_with_state(auth, auth_and_meter))
119}
120
121/// Authenticates the request, runs the inner service, meters the result.
122async fn auth_and_meter(
123    State(auth): State<Arc<AuthMetering>>,
124    mut request: Request,
125    next: Next,
126) -> Response {
127    // One wall-clock reading serves both record fields: `ts` is the request
128    // start and `duration_ms` is `started.elapsed()`. A monotonic `Instant`
129    // would survive a clock step mid-request, but `ts` is a wall-clock reading
130    // by definition, so the one clock is the only one that could be reported.
131    let started = SystemTime::now();
132    let method = request.method().as_str().to_string();
133    let path = request.uri().path().to_string();
134
135    let presented = parse_bearer(
136        request
137            .headers()
138            .get(AUTHORIZATION)
139            .and_then(|value| value.to_str().ok()),
140    );
141    let Some((key_id, secret)) = presented else {
142        return unauthorized();
143    };
144    let verified = match auth.keys.verify(key_id, secret) {
145        Ok(verified) => verified,
146        Err(_) => return unauthorized(),
147    };
148    request
149        .extensions_mut()
150        .insert(KeyId(verified.key_id.clone()));
151    request
152        .extensions_mut()
153        .insert(KeyPlan(verified.plan.clone()));
154
155    // Monthly query quota (paid-path S3). The count is the metering
156    // log itself — a served query appends exactly one record — so
157    // usage and billing can never disagree. Index creation is not a
158    // query and is never refused here. A request refused here is
159    // not metered, like a `401`: it was not served, and metering it
160    // would let refusals consume the very quota they enforce. The
161    // count is of *completed* requests, so a burst already in flight
162    // when the limit lands can carry usage a little past it; the
163    // quota is an abuse envelope, not a billing transaction.
164    let limits = crate::limits::plan_limits(&verified.plan);
165    if let Some(max_queries) = limits.max_queries_per_month {
166        let is_index_op = method.as_str() == "POST" && path.as_str() == "/v1/index";
167        if !is_index_op {
168            let month = crate::metering::current_month();
169            if auth
170                .metering
171                .query_count_reaches(&verified.key_id, &month, max_queries)
172            {
173                let used = auth.metering.usage_in_month(&verified.key_id, &month).queries;
174                return ApiError::query_quota(&verified.plan, max_queries, used, &month)
175                    .into_response();
176            }
177        }
178    }
179
180    let response = next.run(request).await;
181    measure_and_meter(
182        auth.metering.as_ref(),
183        &verified.key_id,
184        &verified.plan,
185        &method,
186        path,
187        response,
188        started,
189    )
190    .await
191}
192
193/// A single `401` for every authentication failure.
194///
195/// The message is intentionally identical for a missing header, a bad scheme, an
196/// unknown `key_id`, and a wrong secret, so the response reveals nothing about
197/// which keys exist. `WWW-Authenticate` is still sent because that is a
198/// protocol requirement, not a disclosure.
199fn unauthorized() -> Response {
200    let mut response = ApiError::Unauthorized.into_response();
201    response.headers_mut().insert(
202        axum::http::header::WWW_AUTHENTICATE,
203        HeaderValue::from_static("Bearer realm=\"codetopo\""),
204    );
205    response
206}
207
208/// Measures the response body, appends one record, and rebuilds the response.
209async fn measure_and_meter(
210    metering: &Metering,
211    key_id: &str,
212    plan: &str,
213    method: &str,
214    path: String,
215    response: Response,
216    started: SystemTime,
217) -> Response {
218    let status = response.status().as_u16();
219    let (parts, body) = response.into_parts();
220
221    let (body, response_bytes) = match declared_length(&parts) {
222        // Declared past the cap: do not buffer. The declared length is the best
223        // available size and the body is served untouched.
224        Some(len) if len > MAX_RESPONSE_BYTES => (body, len as u64),
225        _ => match axum::body::to_bytes(body, MAX_RESPONSE_BYTES).await {
226            Ok(bytes) => {
227                let len = bytes.len() as u64;
228                (Body::from(bytes), len)
229            }
230            // The cap was hit without a declared length to check against, and
231            // `to_bytes` has already consumed the body, so it cannot be
232            // replayed. Answering 500 is the only honest option: serving an
233            // empty 200 would look like a successful empty result.
234            Err(err) => {
235                eprintln!("warning: response body exceeded the metering buffer cap: {err}");
236                return ApiError::Internal(format!(
237                    "response exceeded the {MAX_RESPONSE_BYTES}-byte metering buffer limit"
238                ))
239                .into_response();
240            }
241        },
242    };
243
244    let record = metering.build(key_id, plan, method, &path, status, response_bytes, started);
245    // Appended, not propagated: metering is accounting, and failing a served
246    // read-only query because the log was momentarily unwritable would trade a
247    // correct answer for a correct metric. `append` reports its own failures.
248    metering.append(&record);
249
250    Response::from_parts(parts, body)
251}
252
253/// The `Content-Length` a response declares, if it declares a valid one.
254fn declared_length(parts: &axum::http::response::Parts) -> Option<usize> {
255    parts
256        .headers
257        .get(CONTENT_LENGTH)
258        .and_then(|value| value.to_str().ok())
259        .and_then(|value| value.parse().ok())
260}
261
262#[cfg(test)]
263mod tests {
264    use super::*;
265    use axum::http::HeaderMap;
266
267    fn auth(key: &str) -> (Arc<AuthMetering>, tempfile::TempDir) {
268        let dir = tempfile::tempdir().expect("temp dir");
269        let keys = KeyStore::parse(&format!("{key}:s3cr3t\n"));
270        let metering = Arc::new(Metering::new(&dir.path().join("metering.log")));
271        (Arc::new(AuthMetering::new(Arc::new(keys), metering)), dir)
272    }
273
274    fn bearer(value: &str) -> HeaderMap {
275        let mut headers = HeaderMap::new();
276        headers.insert(AUTHORIZATION, value.parse().expect("header value"));
277        headers
278    }
279
280    /// Resolves the presented credential against `auth`, as the middleware does.
281    fn resolve(auth: &AuthMetering, headers: &HeaderMap) -> Result<String, ApiError> {
282        let Some((key_id, secret)) = parse_bearer(
283            headers
284                .get(AUTHORIZATION)
285                .and_then(|value| value.to_str().ok()),
286        ) else {
287            return Err(ApiError::Unauthorized);
288        };
289        auth.keys.verify(key_id, secret).map(|verified| verified.key_id)
290    }
291
292    #[test]
293    fn key_id_exposes_identifier_not_secret() {
294        let key_id = KeyId("research-agent".to_string());
295        assert_eq!(key_id.as_str(), "research-agent");
296        assert_eq!(key_id.to_string(), "research-agent");
297    }
298
299    #[test]
300    fn key_plan_names_the_self_hosted_plan() {
301        assert!(KeyPlan(crate::keys::SELF_HOSTED_PLAN.to_string()).is_self_hosted());
302        let hosted = KeyPlan("hosted".to_string());
303        assert!(!hosted.is_self_hosted());
304        assert_eq!(hosted.as_str(), "hosted");
305    }
306
307    #[test]
308    fn correct_credential_authenticates() {
309        let (auth, _dir) = auth("agent");
310        assert_eq!(
311            resolve(&auth, &bearer("Bearer agent:s3cr3t")).unwrap(),
312            "agent"
313        );
314    }
315
316    #[test]
317    fn scheme_is_case_insensitive() {
318        let (auth, _dir) = auth("agent");
319        assert_eq!(
320            resolve(&auth, &bearer("bearer agent:s3cr3t")).unwrap(),
321            "agent"
322        );
323    }
324
325    #[test]
326    fn wrong_secret_is_unauthorized() {
327        let (auth, _dir) = auth("agent");
328        assert!(matches!(
329            resolve(&auth, &bearer("Bearer agent:wrong")),
330            Err(ApiError::Unauthorized)
331        ));
332    }
333
334    #[test]
335    fn unknown_key_is_unauthorized() {
336        let (auth, _dir) = auth("agent");
337        assert!(matches!(
338            resolve(&auth, &bearer("Bearer nobody:s3cr3t")),
339            Err(ApiError::Unauthorized)
340        ));
341    }
342
343    #[test]
344    fn missing_header_is_unauthorized() {
345        let (auth, _dir) = auth("agent");
346        assert!(matches!(
347            resolve(&auth, &HeaderMap::new()),
348            Err(ApiError::Unauthorized)
349        ));
350    }
351
352    #[test]
353    fn wrong_scheme_is_unauthorized() {
354        let (auth, _dir) = auth("agent");
355        assert!(matches!(
356            resolve(&auth, &bearer("Basic agent:s3cr3t")),
357            Err(ApiError::Unauthorized)
358        ));
359    }
360
361    #[test]
362    fn token_without_colon_is_unauthorized() {
363        let (auth, _dir) = auth("agent");
364        assert!(matches!(
365            resolve(&auth, &bearer("Bearer agentonly")),
366            Err(ApiError::Unauthorized)
367        ));
368    }
369
370    #[test]
371    fn declared_length_reads_content_length() {
372        let response = Response::builder()
373            .header(CONTENT_LENGTH, "42")
374            .body(Body::empty())
375            .unwrap();
376        let (parts, _body) = response.into_parts();
377        assert_eq!(declared_length(&parts), Some(42));
378    }
379
380    #[test]
381    fn declared_length_absent_when_undeclared() {
382        let (parts, _body) = Response::new(Body::empty()).into_parts();
383        assert_eq!(declared_length(&parts), None);
384    }
385
386    #[test]
387    fn metering_records_the_served_body() {
388        let (auth, dir) = auth("agent");
389        let response = Response::builder()
390            .status(200)
391            .header(CONTENT_LENGTH, "5")
392            .body(Body::from("hello"))
393            .unwrap();
394        let out = tokio::runtime::Builder::new_current_thread()
395            .enable_all()
396            .build()
397            .unwrap()
398            .block_on(measure_and_meter(
399                auth.metering.as_ref(),
400                "agent",
401                "self-hosted",
402                "GET",
403                "/v1/graphs/x/stats".to_string(),
404                response,
405                SystemTime::now(),
406            ));
407        assert_eq!(out.status(), 200);
408        let log = std::fs::read_to_string(dir.path().join("metering.log")).unwrap();
409        let record: serde_json::Value = serde_json::from_str(log.trim()).unwrap();
410        assert_eq!(record["key_id"], "agent");
411        assert_eq!(record["plan"], "self-hosted");
412        assert_eq!(record["method"], "GET");
413        assert_eq!(record["path"], "/v1/graphs/x/stats");
414        assert_eq!(record["status"], 200);
415        assert_eq!(record["response_bytes"], 5);
416    }
417}
codetopo_server/auth.rs

codetopo_server/
auth.rs
