madbase/auth/src/middleware.rs

use axum::{
    extract::{Request, State},
    http::StatusCode,
    middleware::Next,
    response::Response,
};
use common::{Config, ProjectContext};
use jsonwebtoken::{decode, Algorithm, DecodingKey, Validation};
use serde::{Deserialize, Serialize};

#[derive(Clone)]
pub struct AuthMiddlewareState {
    pub config: Config,
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct Claims {
    pub sub: String,
    pub email: Option<String>,
    pub role: String,
    pub exp: usize,
    pub iss: String,
    pub aud: Option<String>,
}

#[derive(Clone)]
pub struct AuthContext {
    pub claims: Option<Claims>,
    pub role: String,
}

pub async fn auth_middleware(
    State(state): State<AuthMiddlewareState>,
    mut req: Request,
    next: Next,
) -> Result<Response, StatusCode> {
    // 1. Try to get ProjectContext (if available)
    // If we are running in multi-tenant mode, ProjectContext should be present.
    // If not, we fall back to global config (legacy/single-tenant).
    let project_ctx = req.extensions().get::<ProjectContext>().cloned();

    // Allow public OAuth routes
    let path = req.uri().path();
    if path.contains("/authorize") || path.contains("/callback") {
         return Ok(next.run(req).await);
    }
    
    // Allow public Signed URL access (GET only)
    if path.contains("/object/sign/") && req.method() == axum::http::Method::GET {
        return Ok(next.run(req).await);
    }

    // Determine the secret to use
    let jwt_secret = if let Some(ctx) = &project_ctx {
        tracing::debug!("Using project-specific JWT secret");
        ctx.jwt_secret.clone()
    } else {
        tracing::debug!("ProjectContext not found, using global JWT secret");
        state.config.jwt_secret.clone()
    };

    let auth_header = req
        .headers()
        .get("Authorization")
        .and_then(|h| h.to_str().ok())
        .map(|s| s.to_string());

    let apikey_header = req
        .headers()
        .get("apikey")
        .and_then(|h| h.to_str().ok())
        .map(|s| s.to_string());

    // Logic:
    // 1. Bearer Token takes precedence for identity (Claims).
    // 2. API Key is checked if no Bearer token, OR it acts as the "Client Key" (anon/service).
    // Usually Supabase requires 'apikey' header ALWAYS, and Authorization header OPTIONAL (for user context).

    let token = if let Some(auth) = auth_header {
        auth.strip_prefix("Bearer ").map(|t| t.to_string())
    } else {
        // If no Auth header, check apikey header as fallback (e.g. for anon requests)
        apikey_header.clone()
    };

    if let Some(token) = token {
        let mut validation = Validation::new(Algorithm::HS256);
        validation.validate_exp = true;
        validation.validate_aud = false;
        // validation.set_audience(&["authenticated"]); // If we used audience

        match decode::<Claims>(
            &token,
            &DecodingKey::from_secret(jwt_secret.as_bytes()),
            &validation,
        ) {
            Ok(token_data) => {
                let claims = token_data.claims;
                let role = claims.role.clone();

                let ctx = AuthContext {
                    claims: Some(claims),
                    role,
                };
                req.extensions_mut().insert(ctx);
                return Ok(next.run(req).await);
            }
            Err(e) => {
                // Invalid token
                tracing::error!("Token validation failed: {}", e);
                return Err(StatusCode::UNAUTHORIZED);
            }
        }
    }

    // No valid token found.
    // Assign "anon" role if apikey is valid anon key?
    // Or just default to "anon" role without claims?
    // Supabase usually requires a valid JWT even for anon. The 'anon key' IS a JWT with role='anon'.

    // So if decoding failed above, we returned Unauthorized.
    // If no header provided at all?
    // We should allow public routes to proceed?
    // But this middleware is applied to ALL routes in /rest, /auth etc.
    // /auth/v1/signup needs to be accessible.
    // But wait, even signup requires the 'anon' key in Supabase.

    // So: strict check.
    Err(StatusCode::UNAUTHORIZED)
}