Skip to main content

ai_crew_sync/
auth.rs

1use axum::{
2    extract::{Request, State},
3    http::StatusCode,
4    middleware::Next,
5    response::{IntoResponse, Response},
6};
7use rand::RngCore;
8use sha2::{Digest, Sha256};
9use sqlx::PgPool;
10use uuid::Uuid;
11
12pub const TOKEN_PREFIX: &str = "acs_";
13
14/// Identity resolved from the bearer token, injected into the HTTP request
15/// extensions so tool handlers can read it. Every tool call is scoped to this.
16#[derive(Clone, Debug)]
17pub struct AuthCtx {
18    pub agent_id: Uuid,
19    pub agent_name: String,
20    pub team_id: Uuid,
21    pub team_slug: String,
22}
23
24/// Generate a fresh opaque token. Returned once, never stored in the clear.
25pub fn generate_token() -> String {
26    let mut bytes = [0u8; 32];
27    rand::thread_rng().fill_bytes(&mut bytes);
28    format!("{TOKEN_PREFIX}{}", hex::encode(bytes))
29}
30
31pub fn hash_token(raw: &str) -> Vec<u8> {
32    Sha256::digest(raw.trim().as_bytes()).to_vec()
33}
34
35/// First 12 characters, kept in plaintext purely so humans can tell tokens
36/// apart in `ai-crew-sync token list`.
37pub fn token_prefix(raw: &str) -> String {
38    raw.chars().take(12).collect()
39}
40
41struct AuthRow {
42    token_id: Uuid,
43    agent_id: Uuid,
44    agent_name: String,
45    agent_disabled: bool,
46    team_id: Uuid,
47    team_slug: String,
48}
49
50pub async fn resolve_token(pool: &PgPool, raw: &str) -> Result<AuthCtx, AuthError> {
51    if !raw.starts_with(TOKEN_PREFIX) {
52        return Err(AuthError::Invalid);
53    }
54    let hash = hash_token(raw);
55
56    let row = sqlx::query_as::<
57        _,
58        (
59            Uuid,
60            Uuid,
61            String,
62            Option<chrono::DateTime<chrono::Utc>>,
63            Uuid,
64            String,
65        ),
66    >(
67        r#"
68        SELECT t.id, a.id, a.name, a.disabled_at, tm.id, tm.slug
69        FROM api_tokens t
70        JOIN agents a ON a.id = t.agent_id
71        JOIN teams tm ON tm.id = a.team_id
72        WHERE t.token_hash = $1 AND t.revoked_at IS NULL
73        "#,
74    )
75    .bind(&hash)
76    .fetch_optional(pool)
77    .await
78    .map_err(|e| {
79        tracing::error!(error = %e, "token lookup failed");
80        AuthError::Internal
81    })?;
82
83    let Some((token_id, agent_id, agent_name, disabled_at, team_id, team_slug)) = row else {
84        return Err(AuthError::Invalid);
85    };
86    let row = AuthRow {
87        token_id,
88        agent_id,
89        agent_name,
90        agent_disabled: disabled_at.is_some(),
91        team_id,
92        team_slug,
93    };
94
95    if row.agent_disabled {
96        return Err(AuthError::Disabled);
97    }
98
99    // Best-effort: record usage without blocking the request path on failure.
100    let _ = sqlx::query("UPDATE api_tokens SET last_used_at = now() WHERE id = $1")
101        .bind(row.token_id)
102        .execute(pool)
103        .await;
104
105    Ok(AuthCtx {
106        agent_id: row.agent_id,
107        agent_name: row.agent_name,
108        team_id: row.team_id,
109        team_slug: row.team_slug,
110    })
111}
112
113#[derive(Debug)]
114pub enum AuthError {
115    Missing,
116    Invalid,
117    Disabled,
118    Internal,
119}
120
121impl IntoResponse for AuthError {
122    fn into_response(self) -> Response {
123        let (status, msg) = match self {
124            AuthError::Missing => (StatusCode::UNAUTHORIZED, "missing bearer token"),
125            AuthError::Invalid => (StatusCode::UNAUTHORIZED, "invalid or revoked token"),
126            AuthError::Disabled => (StatusCode::FORBIDDEN, "agent is disabled"),
127            AuthError::Internal => (StatusCode::INTERNAL_SERVER_ERROR, "internal error"),
128        };
129        let body = serde_json::json!({ "error": msg });
130        let mut resp = (status, axum::Json(body)).into_response();
131        if matches!(self, AuthError::Missing | AuthError::Invalid) {
132            resp.headers_mut().insert(
133                axum::http::header::WWW_AUTHENTICATE,
134                axum::http::HeaderValue::from_static("Bearer"),
135            );
136        }
137        resp
138    }
139}
140
141/// Axum middleware: validates `Authorization: Bearer <token>` and inserts the
142/// resulting [`AuthCtx`] into the request extensions, where rmcp tool handlers
143/// pick it up via `RequestContext -> http::request::Parts -> extensions`.
144pub async fn require_bearer(
145    State(pool): State<PgPool>,
146    mut req: Request,
147    next: Next,
148) -> Result<Response, AuthError> {
149    let raw = req
150        .headers()
151        .get(axum::http::header::AUTHORIZATION)
152        .and_then(|v| v.to_str().ok())
153        .and_then(|v| {
154            v.strip_prefix("Bearer ")
155                .or_else(|| v.strip_prefix("bearer "))
156        })
157        .map(str::trim)
158        .filter(|v| !v.is_empty())
159        .ok_or(AuthError::Missing)?
160        .to_owned();
161
162    let ctx = resolve_token(&pool, &raw).await?;
163    tracing::debug!(agent = %ctx.agent_name, team = %ctx.team_slug, "authenticated");
164    req.extensions_mut().insert(ctx);
165    Ok(next.run(req).await)
166}