use std::convert::Infallible;
use std::future::Future;
use std::marker::PhantomData;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use axum::extract::FromRequestParts;
use axum::response::{IntoResponse, Response};
use serde::de::DeserializeOwned;
use serde::{Deserialize, Serialize};
use tower_sessions::Session as TowerSession;
use crate::auth::AuthUser;
const ABSOLUTE_AUTH_AT_KEY: &str = "__arcature_absolute_auth_at";
fn now_unix_millis() -> i64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_millis() as i64)
.unwrap_or(0)
}
pub struct Auth<U: AuthUser>(pub U);
impl<U: AuthUser> Auth<U> {
#[must_use]
pub fn into_inner(self) -> U {
self.0
}
#[must_use]
pub fn user(&self) -> &U {
&self.0
}
pub fn authorize<M, P: Policy<M, User = U>>(
&self,
action: &str,
resource: &M,
) -> Result<(), AuthzError> {
if P::check(&self.0, action, resource) {
Ok(())
} else {
Err(AuthzError::Forbidden)
}
}
}
impl<U, S> FromRequestParts<S> for Auth<U>
where
U: UserLoader<S>,
S: Send + Sync,
{
type Rejection = Response;
async fn from_request_parts(
parts: &mut axum::http::request::Parts,
state: &S,
) -> Result<Self, Self::Rejection> {
let user = load_user::<U, S>(parts, state).await?;
user.map(Auth).ok_or_else(|| {
(
axum::http::StatusCode::UNAUTHORIZED,
"Authentication required",
)
.into_response()
})
}
}
pub struct OptionalAuth<U: AuthUser>(pub Option<U>);
impl<U: AuthUser> OptionalAuth<U> {
#[must_use]
pub fn user(&self) -> Option<&U> {
self.0.as_ref()
}
#[must_use]
pub fn is_authenticated(&self) -> bool {
self.0.is_some()
}
}
impl<U, S> FromRequestParts<S> for OptionalAuth<U>
where
U: UserLoader<S>,
S: Send + Sync,
{
type Rejection = Infallible;
async fn from_request_parts(
parts: &mut axum::http::request::Parts,
state: &S,
) -> Result<Self, Self::Rejection> {
let user = load_user::<U, S>(parts, state).await.unwrap_or(None);
Ok(OptionalAuth(user))
}
}
pub type Current<U> = Auth<U>;
pub type OptionalCurrent<U> = OptionalAuth<U>;
pub struct AuthManager<U: AuthUser> {
session: TowerSession,
_marker: PhantomData<U>,
}
impl<U: AuthUser> AuthManager<U> {
#[must_use]
pub fn login(&self, user: &U) -> LoginBuilder<'_, U> {
LoginBuilder {
session: &self.session,
user_id: user.id().clone(),
remember: false,
}
}
pub async fn logout(&self) -> Result<(), AuthError> {
self.session
.flush()
.await
.map_err(|e| AuthError::Session(e.to_string()))?;
Ok(())
}
pub async fn regenerate(&self) -> Result<(), AuthError> {
self.session
.cycle_id()
.await
.map_err(|e| AuthError::Session(e.to_string()))?;
Ok(())
}
}
impl<U, S> FromRequestParts<S> for AuthManager<U>
where
U: AuthUser,
S: Send + Sync,
{
type Rejection = Infallible;
async fn from_request_parts(
parts: &mut axum::http::request::Parts,
state: &S,
) -> Result<Self, Self::Rejection> {
let session = TowerSession::from_request_parts(parts, state)
.await
.map_err(|_| unreachable!("Session extraction is infallible"))?;
Ok(AuthManager {
session,
_marker: PhantomData,
})
}
}
pub struct LoginBuilder<'a, U: AuthUser> {
session: &'a TowerSession,
user_id: U::Id,
remember: bool,
}
impl<'a, U: AuthUser> LoginBuilder<'a, U> {
#[must_use]
pub fn remember(mut self, remember: bool) -> Self {
self.remember = remember;
self
}
}
impl<'a, U: AuthUser> std::future::IntoFuture for LoginBuilder<'a, U> {
type Output = Result<(), AuthError>;
type IntoFuture = std::pin::Pin<
std::boxed::Box<dyn std::future::Future<Output = Result<(), AuthError>> + Send + 'a>,
>;
fn into_future(self) -> Self::IntoFuture {
Box::pin(async move {
self.session
.cycle_id()
.await
.map_err(|e| AuthError::Session(e.to_string()))?;
self.session
.insert(U::SESSION_KEY, &self.user_id)
.await
.map_err(|e| AuthError::Session(e.to_string()))?;
self.session
.insert(ABSOLUTE_AUTH_AT_KEY, now_unix_millis())
.await
.map_err(|e| AuthError::Session(e.to_string()))?;
if self.remember {
self.session
.insert("remember_me", true)
.await
.map_err(|e| AuthError::Session(e.to_string()))?;
}
Ok(())
})
}
}
#[derive(Debug)]
pub enum AuthError {
Session(String),
}
impl std::fmt::Display for AuthError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Session(msg) => write!(f, "session error: {msg}"),
}
}
}
impl std::error::Error for AuthError {}
pub trait UserLoader<S>: AuthUser + Sized {
type Error: std::error::Error + Send + Sync + 'static;
fn load_user(
id: &Self::Id,
state: &S,
) -> impl Future<Output = Result<Option<Self>, Self::Error>> + Send;
#[must_use]
fn absolute_max_age() -> Duration {
Duration::from_secs(60 * 60 * 24 * 30)
}
}
#[allow(clippy::result_large_err)]
async fn load_user<U, S>(
parts: &mut axum::http::request::Parts,
state: &S,
) -> Result<Option<U>, Response>
where
U: UserLoader<S>,
S: Send + Sync,
{
let session = TowerSession::from_request_parts(parts, state)
.await
.map_err(|_| {
(
axum::http::StatusCode::INTERNAL_SERVER_ERROR,
"session extraction failed",
)
.into_response()
})?;
let user_id: Option<U::Id> = session.get(U::SESSION_KEY).await.map_err(|_err| {
(
axum::http::StatusCode::INTERNAL_SERVER_ERROR,
"session read failed",
)
.into_response()
})?;
let user_id = match user_id {
Some(id) => id,
None => return Ok(None),
};
let absolute_max_millis: i64 =
i64::try_from(U::absolute_max_age().as_millis()).unwrap_or(i64::MAX);
let auth_at: Option<i64> = session.get(ABSOLUTE_AUTH_AT_KEY).await.map_err(|_err| {
(
axum::http::StatusCode::INTERNAL_SERVER_ERROR,
"session read failed",
)
.into_response()
})?;
match auth_at {
Some(auth_at) => {
if now_unix_millis().saturating_sub(auth_at) > absolute_max_millis {
session.flush().await.map_err(|_err| {
(
axum::http::StatusCode::INTERNAL_SERVER_ERROR,
"session flush failed",
)
.into_response()
})?;
return Ok(None);
}
}
None => {
session
.insert(ABSOLUTE_AUTH_AT_KEY, now_unix_millis())
.await
.map_err(|_err| {
(
axum::http::StatusCode::INTERNAL_SERVER_ERROR,
"session write failed",
)
.into_response()
})?;
}
}
let user = U::load_user(&user_id, state).await.map_err(|_err| {
(
axum::http::StatusCode::INTERNAL_SERVER_ERROR,
"user load failed",
)
.into_response()
})?;
Ok(user)
}
pub struct Session(pub(crate) TowerSession);
impl Session {
pub async fn put<T: Serialize>(&self, key: &str, value: T) -> Result<(), SessionError> {
self.0
.insert(key, value)
.await
.map_err(|e| SessionError(e.to_string()))
}
pub async fn get<T: DeserializeOwned>(&self, key: &str) -> Result<Option<T>, SessionError> {
self.0
.get(key)
.await
.map_err(|e| SessionError(e.to_string()))
}
pub async fn forget<T: DeserializeOwned>(&self, key: &str) -> Result<Option<T>, SessionError> {
self.0
.remove(key)
.await
.map_err(|e| SessionError(e.to_string()))
}
pub async fn regenerate(&self) -> Result<(), SessionError> {
self.0
.cycle_id()
.await
.map_err(|e| SessionError(e.to_string()))
}
pub async fn flush(&self) -> Result<(), SessionError> {
self.0
.flush()
.await
.map_err(|e| SessionError(e.to_string()))
}
#[must_use]
pub fn raw(&self) -> &TowerSession {
&self.0
}
}
#[derive(Debug)]
pub struct SessionError(pub String);
impl std::fmt::Display for SessionError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "session error: {}", self.0)
}
}
impl std::error::Error for SessionError {}
impl<S> FromRequestParts<S> for Session
where
S: Send + Sync,
{
type Rejection = Infallible;
async fn from_request_parts(
parts: &mut axum::http::request::Parts,
state: &S,
) -> Result<Self, Self::Rejection> {
let session = TowerSession::from_request_parts(parts, state)
.await
.map_err(|_| unreachable!("Session extraction is infallible"))?;
Ok(Session(session))
}
}
pub struct Flash {
session: TowerSession,
messages: Vec<FlashMessage>,
data: std::collections::BTreeMap<String, String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FlashMessage {
pub level: FlashLevel,
pub message: String,
}
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum FlashLevel {
Success,
Error,
Warning,
Info,
}
const FLASH_KEY: &str = "_flash";
pub(crate) const FLASH_DATA_KEY: &str = "_flash_data";
impl Flash {
#[must_use]
pub fn messages(&self) -> &[FlashMessage] {
&self.messages
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.messages.is_empty() && self.data.is_empty()
}
#[must_use]
pub fn get(&self, key: &str) -> Option<&str> {
self.data.get(key).map(String::as_str)
}
#[must_use]
pub fn data(&self) -> &std::collections::BTreeMap<String, String> {
&self.data
}
pub async fn success(&self, message: &str) -> Result<(), FlashError> {
self.add(FlashLevel::Success, message).await
}
pub async fn error(&self, message: &str) -> Result<(), FlashError> {
self.add(FlashLevel::Error, message).await
}
pub async fn warning(&self, message: &str) -> Result<(), FlashError> {
self.add(FlashLevel::Warning, message).await
}
pub async fn info(&self, message: &str) -> Result<(), FlashError> {
self.add(FlashLevel::Info, message).await
}
async fn add(&self, level: FlashLevel, message: &str) -> Result<(), FlashError> {
let mut messages: Vec<FlashMessage> = self
.session
.get(FLASH_KEY)
.await
.map_err(|e| FlashError::Session(e.to_string()))?
.unwrap_or_default();
messages.push(FlashMessage {
level,
message: message.to_string(),
});
self.session
.insert(FLASH_KEY, &messages)
.await
.map_err(|e| FlashError::Session(e.to_string()))?;
Ok(())
}
}
impl<S> FromRequestParts<S> for Flash
where
S: Send + Sync,
{
type Rejection = Infallible;
async fn from_request_parts(
parts: &mut axum::http::request::Parts,
state: &S,
) -> Result<Self, Self::Rejection> {
let session = TowerSession::from_request_parts(parts, state)
.await
.map_err(|_| unreachable!("Session extraction is infallible"))?;
let messages: Vec<FlashMessage> = session
.get(FLASH_KEY)
.await
.map_err(|e| FlashError::Session(e.to_string()))
.unwrap_or(None)
.unwrap_or_default();
let _ = session.remove::<Vec<FlashMessage>>(FLASH_KEY).await;
let data: std::collections::BTreeMap<String, String> = session
.get(FLASH_DATA_KEY)
.await
.unwrap_or(None)
.unwrap_or_default();
if !data.is_empty() {
let _ = session
.remove::<std::collections::BTreeMap<String, String>>(FLASH_DATA_KEY)
.await;
}
Ok(Flash {
session,
messages,
data,
})
}
}
#[derive(Debug)]
pub enum FlashError {
Session(String),
}
impl std::fmt::Display for FlashError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Session(msg) => write!(f, "session error: {msg}"),
}
}
}
impl std::error::Error for FlashError {}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AuthzError {
Forbidden,
}
impl std::fmt::Display for AuthzError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Forbidden => write!(f, "forbidden: policy denied the action"),
}
}
}
impl std::error::Error for AuthzError {}
impl IntoResponse for AuthzError {
fn into_response(self) -> Response {
(axum::http::StatusCode::FORBIDDEN, "Forbidden").into_response()
}
}
pub trait Policy<M>: Send + Sync + 'static {
type User: AuthUser;
fn check(user: &Self::User, action: &str, resource: &M) -> bool;
}