synops/maskinrommet/src/auth.rs

use axum::{
    extract::FromRequestParts,
    http::{request::Parts, StatusCode},
    response::{IntoResponse, Response},
    Json,
};
use jsonwebtoken::{decode, decode_header, Algorithm, DecodingKey, Validation};
use serde::{Deserialize, Serialize};
use uuid::Uuid;

use crate::AppState;

// ---------------------------------------------------------------------------
// JWKS types
// ---------------------------------------------------------------------------

#[derive(Debug, Deserialize)]
pub struct JwksResponse {
    pub keys: Vec<JwkKey>,
}

#[derive(Debug, Clone, Deserialize)]
pub struct JwkKey {
    pub kid: Option<String>,
    pub kty: String,
    pub n: String,
    pub e: String,
}

/// Cached JWKS keys and OIDC config fetched from Authentik at startup.
#[derive(Debug, Clone)]
pub struct JwksKeys {
    pub keys: Vec<JwkKey>,
    pub issuer: String,
    pub audience: String,
}

impl JwksKeys {
    /// Fetch JWKS from Authentik's OIDC discovery endpoint.
    pub async fn fetch(issuer: &str, audience: &str) -> Result<Self, String> {
        let jwks_url = format!("{}jwks/", issuer.trim_end_matches('/').to_owned() + "/");
        tracing::info!("Henter JWKS fra {jwks_url}");

        let resp = reqwest::get(&jwks_url)
            .await
            .map_err(|e| format!("Kunne ikke hente JWKS: {e}"))?;

        let jwks: JwksResponse = resp
            .json()
            .await
            .map_err(|e| format!("Kunne ikke parse JWKS: {e}"))?;

        if jwks.keys.is_empty() {
            return Err("JWKS inneholder ingen nøkler".to_string());
        }

        tracing::info!("Lastet {} JWKS-nøkler", jwks.keys.len());

        Ok(Self {
            keys: jwks.keys,
            issuer: issuer.to_string(),
            audience: audience.to_string(),
        })
    }

    /// Find decoding key by kid, or use the first key if no kid in header.
    fn decoding_key(&self, kid: Option<&str>) -> Result<DecodingKey, String> {
        let key = match kid {
            Some(kid) => self
                .keys
                .iter()
                .find(|k| k.kid.as_deref() == Some(kid))
                .ok_or_else(|| format!("Ukjent kid: {kid}"))?,
            None => self.keys.first().ok_or("Ingen JWKS-nøkler tilgjengelig")?,
        };

        DecodingKey::from_rsa_components(&key.n, &key.e)
            .map_err(|e| format!("Ugyldig RSA-nøkkel: {e}"))
    }
}

// ---------------------------------------------------------------------------
// JWT claims
// ---------------------------------------------------------------------------

#[derive(Debug, Deserialize)]
struct Claims {
    sub: String,
    // iss and exp are validated by jsonwebtoken automatically
}

// ---------------------------------------------------------------------------
// AuthUser extractor
// ---------------------------------------------------------------------------

/// Authenticated user extracted from a valid JWT.
/// Use as an axum extractor on protected endpoints.
#[derive(Debug, Clone)]
pub struct AuthUser {
    pub node_id: Uuid,
    pub authentik_sub: String,
}

/// Error response for auth failures.
#[derive(Serialize)]
struct AuthError {
    error: String,
}

impl IntoResponse for AuthErrorKind {
    fn into_response(self) -> Response {
        let (status, message) = match self {
            AuthErrorKind::MissingToken => (StatusCode::UNAUTHORIZED, "Mangler Authorization-header"),
            AuthErrorKind::InvalidToken(ref _e) => (StatusCode::UNAUTHORIZED, "Ugyldig token"),
            AuthErrorKind::UnknownIdentity => {
                (StatusCode::UNAUTHORIZED, "Ukjent brukeridentitet")
            }
            AuthErrorKind::Internal => {
                (StatusCode::INTERNAL_SERVER_ERROR, "Intern feil ved autentisering")
            }
        };

        let body = Json(AuthError {
            error: message.to_string(),
        });

        (status, body).into_response()
    }
}

#[derive(Debug)]
pub enum AuthErrorKind {
    MissingToken,
    InvalidToken(String),
    UnknownIdentity,
    Internal,
}

impl<S> FromRequestParts<S> for AuthUser
where
    S: Send + Sync,
    AppState: FromRef<S>,
{
    type Rejection = AuthErrorKind;

    async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
        let app_state = AppState::from_ref(state);

        // Extract Bearer token from Authorization header
        let auth_header = parts
            .headers
            .get("authorization")
            .and_then(|v| v.to_str().ok())
            .ok_or(AuthErrorKind::MissingToken)?;

        let token = auth_header
            .strip_prefix("Bearer ")
            .ok_or(AuthErrorKind::MissingToken)?;

        // Decode header to get kid
        let header = decode_header(token)
            .map_err(|e| AuthErrorKind::InvalidToken(e.to_string()))?;

        // Get decoding key from JWKS
        let decoding_key = app_state
            .jwks
            .decoding_key(header.kid.as_deref())
            .map_err(|e| {
                tracing::warn!("JWKS-nøkkel ikke funnet: {e}");
                AuthErrorKind::InvalidToken(e)
            })?;

        // Validate JWT (signature, exp, iss, aud)
        let mut validation = Validation::new(Algorithm::RS256);
        validation.set_issuer(&[&app_state.jwks.issuer]);
        validation.set_audience(&[&app_state.jwks.audience]);

        let token_data = decode::<Claims>(token, &decoding_key, &validation).map_err(|e| {
            tracing::debug!("JWT-validering feilet: {e}");
            AuthErrorKind::InvalidToken(e.to_string())
        })?;

        let authentik_sub = token_data.claims.sub;

        // Look up auth_identities → node_id
        let row = sqlx::query_scalar::<_, Uuid>(
            "SELECT node_id FROM auth_identities WHERE authentik_sub = $1",
        )
        .bind(&authentik_sub)
        .fetch_optional(&app_state.db)
        .await
        .map_err(|e| {
            tracing::error!("DB-feil ved oppslag av auth_identities: {e}");
            AuthErrorKind::Internal
        })?;

        let node_id = row.ok_or_else(|| {
            tracing::warn!("Ingen auth_identity for sub={authentik_sub}");
            AuthErrorKind::UnknownIdentity
        })?;

        tracing::debug!("Autentisert bruker: node_id={node_id}, sub={authentik_sub}");

        Ok(AuthUser {
            node_id,
            authentik_sub,
        })
    }
}

// We need FromRef to extract AppState from the state
use axum::extract::FromRef;