pinakes/crates/pinakes-server/src/auth.rs

use axum::{
  extract::{Request, State},
  http::StatusCode,
  middleware::Next,
  response::{IntoResponse, Response},
};
use pinakes_core::config::UserRole;

use crate::state::AppState;

/// Constant-time string comparison to prevent timing attacks on API keys.
///
/// Always iterates to `max(len_a, len_b)` so that neither a length difference
/// nor a byte mismatch causes an early return.
fn constant_time_eq(a: &str, b: &str) -> bool {
  let a = a.as_bytes();
  let b = b.as_bytes();
  let len = a.len().max(b.len());
  let mut result = a.len() ^ b.len(); // non-zero if lengths differ
  for i in 0..len {
    let ab = a.get(i).copied().unwrap_or(0);
    let bb = b.get(i).copied().unwrap_or(0);
    result |= usize::from(ab ^ bb);
  }
  result == 0
}

/// Axum middleware that checks for a valid Bearer token.
///
/// If `accounts.enabled == true`: look up bearer token in database session
/// store. If `accounts.enabled == false`: use existing `api_key` logic
/// (unchanged behavior). Skips authentication for the `/health` and
/// `/auth/login` path suffixes.
pub async fn require_auth(
  State(state): State<AppState>,
  mut request: Request,
  next: Next,
) -> Response {
  let path = request.uri().path().to_string();

  // Always allow health and login endpoints
  if path.ends_with("/health") || path.ends_with("/auth/login") {
    return next.run(request).await;
  }

  let config = state.config.read().await;

  // Check if authentication is explicitly disabled
  if config.server.authentication_disabled {
    drop(config);
    tracing::warn!("authentication is disabled - allowing all requests");
    request.extensions_mut().insert(UserRole::Admin);
    request.extensions_mut().insert("admin".to_string());
    return next.run(request).await;
  }

  if config.accounts.enabled {
    drop(config);

    // Session-based auth using database
    let token = request
      .headers()
      .get("authorization")
      .and_then(|v| v.to_str().ok())
      .and_then(|s| s.strip_prefix("Bearer "))
      .map(std::string::ToString::to_string);

    let Some(token) = token else {
      tracing::debug!(path = %path, "rejected: missing Authorization header");
      return unauthorized("missing Authorization header");
    };

    // Look up session in database
    let session_result = state.storage.get_session(&token).await;
    let session = match session_result {
      Ok(Some(session)) => session,
      Ok(None) => {
        tracing::debug!(path = %path, "rejected: invalid session token");
        return unauthorized("invalid or expired session token");
      },
      Err(e) => {
        tracing::error!(error = %e, "failed to query session from database");
        return (StatusCode::INTERNAL_SERVER_ERROR, "database error")
          .into_response();
      },
    };

    // Check session expiry
    let now = chrono::Utc::now();
    if session.expires_at < now {
      let username = session.username;
      // Delete expired session in a bounded background task
      if let Ok(permit) = state.session_semaphore.clone().try_acquire_owned() {
        let storage = state.storage.clone();
        let token_owned = token.clone();
        tokio::spawn(async move {
          if let Err(e) = storage.delete_session(&token_owned).await {
            tracing::error!(error = %e, "failed to delete expired session");
          }
          drop(permit);
        });
      }
      tracing::info!(username = %username, "session expired");
      return unauthorized("session expired");
    }

    // Update last_accessed timestamp in a bounded background task
    if let Ok(permit) = state.session_semaphore.clone().try_acquire_owned() {
      let storage = state.storage.clone();
      let token_owned = token.clone();
      tokio::spawn(async move {
        if let Err(e) = storage.touch_session(&token_owned).await {
          tracing::warn!(error = %e, "failed to update session last_accessed");
        }
        drop(permit);
      });
    }

    // Parse role from string
    let role = match session.role.as_str() {
      "admin" => UserRole::Admin,
      "editor" => UserRole::Editor,
      "viewer" => UserRole::Viewer,
      _ => {
        tracing::warn!(role = %session.role, "unknown role, defaulting to viewer");
        UserRole::Viewer
      },
    };

    // Inject role and username into request extensions
    request.extensions_mut().insert(role);
    request.extensions_mut().insert(session.username);
  } else {
    // Legacy API key auth
    let api_key = std::env::var("PINAKES_API_KEY")
      .ok()
      .or_else(|| config.server.api_key.clone());
    drop(config);

    let Some(ref expected_key) = api_key else {
      tracing::error!("no authentication configured");
      return unauthorized("authentication not configured");
    };

    if expected_key.is_empty() {
      // Empty key is not allowed - must use authentication_disabled flag
      tracing::error!(
        "empty api_key rejected, use authentication_disabled flag instead"
      );
      return unauthorized("authentication not properly configured");
    }

    let auth_header = request
      .headers()
      .get("authorization")
      .and_then(|v| v.to_str().ok());

    match auth_header {
      Some(header) if header.starts_with("Bearer ") => {
        let token = &header[7..];
        if !constant_time_eq(token, expected_key.as_str()) {
          tracing::warn!(path = %path, "rejected: invalid API key");
          return unauthorized("invalid api key");
        }
      },
      _ => {
        return unauthorized(
          "missing or malformed Authorization header, expected: Bearer \
           <api_key>",
        );
      },
    }

    // API key matches, grant admin
    request.extensions_mut().insert(UserRole::Admin);
    request.extensions_mut().insert("admin".to_string());
  }

  next.run(request).await
}

/// Middleware: requires Editor or Admin role.
pub async fn require_editor(request: Request, next: Next) -> Response {
  let role = request
    .extensions()
    .get::<UserRole>()
    .copied()
    .unwrap_or(UserRole::Viewer);
  if role.can_write() {
    next.run(request).await
  } else {
    forbidden("editor role required")
  }
}

/// Middleware: requires Admin role.
pub async fn require_admin(request: Request, next: Next) -> Response {
  let role = request
    .extensions()
    .get::<UserRole>()
    .copied()
    .unwrap_or(UserRole::Viewer);
  if role.can_admin() {
    next.run(request).await
  } else {
    forbidden("admin role required")
  }
}

/// Resolve the authenticated username (from request extensions) to a `UserId`.
///
/// Returns an error if the user cannot be found.
pub async fn resolve_user_id(
  storage: &pinakes_core::storage::DynStorageBackend,
  username: &str,
) -> Result<pinakes_core::users::UserId, crate::error::ApiError> {
  match storage.get_user_by_username(username).await {
    Ok(user) => Ok(user.id),
    Err(e) => {
      tracing::warn!(username = %username, error = ?e, "failed to resolve user");
      Err(crate::error::ApiError(
        pinakes_core::error::PinakesError::Authentication(
          "user not found".into(),
        ),
      ))
    },
  }
}

fn unauthorized(message: &str) -> Response {
  let body = format!(r#"{{"error":"{message}"}}"#);
  (
    StatusCode::UNAUTHORIZED,
    [("content-type", "application/json")],
    body,
  )
    .into_response()
}

fn forbidden(message: &str) -> Response {
  let body = format!(r#"{{"error":"{message}"}}"#);
  (
    StatusCode::FORBIDDEN,
    [("content-type", "application/json")],
    body,
  )
    .into_response()
}