use std::sync::Arc;
use std::time::Duration as StdDuration;
use axum::{
body::Body,
extract::{Query, State},
http::{header, StatusCode},
response::{IntoResponse, Redirect, Response},
};
use axum_extra::extract::cookie::{Cookie, SameSite};
use pep::oidc_client::OidcClient;
use pep::oidc_resource_server::ResourceServerClient;
use pep::oidc::pkce_cookie::PkceCookieManager;
use pep::session_manager::WebSessionManager;
use pep::{DevConfig, JwtClaims, JwtValidationOptions, OidcClientConfig};
use serde::Deserialize;
use time::Duration as TimeDuration;
#[derive(Debug, Clone)]
pub struct AuthConfig {
pub issuer_url: String,
pub client_id: String,
pub client_secret: Option<String>,
pub redirect_uri: String,
pub scope: String,
pub cookie_name: String,
pub dev_config: DevConfig,
pub validation_options: JwtValidationOptions,
pub pkce_cookie_secret: String,
}
impl AuthConfig {
pub fn from_toml(config_toml: &str) -> Option<Self> {
let table: toml::Table = toml::from_str(config_toml).ok()?;
let dev_config = table.get("dev").and_then(|d| d.as_table()).map(|d| {
DevConfig {
local_dev_mode: d.get("local_dev_mode").and_then(|v| v.as_bool()).unwrap_or(false),
local_dev_email: d.get("local_dev_email").and_then(|v| v.as_str()).map(String::from),
local_dev_name: d.get("local_dev_name").and_then(|v| v.as_str()).map(String::from),
local_dev_username: d.get("local_dev_username").and_then(|v| v.as_str()).map(String::from),
}
});
if let Some(ref dc) = dev_config {
if dc.local_dev_mode {
let oidc = Self::parse_oidc_section(&table);
return Some(Self {
issuer_url: oidc.as_ref().map(|o| o.0.clone()).unwrap_or_else(|| "https://auth.example.com".into()),
client_id: oidc.as_ref().map(|o| o.1.clone()).unwrap_or_else(|| "trustee".into()),
client_secret: oidc.as_ref().and_then(|o| o.2.clone()),
redirect_uri: oidc.as_ref().map(|o| o.3.clone()).unwrap_or_else(|| "http://localhost:3000/auth/callback".into()),
scope: oidc.as_ref().map(|o| o.4.clone()).unwrap_or_else(|| "openid profile email".into()),
cookie_name: "trustee_token".into(),
dev_config: dc.clone(),
validation_options: JwtValidationOptions::default(),
pkce_cookie_secret: oidc.as_ref().map(|o| o.6.clone()).unwrap_or_else(|| "trustee-default-pkce-secret-change-me".into()),
});
}
}
let (issuer_url, client_id, client_secret, redirect_uri, scope, validation_options, pkce_secret) =
Self::parse_oidc_section(&table)?;
Some(Self {
issuer_url,
client_id,
client_secret,
redirect_uri,
scope,
cookie_name: "trustee_token".into(),
dev_config: dev_config.unwrap_or_default(),
validation_options,
pkce_cookie_secret: pkce_secret,
})
}
fn parse_oidc_section(
table: &toml::Table,
) -> Option<(String, String, Option<String>, String, String, JwtValidationOptions, String)> {
let oidc = table.get("oidc")?.as_table()?;
let issuer_url = oidc.get("issuer_url")?.as_str()?.to_string();
let client_id = oidc.get("client_id")?.as_str()?.to_string();
let client_secret = oidc.get("client_secret").and_then(|v| v.as_str()).map(String::from);
let redirect_uri = oidc
.get("redirect_url")
.and_then(|v| v.as_str())
.unwrap_or("http://localhost:3000/auth/callback")
.to_string();
let scope = oidc
.get("scope")
.and_then(|v| v.as_str())
.unwrap_or("openid profile email")
.to_string();
let mut validation_options = JwtValidationOptions::default();
if let Some(skip) = oidc.get("skip_issuer_validation").and_then(|v| v.as_bool()) {
validation_options.skip_issuer_validation = skip;
}
if let Some(skip) = oidc.get("skip_audience_validation").and_then(|v| v.as_bool()) {
validation_options.skip_audience_validation = skip;
}
validation_options.expected_audience = oidc
.get("expected_audience")
.and_then(|v| v.as_str())
.map(String::from);
let pkce_secret = oidc
.get("pkce_cookie_secret")
.and_then(|v| v.as_str())
.unwrap_or("trustee-default-pkce-secret-change-me")
.to_string();
Some((issuer_url, client_id, client_secret, redirect_uri, scope, validation_options, pkce_secret))
}
pub fn oidc_client_config(&self) -> OidcClientConfig {
OidcClientConfig {
issuer_url: self.issuer_url.clone(),
client_id: self.client_id.clone(),
client_secret: self.client_secret.clone(),
redirect_uri: self.redirect_uri.clone(),
scope: self.scope.clone(),
code_challenge_method: "S256".to_string(),
}
}
}
#[derive(Clone)]
pub struct AuthState {
pub oidc_client: OidcClient,
pub resource_server: ResourceServerClient,
pub client_config: OidcClientConfig,
pub config: AuthConfig,
pub pkce_manager: PkceCookieManager,
pub session_manager: Arc<WebSessionManager>,
}
impl AuthState {
pub fn new(config: AuthConfig) -> Self {
let pkce_manager = PkceCookieManager::new(
config.pkce_cookie_secret.as_bytes(),
"trustee_pkce_state",
StdDuration::from_secs(600),
);
let session_manager = Arc::new(WebSessionManager::new(
OidcClient::new(),
config.issuer_url.clone(),
config.client_id.clone(),
config.client_secret.clone(),
config.scope.clone(),
));
Self {
oidc_client: OidcClient::new(),
resource_server: ResourceServerClient::new(),
client_config: config.oidc_client_config(),
pkce_manager,
session_manager,
config,
}
}
pub fn is_dev_mode(&self) -> bool {
self.config.dev_config.local_dev_mode
}
pub async fn validate_token(&self, token: &str) -> anyhow::Result<JwtClaims> {
let mut claims = self
.resource_server
.validate_jwt_with_options(
token,
&self.config.issuer_url,
&self.config.client_id,
&self.config.validation_options,
)
.await
.map_err(|e| anyhow::anyhow!("Token validation failed: {}", e))?;
let _ = self
.resource_server
.enrich_claims_with_userinfo(&mut claims, token, &self.config.issuer_url, None)
.await;
if claims.name.is_none() || claims.email.is_none() {
self.fill_userinfo_fields(&mut claims, token).await;
}
Ok(claims)
}
async fn fill_userinfo_fields(&self, claims: &mut JwtClaims, token: &str) {
let userinfo_url = format!("{}/userinfo", self.config.issuer_url.trim_end_matches('/'));
let client = reqwest::Client::new();
let resp = client
.get(&userinfo_url)
.header("Authorization", format!("Bearer {}", token))
.header("Accept", "application/json")
.send()
.await;
let Ok(resp) = resp else {
tracing::debug!("Userinfo request failed for name/email enrichment");
return;
};
if !resp.status().is_success() {
tracing::debug!("Userinfo returned {} for name/email enrichment", resp.status());
return;
}
let Ok(userinfo): Result<serde_json::Map<String, serde_json::Value>, _> = resp.json().await else {
return;
};
tracing::debug!("Userinfo keys: {:?}", userinfo.keys().collect::<Vec<_>>());
if claims.name.is_none() {
if let Some(name) = userinfo.get("name").and_then(|v| v.as_str()) {
claims.name = Some(name.to_string());
}
}
if claims.email.is_none() {
if let Some(email) = userinfo.get("email").and_then(|v| v.as_str()) {
claims.email = Some(email.to_string());
}
}
if claims.preferred_username.is_none() {
if let Some(uname) = userinfo.get("preferred_username").and_then(|v| v.as_str()) {
claims.preferred_username = Some(uname.to_string());
}
}
}
}
#[derive(Debug, Clone)]
pub struct AuthUser {
pub sub: String,
pub email: Option<String>,
pub name: Option<String>,
pub username: Option<String>,
pub is_dev: bool,
}
impl From<JwtClaims> for AuthUser {
fn from(claims: JwtClaims) -> Self {
Self {
sub: claims.sub,
email: claims.email,
name: claims.name,
username: claims.preferred_username,
is_dev: false,
}
}
}
const SESSION_COOKIE_MAX_AGE: StdDuration = StdDuration::from_secs(3600);
pub async fn check_auth(
auth: &Option<Arc<AuthState>>,
headers: &axum::http::HeaderMap,
) -> Result<Option<String>, StatusCode> {
let Some(auth) = auth.as_ref() else {
return Ok(None); };
if let Some(token) = headers
.get(header::AUTHORIZATION)
.and_then(|v| v.to_str().ok())
.and_then(|v| v.strip_prefix("Bearer "))
.map(|s| s.to_string())
{
if token.starts_with("dev:") {
if !auth.config.dev_config.local_dev_mode {
tracing::warn!("Dev token presented but dev mode is disabled — rejecting");
return Err(StatusCode::UNAUTHORIZED);
}
let parts: Vec<&str> = token.splitn(4, ':').collect();
return if parts.len() >= 4 {
Ok(None)
} else {
Err(StatusCode::UNAUTHORIZED)
};
}
return match auth.validate_token(&token).await {
Ok(_) => Ok(None),
Err(e) => {
tracing::warn!("Bearer token validation failed: {}", e);
Err(StatusCode::UNAUTHORIZED)
}
};
}
let session_id = headers
.get(header::COOKIE)
.and_then(|v| v.to_str().ok())
.and_then(|cookies| extract_token_from_cookies(cookies, &auth.config.cookie_name));
let Some(session_id) = session_id else {
tracing::warn!("No auth token found in request");
return Err(StatusCode::UNAUTHORIZED);
};
if session_id.starts_with("dev:") {
if !auth.config.dev_config.local_dev_mode {
tracing::warn!("Dev cookie presented but dev mode is disabled — rejecting");
return Err(StatusCode::UNAUTHORIZED);
}
let parts: Vec<&str> = session_id.splitn(4, ':').collect();
return if parts.len() >= 4 {
Ok(None)
} else {
Err(StatusCode::UNAUTHORIZED)
};
}
match auth.session_manager.get_token(&session_id).await {
Ok(access_token) => match auth.validate_token(&access_token).await {
Ok(_) => {
let secure = auth.client_config.redirect_uri.starts_with("https");
let cookie = create_auth_cookie(
&auth.config.cookie_name,
&session_id,
SESSION_COOKIE_MAX_AGE,
secure,
);
Ok(Some(cookie.to_string()))
}
Err(e) => {
tracing::warn!("Session token validation failed: {} — attempting force-refresh", e);
match auth.session_manager.force_refresh(&session_id).await {
Ok(new_token) => match auth.validate_token(&new_token).await {
Ok(_) => {
let secure = auth.client_config.redirect_uri.starts_with("https");
let cookie = create_auth_cookie(
&auth.config.cookie_name,
&session_id,
SESSION_COOKIE_MAX_AGE,
secure,
);
Ok(Some(cookie.to_string()))
}
Err(e2) => {
tracing::warn!("Session token still invalid after force-refresh: {}", e2);
Err(StatusCode::UNAUTHORIZED)
}
},
Err(e2) => {
tracing::warn!("Force-refresh failed: {}", e2);
Err(StatusCode::UNAUTHORIZED)
}
}
}
},
Err(e) => {
tracing::warn!("Session lookup/refresh failed: {}", e);
Err(StatusCode::UNAUTHORIZED)
}
}
}
async fn resolve_access_token(
auth: &AuthState,
headers: &axum::http::HeaderMap,
) -> Result<String, StatusCode> {
if let Some(token) = headers
.get(header::AUTHORIZATION)
.and_then(|v| v.to_str().ok())
.and_then(|v| v.strip_prefix("Bearer "))
.map(|s| s.to_string())
{
return Ok(token);
}
let session_id = headers
.get(header::COOKIE)
.and_then(|v| v.to_str().ok())
.and_then(|cookies| extract_token_from_cookies(cookies, &auth.config.cookie_name));
match session_id {
Some(sid) if sid.starts_with("dev:") => {
if !auth.config.dev_config.local_dev_mode {
tracing::warn!("Dev cookie in resolve_access_token but dev mode is disabled — rejecting");
Err(StatusCode::UNAUTHORIZED)
} else {
Ok(sid)
}
}
Some(sid) => auth.session_manager.get_token(&sid).await.map_err(|e| {
tracing::warn!("Failed to resolve session token: {}", e);
StatusCode::UNAUTHORIZED
}),
None => Err(StatusCode::UNAUTHORIZED),
}
}
fn extract_token_from_cookies(cookie_header: &str, cookie_name: &str) -> Option<String> {
for cookie in cookie_header.split(';') {
let cookie = cookie.trim();
if let Some(value) = cookie.strip_prefix(&format!("{}=", cookie_name)) {
return Some(value.to_string());
}
}
None
}
pub fn auth_routes() -> axum::Router<crate::ServerState> {
axum::Router::new()
.route("/login", axum::routing::get(login_handler))
.route("/callback", axum::routing::get(callback_handler))
.route("/me", axum::routing::get(me_handler))
.route("/logout", axum::routing::post(logout_handler))
}
#[derive(Debug, Deserialize)]
pub struct CallbackQuery {
pub code: Option<String>,
pub state: Option<String>,
pub error: Option<String>,
pub error_description: Option<String>,
}
async fn login_handler(
State(state): State<crate::ServerState>,
) -> Result<Response, AuthError> {
let auth = state.auth.as_ref().ok_or(AuthError::AuthNotConfigured)?;
if auth.is_dev_mode() {
tracing::info!("Dev mode: creating dev session");
let dev = &auth.config.dev_config;
let dev_token = format!(
"dev:{}:{}:{}",
dev.local_dev_email.as_deref().unwrap_or("dev@localhost"),
dev.local_dev_name.as_deref().unwrap_or("Dev User"),
dev.local_dev_username.as_deref().unwrap_or("dev")
);
let cookie = create_auth_cookie(&auth.config.cookie_name, &dev_token, StdDuration::from_secs(86400), false);
return Ok(Response::builder()
.status(StatusCode::FOUND)
.header(header::LOCATION, "/")
.header(header::SET_COOKIE, cookie.to_string())
.body(Body::empty())
.unwrap());
}
let pkce_session = auth.pkce_manager.create();
let challenge = OidcClient::generate_code_challenge(&pkce_session.verifier);
let auth_url = auth
.oidc_client
.build_authorization_url(&auth.client_config, &pkce_session.state, Some(&challenge))
.await
.map_err(|e| AuthError::OidcError(e.to_string()))?;
let secure = auth.client_config.redirect_uri.starts_with("https");
let pkce_cookie = Cookie::build((
auth.pkce_manager.cookie_name().to_string(),
pkce_session.cookie_value,
))
.path("/")
.http_only(true)
.same_site(SameSite::Lax)
.secure(secure)
.max_age(TimeDuration::seconds(auth.pkce_manager.ttl().as_secs() as i64))
.build();
Ok(Response::builder()
.status(StatusCode::TEMPORARY_REDIRECT)
.header(header::LOCATION, &auth_url)
.header(header::SET_COOKIE, pkce_cookie.to_string())
.body(Body::empty())
.unwrap())
}
async fn callback_handler(
State(state): State<crate::ServerState>,
Query(query): Query<CallbackQuery>,
headers: axum::http::HeaderMap,
) -> Result<Response, AuthError> {
let auth = state.auth.as_ref().ok_or(AuthError::AuthNotConfigured)?;
if let Some(error) = query.error {
let desc = query.error_description.unwrap_or_default();
tracing::error!("OIDC error: {} - {}", error, desc);
return Ok(Redirect::temporary(&format!(
"/?error={}&error_description={}",
urlencoding::encode(&error),
urlencoding::encode(&desc)
))
.into_response());
}
let code = query.code.ok_or(AuthError::MissingCode)?;
let oauth_state = query.state.ok_or(AuthError::MissingState)?;
let cookie_header = headers
.get(header::COOKIE)
.and_then(|v| v.to_str().ok())
.unwrap_or("");
let pkce_value = extract_token_from_cookies(cookie_header, auth.pkce_manager.cookie_name())
.ok_or(AuthError::InvalidState)?;
let verifier = auth
.pkce_manager
.verify(&pkce_value, &oauth_state)
.ok_or(AuthError::InvalidState)?;
tracing::info!("Exchanging authorization code for tokens");
let token_response = auth
.oidc_client
.exchange_code_for_tokens(&auth.client_config, &code, Some(&verifier))
.await
.map_err(|e| AuthError::TokenExchangeFailed(e.to_string()))?;
let session_id = auth
.session_manager
.create_session(&token_response)
.await
.map_err(|e| AuthError::TokenExchangeFailed(format!("Session creation failed: {}", e)))?;
let max_age = SESSION_COOKIE_MAX_AGE;
let secure = auth.client_config.redirect_uri.starts_with("https");
let cookie = create_auth_cookie(&auth.config.cookie_name, &session_id, max_age, secure);
let clear_pkce = Cookie::build((auth.pkce_manager.cookie_name().to_string(), ""))
.path("/")
.http_only(true)
.same_site(SameSite::Lax)
.max_age(TimeDuration::seconds(-1))
.build();
tracing::info!("Authentication successful, redirecting to /");
Ok(Response::builder()
.status(StatusCode::FOUND)
.header(header::LOCATION, "/")
.header(header::SET_COOKIE, cookie.to_string())
.header(header::SET_COOKIE, clear_pkce.to_string())
.body(Body::empty())
.unwrap())
}
async fn me_handler(
State(state): State<crate::ServerState>,
headers: axum::http::HeaderMap,
) -> Response {
let Some(ref auth) = state.auth else {
return axum::Json(serde_json::json!({
"authenticated": true,
"auth_enabled": false
}))
.into_response();
};
let cookie_header = headers
.get(header::COOKIE)
.and_then(|v| v.to_str().ok())
.unwrap_or("");
let bearer = headers
.get(header::AUTHORIZATION)
.and_then(|v| v.to_str().ok())
.and_then(|v| v.strip_prefix("Bearer "))
.map(String::from);
let token = bearer.clone().or_else(|| extract_token_from_cookies(cookie_header, &auth.config.cookie_name));
let Some(cookie_value) = token else {
return axum::Json(serde_json::json!({
"authenticated": false,
"auth_enabled": true
}))
.into_response();
};
if cookie_value.starts_with("dev:") && auth.config.dev_config.local_dev_mode {
let parts: Vec<&str> = cookie_value.splitn(4, ':').collect();
if parts.len() >= 4 {
return axum::Json(serde_json::json!({
"authenticated": true,
"auth_enabled": true,
"email": parts[1],
"name": parts[2],
"username": parts[3],
"dev_mode": true
}))
.into_response();
}
}
let access_token = if bearer.is_some() {
cookie_value
} else {
match auth.session_manager.get_token(&cookie_value).await {
Ok(token) => token,
Err(e) => {
tracing::debug!("Session token resolution failed for /auth/me: {}", e);
return axum::Json(serde_json::json!({
"authenticated": false,
"auth_enabled": true
}))
.into_response();
}
}
};
match auth.validate_token(&access_token).await {
Ok(claims) => axum::Json(serde_json::json!({
"authenticated": true,
"auth_enabled": true,
"sub": claims.sub,
"email": claims.email,
"name": claims.name,
"username": claims.preferred_username,
"dev_mode": false
}))
.into_response(),
Err(e) => {
tracing::debug!("Token validation failed for /auth/me: {}", e);
axum::Json(serde_json::json!({
"authenticated": false,
"auth_enabled": true
}))
.into_response()
}
}
}
async fn logout_handler(
State(state): State<crate::ServerState>,
headers: axum::http::HeaderMap,
) -> Response {
let cookie_name = state
.auth
.as_ref()
.map(|a| a.config.cookie_name.as_str())
.unwrap_or("trustee_token");
if let Some(ref auth) = state.auth {
if let Some(cookie_header) = headers.get(header::COOKIE).and_then(|v| v.to_str().ok()) {
if let Some(session_id) = extract_token_from_cookies(cookie_header, cookie_name) {
if !session_id.starts_with("dev:") {
let _ = auth.session_manager.destroy_session(&session_id);
}
}
}
}
let cookie = Cookie::build((cookie_name.to_string(), ""))
.path("/")
.http_only(true)
.same_site(SameSite::Lax)
.max_age(TimeDuration::seconds(-1))
.build();
Response::builder()
.status(StatusCode::FOUND)
.header(header::LOCATION, "/")
.header(header::SET_COOKIE, cookie.to_string())
.body(Body::empty())
.unwrap()
}
fn create_auth_cookie(name: &str, value: &str, max_age: StdDuration, secure: bool) -> Cookie<'static> {
Cookie::build((name.to_string(), value.to_string()))
.path("/")
.http_only(true)
.same_site(SameSite::Lax)
.secure(secure)
.max_age(TimeDuration::seconds(max_age.as_secs() as i64))
.build()
}
#[derive(Debug)]
pub enum AuthError {
MissingCode,
MissingState,
InvalidState,
OidcError(String),
TokenExchangeFailed(String),
AuthNotConfigured,
}
impl IntoResponse for AuthError {
fn into_response(self) -> Response {
let (_status, msg) = match self {
AuthError::MissingCode => (StatusCode::BAD_REQUEST, "Missing authorization code"),
AuthError::MissingState => (StatusCode::BAD_REQUEST, "Missing state parameter"),
AuthError::InvalidState => (StatusCode::BAD_REQUEST, "Invalid or expired state"),
AuthError::OidcError(_) => (StatusCode::SERVICE_UNAVAILABLE, "Authentication service error"),
AuthError::TokenExchangeFailed(_) => (StatusCode::BAD_REQUEST, "Token exchange failed"),
AuthError::AuthNotConfigured => (StatusCode::NOT_IMPLEMENTED, "Authentication not configured"),
};
Redirect::temporary(&format!("/?error={}", urlencoding::encode(msg))).into_response()
}
}