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#[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
24pub 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
35pub 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 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
141pub 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}