postrust_proxy/saas/
auth.rs

1//! Authentication middleware for SaaS domain management.
2//!
3//! Supports both JWT and API key authentication.
4
5use crate::saas::api_keys::hash_api_key;
6use crate::saas::db;
7use axum::{
8    extract::{Request, State},
9    http::{header, StatusCode},
10    middleware::Next,
11    response::{IntoResponse, Response},
12};
13use serde::Deserialize;
14use sqlx::PgPool;
15use std::sync::Arc;
16use uuid::Uuid;
17
18/// Authentication context extracted from the request.
19#[derive(Clone, Debug)]
20pub struct AuthContext {
21    /// The tenant ID
22    pub tenant_id: Uuid,
23    /// The authentication type used
24    pub auth_type: AuthType,
25    /// Scopes available to this authentication
26    pub scopes: Vec<String>,
27}
28
29/// Type of authentication used.
30#[derive(Clone, Debug)]
31pub enum AuthType {
32    /// JWT authentication
33    Jwt {
34        /// User ID from JWT claims
35        user_id: String,
36        /// Role from JWT claims
37        role: Option<String>,
38    },
39    /// API key authentication
40    ApiKey {
41        /// API key ID
42        key_id: Uuid,
43    },
44}
45
46impl AuthContext {
47    /// Check if the context has a specific scope.
48    pub fn has_scope(&self, scope: &str) -> bool {
49        self.scopes.iter().any(|s| s == scope || s == "*")
50    }
51
52    /// Check if the context can read a resource.
53    pub fn can_read(&self, resource: &str) -> bool {
54        self.has_scope(&format!("{}:read", resource))
55            || self.has_scope(&format!("{}:write", resource))
56            || self.has_scope("*")
57    }
58
59    /// Check if the context can write to a resource.
60    pub fn can_write(&self, resource: &str) -> bool {
61        self.has_scope(&format!("{}:write", resource)) || self.has_scope("*")
62    }
63}
64
65/// SaaS authentication layer state.
66#[derive(Clone)]
67pub struct SaasAuthLayer {
68    pool: PgPool,
69    jwt_secret: Option<String>,
70}
71
72impl SaasAuthLayer {
73    /// Create a new SaaS auth layer.
74    pub fn new(pool: PgPool, jwt_secret: Option<String>) -> Self {
75        Self { pool, jwt_secret }
76    }
77}
78
79/// Authentication middleware function.
80pub async fn auth_middleware(
81    State(auth_layer): State<Arc<SaasAuthLayer>>,
82    mut request: Request,
83    next: Next,
84) -> Result<Response, AuthError> {
85    let auth_header = request
86        .headers()
87        .get(header::AUTHORIZATION)
88        .and_then(|h| h.to_str().ok());
89
90    let auth_context = match auth_header {
91        Some(header) if header.starts_with("Bearer ") => {
92            // JWT authentication
93            let token = header.strip_prefix("Bearer ").unwrap();
94            validate_jwt(token, &auth_layer).await?
95        }
96        Some(header) if header.starts_with("ApiKey ") => {
97            // API key authentication
98            let key = header.strip_prefix("ApiKey ").unwrap();
99            validate_api_key(key, &auth_layer.pool).await?
100        }
101        Some(header) if header.starts_with("Basic ") => {
102            // Could support basic auth in the future
103            return Err(AuthError::UnsupportedAuthMethod);
104        }
105        Some(_) => {
106            return Err(AuthError::InvalidFormat);
107        }
108        None => {
109            return Err(AuthError::MissingHeader);
110        }
111    };
112
113    // Check if tenant is active
114    if !db::is_tenant_active(&auth_layer.pool, auth_context.tenant_id).await? {
115        return Err(AuthError::TenantSuspended);
116    }
117
118    // Insert auth context into request extensions
119    request.extensions_mut().insert(auth_context);
120
121    Ok(next.run(request).await)
122}
123
124/// Validate a JWT token.
125async fn validate_jwt(token: &str, auth_layer: &SaasAuthLayer) -> Result<AuthContext, AuthError> {
126    let secret = auth_layer
127        .jwt_secret
128        .as_ref()
129        .ok_or(AuthError::JwtNotConfigured)?;
130
131    let claims = decode_jwt(token, secret)?;
132
133    // Get tenant ID from claims
134    let tenant_id = claims
135        .tenant_id
136        .ok_or_else(|| AuthError::MissingClaim("tenant_id".into()))?;
137
138    Ok(AuthContext {
139        tenant_id,
140        auth_type: AuthType::Jwt {
141            user_id: claims.sub,
142            role: claims.role,
143        },
144        scopes: claims.scopes.unwrap_or_else(|| vec!["*".to_string()]),
145    })
146}
147
148/// Validate an API key.
149async fn validate_api_key(key: &str, pool: &PgPool) -> Result<AuthContext, AuthError> {
150    let key_hash = hash_api_key(key);
151
152    let validation = db::validate_api_key_by_hash(pool, &key_hash)
153        .await?
154        .ok_or(AuthError::InvalidApiKey)?;
155
156    if !validation.enabled {
157        return Err(AuthError::ApiKeyDisabled);
158    }
159
160    // Update last used timestamp asynchronously
161    let pool = pool.clone();
162    let key_id = validation.id;
163    tokio::spawn(async move {
164        let _ = db::update_last_used(&pool, key_id).await;
165    });
166
167    Ok(AuthContext {
168        tenant_id: validation.tenant_id,
169        auth_type: AuthType::ApiKey {
170            key_id: validation.id,
171        },
172        scopes: validation.scopes,
173    })
174}
175
176/// JWT claims structure.
177#[derive(Debug, Deserialize)]
178struct JwtClaims {
179    /// Subject (user ID)
180    sub: String,
181    /// Tenant ID
182    tenant_id: Option<Uuid>,
183    /// User role
184    role: Option<String>,
185    /// Scopes
186    scopes: Option<Vec<String>>,
187    /// Expiration time
188    exp: Option<i64>,
189}
190
191/// Decode and validate a JWT token (HMAC-SHA256 signature, then expiry).
192fn decode_jwt(token: &str, secret: &str) -> Result<JwtClaims, AuthError> {
193    use jsonwebtoken::{decode, Algorithm, DecodingKey, Validation};
194
195    let mut validation = Validation::new(Algorithm::HS256);
196    // `exp` is optional for these tokens: don't require it, and check it
197    // manually below when present so absence isn't an error.
198    validation.required_spec_claims.clear();
199    validation.validate_exp = false;
200
201    let token_data = decode::<JwtClaims>(
202        token,
203        &DecodingKey::from_secret(secret.as_bytes()),
204        &validation,
205    )
206    .map_err(|e| match e.kind() {
207        jsonwebtoken::errors::ErrorKind::ExpiredSignature => AuthError::TokenExpired,
208        _ => AuthError::InvalidToken(e.to_string()),
209    })?;
210
211    let claims = token_data.claims;
212
213    // Check expiration
214    if let Some(exp) = claims.exp {
215        let now = chrono::Utc::now().timestamp();
216        if exp < now {
217            return Err(AuthError::TokenExpired);
218        }
219    }
220
221    Ok(claims)
222}
223
224/// Authentication errors.
225#[derive(Debug)]
226pub enum AuthError {
227    MissingHeader,
228    InvalidFormat,
229    InvalidToken(String),
230    TokenExpired,
231    MissingClaim(String),
232    InvalidApiKey,
233    ApiKeyDisabled,
234    ApiKeyExpired,
235    TenantSuspended,
236    JwtNotConfigured,
237    UnsupportedAuthMethod,
238    Database(sqlx::Error),
239}
240
241impl From<sqlx::Error> for AuthError {
242    fn from(err: sqlx::Error) -> Self {
243        AuthError::Database(err)
244    }
245}
246
247impl From<crate::error::ProxyError> for AuthError {
248    fn from(err: crate::error::ProxyError) -> Self {
249        match err {
250            crate::error::ProxyError::Database(e) => AuthError::Database(e),
251            _ => AuthError::InvalidToken(err.to_string()),
252        }
253    }
254}
255
256impl IntoResponse for AuthError {
257    fn into_response(self) -> Response {
258        let (status, message) = match self {
259            AuthError::MissingHeader => (
260                StatusCode::UNAUTHORIZED,
261                "Missing Authorization header".to_string(),
262            ),
263            AuthError::InvalidFormat => (
264                StatusCode::UNAUTHORIZED,
265                "Invalid Authorization header format".to_string(),
266            ),
267            AuthError::InvalidToken(msg) => {
268                (StatusCode::UNAUTHORIZED, format!("Invalid token: {}", msg))
269            }
270            AuthError::TokenExpired => (StatusCode::UNAUTHORIZED, "Token has expired".to_string()),
271            AuthError::MissingClaim(claim) => (
272                StatusCode::UNAUTHORIZED,
273                format!("Missing required claim: {}", claim),
274            ),
275            AuthError::InvalidApiKey => (StatusCode::UNAUTHORIZED, "Invalid API key".to_string()),
276            AuthError::ApiKeyDisabled => {
277                (StatusCode::UNAUTHORIZED, "API key is disabled".to_string())
278            }
279            AuthError::ApiKeyExpired => {
280                (StatusCode::UNAUTHORIZED, "API key has expired".to_string())
281            }
282            AuthError::TenantSuspended => (
283                StatusCode::FORBIDDEN,
284                "Tenant account is suspended".to_string(),
285            ),
286            AuthError::JwtNotConfigured => (
287                StatusCode::INTERNAL_SERVER_ERROR,
288                "JWT authentication not configured".to_string(),
289            ),
290            AuthError::UnsupportedAuthMethod => (
291                StatusCode::UNAUTHORIZED,
292                "Unsupported authentication method".to_string(),
293            ),
294            AuthError::Database(_) => (
295                StatusCode::INTERNAL_SERVER_ERROR,
296                "Internal server error".to_string(),
297            ),
298        };
299
300        let body = serde_json::json!({
301            "success": false,
302            "error": message
303        });
304
305        (status, axum::Json(body)).into_response()
306    }
307}
308
309/// Extractor for AuthContext from request extensions.
310#[derive(Clone, Debug)]
311pub struct Auth(pub AuthContext);
312
313impl<S> axum::extract::FromRequestParts<S> for Auth
314where
315    S: Send + Sync,
316{
317    type Rejection = AuthError;
318
319    async fn from_request_parts(
320        parts: &mut axum::http::request::Parts,
321        _state: &S,
322    ) -> Result<Self, Self::Rejection> {
323        parts
324            .extensions
325            .get::<AuthContext>()
326            .cloned()
327            .map(Auth)
328            .ok_or(AuthError::MissingHeader)
329    }
330}
331
332#[cfg(test)]
333mod tests {
334    use super::*;
335    use base64::Engine;
336    use jsonwebtoken::{encode, EncodingKey, Header};
337
338    const SECRET: &str = "test-secret";
339    const TENANT_ID: &str = "6ba7b810-9dad-11d1-80b4-00c04fd430c8";
340
341    fn sign(claims: &serde_json::Value, secret: &str) -> String {
342        encode(
343            &Header::default(),
344            claims,
345            &EncodingKey::from_secret(secret.as_bytes()),
346        )
347        .unwrap()
348    }
349
350    fn valid_claims() -> serde_json::Value {
351        serde_json::json!({
352            "sub": "user-1",
353            "tenant_id": TENANT_ID,
354            "role": "admin",
355            "scopes": ["domains:read"],
356            "exp": chrono::Utc::now().timestamp() + 3600,
357        })
358    }
359
360    fn b64url(data: &[u8]) -> String {
361        base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(data)
362    }
363
364    #[test]
365    fn valid_token_decodes() {
366        let token = sign(&valid_claims(), SECRET);
367        let claims = decode_jwt(&token, SECRET).unwrap();
368        assert_eq!(claims.sub, "user-1");
369        assert_eq!(claims.tenant_id, Some(TENANT_ID.parse().unwrap()));
370        assert_eq!(claims.role.as_deref(), Some("admin"));
371        assert_eq!(claims.scopes, Some(vec!["domains:read".to_string()]));
372    }
373
374    #[test]
375    fn token_without_exp_is_accepted() {
376        let mut claims = valid_claims();
377        claims.as_object_mut().unwrap().remove("exp");
378        let token = sign(&claims, SECRET);
379        assert!(decode_jwt(&token, SECRET).is_ok());
380    }
381
382    #[test]
383    fn token_signed_with_wrong_secret_is_rejected() {
384        let token = sign(&valid_claims(), "some-other-secret");
385        assert!(matches!(
386            decode_jwt(&token, SECRET),
387            Err(AuthError::InvalidToken(_))
388        ));
389    }
390
391    #[test]
392    fn tampered_payload_is_rejected() {
393        // Re-encode the payload with an escalated tenant_id while keeping the
394        // original (now mismatched) signature.
395        let token = sign(&valid_claims(), SECRET);
396        let parts: Vec<&str> = token.split('.').collect();
397
398        let mut claims = valid_claims();
399        claims["tenant_id"] = serde_json::json!("00000000-0000-0000-0000-000000000001");
400        let forged_payload = b64url(serde_json::to_vec(&claims).unwrap().as_slice());
401
402        let forged = format!("{}.{}.{}", parts[0], forged_payload, parts[2]);
403        assert!(matches!(
404            decode_jwt(&forged, SECRET),
405            Err(AuthError::InvalidToken(_))
406        ));
407    }
408
409    #[test]
410    fn unsigned_alg_none_token_is_rejected() {
411        // A classic "alg": "none" forgery must not pass HS256 validation,
412        // with or without a trailing signature segment.
413        let header = b64url(br#"{"alg":"none","typ":"JWT"}"#);
414        let payload = b64url(serde_json::to_vec(&valid_claims()).unwrap().as_slice());
415
416        for token in [
417            format!("{header}.{payload}."),
418            format!("{header}.{payload}"),
419        ] {
420            assert!(
421                matches!(decode_jwt(&token, SECRET), Err(AuthError::InvalidToken(_))),
422                "token {token:?} should be rejected"
423            );
424        }
425    }
426
427    #[test]
428    fn expired_token_is_rejected() {
429        let mut claims = valid_claims();
430        claims["exp"] = serde_json::json!(chrono::Utc::now().timestamp() - 60);
431        let token = sign(&claims, SECRET);
432        assert!(matches!(
433            decode_jwt(&token, SECRET),
434            Err(AuthError::TokenExpired)
435        ));
436    }
437
438    #[test]
439    fn malformed_token_is_rejected() {
440        assert!(matches!(
441            decode_jwt("not-a-jwt", SECRET),
442            Err(AuthError::InvalidToken(_))
443        ));
444    }
445}
postrust_proxy/saas/auth.rs

postrust_proxy/saas/
auth.rs
