AkurAI-IDP / src/middleware/auth.rs

Latest change 6d386d9876a747c093481f05e926b1d43ccfb711 - fix: remove orphaned auth helper (AKURAI-IDP-11) by Ólafur Búi Ólafsson

use axum::{
    body::Body,
    http::{Request, Response, StatusCode, header},
    middleware::Next,
    response::{IntoResponse, Redirect},
};
use rusqlite::OptionalExtension;
use serde_json::json;
use std::time::{SystemTime, UNIX_EPOCH};

use crate::config;
use crate::db::with_db;
use crate::lib::crypto::constant_time_equal;
use crate::lib::jwt::verify_active_token;

/// User info attached to request extensions after auth.
#[derive(Clone, Debug)]
pub struct AuthUser {
    pub id: String,
    /// Legacy compatibility alias for the canonical organization_id.
    pub tenant_id: String,
    pub organization_id: String,
    pub role: String,
    pub authorization_version: i64,
    pub email: String,
    pub email_verified: bool,
}

#[derive(Clone, Debug)]
pub struct AuthSession {
    pub session_id: String,
    pub user_id: String,
    pub active_organization_id: String,
}

#[derive(Clone, Debug)]
pub struct AuthMachine {
    pub client_id: String,
    pub organization_id: String,
}

// ---------------------------------------------------------------------------
// 1. Session auth — cookie-based, redirects to /login on failure
// ---------------------------------------------------------------------------

pub async fn session_auth(request: Request<Body>, next: Next) -> Response<Body> {
    let session_id = match get_cookie(&request, "idp_session") {
        Some(s) => s,
        None => return Redirect::to("/login").into_response(),
    };

    let now = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap()
        .as_secs() as i64;
    let result: rusqlite::Result<Option<(AuthSession, AuthUser)>> = with_db(|conn| {
        let mut stmt = conn.prepare(
            "SELECT s.id, s.user_id, s.authorization_version,
                    COALESCE(s.active_organization_id, u.tenant_id),
                    u.authorization_version, u.email, u.email_verified, m.role
             FROM sessions s
             JOIN users u ON u.id = s.user_id
             JOIN organization_memberships m
               ON m.organization_id = COALESCE(s.active_organization_id, u.tenant_id)
              AND m.user_id = u.id AND m.status = 'active'
             WHERE s.id = ?1 AND s.expires_at > ?2
               AND s.authorization_version = u.authorization_version",
        )?;
        stmt.query_row(rusqlite::params![session_id, now], |row| {
            Ok((
                AuthSession {
                    session_id: row.get(0)?,
                    user_id: row.get(1)?,
                    active_organization_id: row.get(3)?,
                },
                AuthUser {
                    id: row.get::<_, String>(1)?,
                    tenant_id: row.get(3)?,
                    organization_id: row.get(3)?,
                    role: row.get(7)?,
                    authorization_version: row.get(4)?,
                    email: row.get(5)?,
                    email_verified: row.get::<_, i64>(6)? != 0,
                },
            ))
        })
        .optional()
    });

    match result {
        Ok(Some((session, user))) => {
            let mut request = request;
            request.extensions_mut().insert(session);
            request.extensions_mut().insert(user);
            next.run(request).await
        }
        Ok(None) => Redirect::to("/login").into_response(),
        Err(error) => {
            tracing::error!(%error, "session authorization database failure");
            (StatusCode::INTERNAL_SERVER_ERROR, "internal server error").into_response()
        }
    }
}

// ---------------------------------------------------------------------------
// 1b. Optional session — populates extensions if valid, always continues
// ---------------------------------------------------------------------------

pub async fn session_optional(request: Request<Body>, next: Next) -> Response<Body> {
    let session_id = match get_cookie(&request, "idp_session") {
        Some(s) => s,
        None => return next.run(request).await,
    };

    let now = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap()
        .as_secs() as i64;
    let result: rusqlite::Result<Option<(AuthSession, AuthUser)>> = with_db(|conn| {
        let mut stmt = conn.prepare(
            "SELECT s.id, s.user_id, s.authorization_version,
                    COALESCE(s.active_organization_id, u.tenant_id),
                    u.authorization_version, u.email, u.email_verified, m.role
             FROM sessions s
             JOIN users u ON u.id = s.user_id
             JOIN organization_memberships m
               ON m.organization_id = COALESCE(s.active_organization_id, u.tenant_id)
              AND m.user_id = u.id AND m.status = 'active'
             WHERE s.id = ?1 AND s.expires_at > ?2
               AND s.authorization_version = u.authorization_version",
        )?;
        stmt.query_row(rusqlite::params![session_id, now], |row| {
            Ok((
                AuthSession {
                    session_id: row.get(0)?,
                    user_id: row.get(1)?,
                    active_organization_id: row.get(3)?,
                },
                AuthUser {
                    id: row.get::<_, String>(1)?,
                    tenant_id: row.get(3)?,
                    organization_id: row.get(3)?,
                    role: row.get(7)?,
                    authorization_version: row.get(4)?,
                    email: row.get(5)?,
                    email_verified: row.get::<_, i64>(6)? != 0,
                },
            ))
        })
        .optional()
    });

    let mut request = request;
    match result {
        Ok(Some((session, user))) => {
            request.extensions_mut().insert(session);
            request.extensions_mut().insert(user);
        }
        Ok(None) => {}
        Err(error) => tracing::error!(%error, "optional session database failure"),
    }
    next.run(request).await
}

