use axum::{
extract::{FromRequestParts, TypedHeader},
headers::{authorization::Bearer, Authorization},
http::{request::Parts, StatusCode},
response::{IntoResponse, Response},
Json,
};
use serde::Serialize;
use std::collections::HashSet;
use std::sync::{Arc, RwLock};
#[derive(Debug, Clone)]
pub struct AuthUser {
pub token: String,
}
#[derive(Clone)]
pub struct AuthConfig {
tokens: Arc<RwLock<HashSet<String>>>,
}
impl AuthConfig {
pub fn new() -> Self {
Self {
tokens: Arc::new(RwLock::new(HashSet::new())),
}
}
pub fn add_token(&self, token: String) {
if let Ok(mut tokens) = self.tokens.write() {
tokens.insert(token);
}
}
pub fn is_valid_token(&self, token: &str) -> bool {
if let Ok(tokens) = self.tokens.read() {
tokens.contains(token)
} else {
false
}
}
}
impl Default for AuthConfig {
fn default() -> Self {
Self::new()
}
}
#[derive(Serialize)]
pub struct AuthError {
message: String,
}
impl IntoResponse for AuthError {
fn into_response(self) -> Response {
(StatusCode::UNAUTHORIZED, Json(self)).into_response()
}
}
#[async_trait::async_trait]
impl<S> FromRequestParts<S> for AuthUser
where
S: Send + Sync,
{
type Rejection = AuthError;
async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
if let Ok(TypedHeader(Authorization(bearer))) =
TypedHeader::<Authorization<Bearer>>::from_request_parts(parts, _state).await
{
Ok(AuthUser {
token: bearer.token().to_string(),
})
} else {
Err(AuthError {
message: "Token de autorização ausente ou inválido".into(),
})
}
}
}
#[async_trait::async_trait]
impl<S> FromRequestParts<S> for AuthConfig
where
S: Send + Sync,
{
type Rejection = AuthError;
async fn from_request_parts(_parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
Ok(AuthConfig::new())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_auth_config_add_and_validate_token() {
let config = AuthConfig::new();
config.add_token("test-token".to_string());
assert!(config.is_valid_token("test-token"));
assert!(!config.is_valid_token("invalid-token"));
}
#[test]
fn test_auth_config_multiple_tokens() {
let config = AuthConfig::new();
config.add_token("token1".to_string());
config.add_token("token2".to_string());
config.add_token("token3".to_string());
assert!(config.is_valid_token("token1"));
assert!(config.is_valid_token("token2"));
assert!(config.is_valid_token("token3"));
assert!(!config.is_valid_token("token4"));
}
#[test]
fn test_auth_error_into_response() {
let error = AuthError {
message: "Token inválido".to_string(),
};
let response = error.into_response();
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
}
}