eric / rustwarden

use axum::{
    Json,
    http::StatusCode,
    response::{IntoResponse, Response},
};
use serde_json::json;
use thiserror::Error;

/// Application error types
///
/// This module defines all error types used throughout the application.
/// Errors are converted to appropriate HTTP responses for API endpoints.

#[derive(Error, Debug)]
pub enum AppError {
    /// A SeaORM database operation failed.
    #[error("Database error: {0}")]
    Database(#[from] sea_orm::DbErr),

    /// User-supplied input failed validation. The message describes what was invalid.
    #[error("Validation error: {0}")]
    Validation(String),

    /// The requested resource could not be found.
    #[error("Resource not found: {0}")]
    NotFound(String),

    /// An unexpected internal error occurred. The message contains diagnostic detail.
    #[error("Internal server error: {0}")]
    Internal(String),

    /// The supplied username/password combination is incorrect.
    #[error("Invalid credentials")]
    InvalidCredentials,

    /// The supplied token is malformed or does not match any known record.
    #[error("Invalid token")]
    InvalidToken,

    /// The token was valid but its expiration time has passed.
    #[error("Token expired")]
    TokenExpired,

    /// A refresh token was presented that had already been used, indicating possible token theft.
    #[error("Token reused - possible security breach")]
    TokenReused,

    /// The user must set a new password before accessing the application.
    #[allow(dead_code)]
    #[error("Password change required")]
    MustChangePassword,

    /// The request requires authentication but no valid credentials were provided.
    #[allow(dead_code)]
    #[error("Authentication required")]
    Unauthorized,

    /// The authenticated user does not have permission to perform this action.
    #[allow(dead_code)]
    #[error("Insufficient permissions")]
    Forbidden,

    /// A required configuration value is missing or invalid.
    #[error("Configuration error: {0}")]
    Config(String),

    /// An Askama template failed to render.
    #[error("Template rendering error: {0}")]
    Template(String),

    /// The submitted TOTP code is incorrect.
    #[error("TOTP verification failed")]
    TotpFailed,

    /// The TOTP challenge token has expired or does not exist.
    #[error("TOTP challenge expired or invalid")]
    TotpChallengeInvalid,

    /// The TOTP challenge was failed too many times and is now locked.
    #[error("TOTP challenge max attempts exceeded")]
    TotpMaxAttemptsExceeded,

    /// The account has TOTP enabled; a second-factor code must be provided to continue.
    #[error("Two-factor authentication required")]
    TotpRequired,

    /// TOTP cannot be enabled because it is already active for this user.
    #[allow(dead_code)]
    #[error("TOTP already enabled")]
    TotpAlreadyEnabled,

    /// The requested operation requires TOTP to be enabled, but it is not.
    #[allow(dead_code)]
    #[error("TOTP not enabled")]
    TotpNotEnabled,

    /// Client exceeded a rate limit. Carries the `Retry-After` header value (seconds).
    #[error("Too many requests, retry after {retry_after_secs}s")]
    TooManyRequests {
        /// Number of seconds the client should wait before retrying.
        retry_after_secs: u64,
    },
}

impl IntoResponse for AppError {
    fn into_response(self) -> Response {
        // 429 has a required Retry-After header, so handle it before the generic
        // (status, body) path used by other variants.
        if let AppError::TooManyRequests { retry_after_secs } = self {
            tracing::warn!(
                retry_after_secs,
                "Rate limit exceeded, returning 429 Too Many Requests"
            );
            let body = Json(json!({
                "error": "Too many requests, please retry later",
                "status": StatusCode::TOO_MANY_REQUESTS.as_u16(),
                "retry_after_secs": retry_after_secs,
            }));
            return (
                StatusCode::TOO_MANY_REQUESTS,
                [("Retry-After", retry_after_secs.to_string())],
                body,
            )
                .into_response();
        }

        let (status, error_message) = match self {
            AppError::Database(err) => {
                tracing::error!("Database error: {}", err);
                (
                    StatusCode::INTERNAL_SERVER_ERROR,
                    "Database error occurred".to_string(),
                )
            }
            AppError::Validation(msg) => {
                tracing::warn!("Validation error: {}", msg);
                (
                    StatusCode::BAD_REQUEST,
                    format!("Validation error: {}", msg),
                )
            }
            AppError::NotFound(msg) => {
                tracing::info!("Not found: {}", msg);
                (
                    StatusCode::NOT_FOUND,
                    format!("Resource not found: {}", msg),
                )
            }
            AppError::Internal(msg) => {
                tracing::error!("Internal error: {}", msg);
                (
                    StatusCode::INTERNAL_SERVER_ERROR,
                    "Internal server error".to_string(),
                )
            }
            AppError::InvalidCredentials => {
                tracing::warn!("Authentication attempt with invalid credentials");
                (StatusCode::UNAUTHORIZED, "Invalid credentials".to_string())
            }
            AppError::InvalidToken => {
                tracing::warn!("Invalid token provided");
                (StatusCode::UNAUTHORIZED, "Invalid token".to_string())
            }
            AppError::TokenExpired => {
                tracing::debug!("Token expired");
                (StatusCode::UNAUTHORIZED, "Token expired".to_string())
            }
            AppError::TokenReused => {
                tracing::error!("Refresh token reuse detected - possible security breach");
                (
                    StatusCode::UNAUTHORIZED,
                    "Token reused - please re-authenticate".to_string(),
                )
            }
            AppError::MustChangePassword => {
                tracing::info!("User must change password");
                (
                    StatusCode::FORBIDDEN,
                    "Password change required".to_string(),
                )
            }
            AppError::Unauthorized => {
                tracing::warn!("Unauthorized access attempt");
                (
                    StatusCode::UNAUTHORIZED,
                    "Authentication required".to_string(),
                )
            }
            AppError::Forbidden => {
                tracing::warn!("Forbidden access attempt");
                (
                    StatusCode::FORBIDDEN,
                    "Insufficient permissions".to_string(),
                )
            }
            AppError::Config(msg) => {
                tracing::error!("Configuration error: {}", msg);
                (
                    StatusCode::INTERNAL_SERVER_ERROR,
                    "Server configuration error".to_string(),
                )
            }
            AppError::Template(msg) => {
                tracing::error!("Template rendering error: {}", msg);
                (
                    StatusCode::INTERNAL_SERVER_ERROR,
                    "Template error".to_string(),
                )
            }
            AppError::TotpFailed => {
                tracing::warn!("TOTP verification failed");
                (
                    StatusCode::UNAUTHORIZED,
                    "TOTP verification failed".to_string(),
                )
            }
            AppError::TotpChallengeInvalid => {
                tracing::warn!("TOTP challenge expired or invalid");
                (
                    StatusCode::UNAUTHORIZED,
                    "TOTP challenge expired or invalid".to_string(),
                )
            }
            AppError::TotpMaxAttemptsExceeded => {
                tracing::warn!("TOTP challenge max attempts exceeded");
                (
                    StatusCode::TOO_MANY_REQUESTS,
                    "Too many TOTP attempts".to_string(),
                )
            }
            AppError::TotpRequired => {
                tracing::info!("Two-factor authentication setup required");
                (
                    StatusCode::FORBIDDEN,
                    "Two-factor authentication required".to_string(),
                )
            }
            AppError::TotpAlreadyEnabled => {
                tracing::info!("TOTP already enabled for user");
                (
                    StatusCode::BAD_REQUEST,
                    "TOTP is already enabled".to_string(),
                )
            }
            AppError::TotpNotEnabled => {
                tracing::info!("TOTP not enabled for user");
                (StatusCode::BAD_REQUEST, "TOTP is not enabled".to_string())
            }
            AppError::TooManyRequests { .. } => {
                unreachable!("TooManyRequests is handled above to include the Retry-After header")
            }
        };

        let body = Json(json!({
            "error": error_message,
            "status": status.as_u16()
        }));

        (status, body).into_response()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use axum::body::to_bytes;

    #[tokio::test]
    async fn too_many_requests_returns_429_with_retry_after_header() {
        let err = AppError::TooManyRequests {
            retry_after_secs: 42,
        };
        let response = err.into_response();

        assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
        let retry_after = response
            .headers()
            .get("Retry-After")
            .expect("Retry-After header must be set for 429 responses");
        assert_eq!(retry_after.to_str().unwrap(), "42");
    }

    #[tokio::test]
    async fn too_many_requests_body_contains_error_message() {
        let err = AppError::TooManyRequests {
            retry_after_secs: 1,
        };
        let response = err.into_response();
        let bytes = to_bytes(response.into_body(), 1024).await.unwrap();
        let body: serde_json::Value = serde_json::from_slice(&bytes).unwrap();

        assert_eq!(body["status"], 429);
        assert!(
            body["error"]
                .as_str()
                .unwrap()
                .to_lowercase()
                .contains("too many")
        );
    }
}