// ---------------------------------------------------------------------------
// 2. Bearer auth — header-based, for API endpoints
// ---------------------------------------------------------------------------

pub async fn bearer_auth(request: Request<Body>, next: Next) -> Response<Body> {
    let token = match extract_bearer(&request) {
        Some(t) => t,
        None => return unauthorized("Missing or invalid Authorization header"),
    };

    let cfg = config::get();
    if !cfg.admin_token.is_empty()
        && constant_time_equal(&token, &cfg.admin_token)
        && request.uri().path().starts_with("/organizations")
    {
        return unauthorized("User bearer token required");
    }
    if !cfg.admin_token.is_empty() && constant_time_equal(&token, &cfg.admin_token) {
        return next.run(request).await;
    }
    let claims = match verify_active_token(&token) {
        Some(claims) => claims,
        None => return unauthorized("Invalid or expired token"),
    };

    if claims.get("token_type").and_then(|v| v.as_str()) == Some("client_credentials") {
        let client_id = claims
            .get("client_id")
            .or_else(|| claims.get("sub"))
            .and_then(|v| v.as_str())
            .unwrap_or_default()
            .to_string();
        let organization_id = claims
            .get("organization_id")
            .or_else(|| claims.get("tenant_id"))
            .and_then(|v| v.as_str())
            .unwrap_or_default()
            .to_string();
        let valid = !client_id.is_empty()
            && !organization_id.is_empty()
            && claims.get("aud").and_then(|v| v.as_str()) == Some(client_id.as_str())
            && with_db(|conn| {
                conn.query_row(
                    "SELECT COUNT(*) FROM clients c
                     JOIN organizations o ON o.id = c.tenant_id
                     WHERE c.id = ?1 AND o.id = ?2",
                    rusqlite::params![client_id, organization_id],
                    |row| row.get::<_, i64>(0),
                )
                .map(|count| count == 1)
                .unwrap_or(false)
            });
        if !valid {
            return unauthorized("Invalid or expired token");
        }
        if request.uri().path().starts_with("/organizations") {
            return unauthorized("User bearer token required");
        }
        let mut request = request;
        request.extensions_mut().insert(AuthMachine {
            client_id,
            organization_id,
        });
        return next.run(request).await;
    }

    let user = (|| {
        let id = claims.get("sub")?.as_str()?.to_string();
        let organization_id = claims
            .get("organization_id")
            .or_else(|| claims.get("tenant_id"))?
            .as_str()?
            .to_string();
        let claims_version = claims.get("claims_version")?.as_i64()?;
        let authorization_version = claims.get("authorization_version")?.as_i64()?;
        if claims_version != crate::lib::organizations::CLAIMS_VERSION {
            return None;
        }
        with_db(|conn| {
            let membership =
                crate::lib::organizations::active_membership(conn, &id, &organization_id)
                    .ok()??;
            if membership.authorization_version != authorization_version {
                return None;
            }
            let (email, email_verified): (String, bool) = conn
                .query_row(
                    "SELECT email, email_verified FROM users WHERE id = ?1",
                    [&id],
                    |row| Ok((row.get(0)?, row.get::<_, i64>(1)? != 0)),
                )
                .ok()?;
            Some(AuthUser {
                id,
                tenant_id: organization_id.clone(),
                organization_id,
                role: membership.role,
                authorization_version,
                email,
                email_verified,
            })
        })
    })();

    match user {
        Some(u) => {
            let mut request = request;
            request.extensions_mut().insert(u);
            next.run(request).await
        }
        None => unauthorized("Invalid or expired token"),
    }
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

fn extract_bearer(request: &Request<Body>) -> Option<String> {
    let header = request.headers().get(header::AUTHORIZATION)?;
    let value = header.to_str().ok()?;
    let token = value.strip_prefix("Bearer ")?;
    if token.is_empty() {
        return None;
    }
    Some(token.to_string())
}

fn get_cookie(request: &Request<Body>, name: &str) -> Option<String> {
    let cookies = request.headers().get(header::COOKIE)?.to_str().ok()?;
    let prefix = format!("{name}=");
    for part in cookies.split(';') {
        let part = part.trim();
        if let Some(value) = part.strip_prefix(&prefix) {
            return Some(value.to_string());
        }
    }
    None
}

fn unauthorized(msg: &str) -> Response<Body> {
    let body = json!({"error": "unauthorized", "error_description": msg});
    (StatusCode::UNAUTHORIZED, axum::Json(body)).into_response()
}