use std::time::{Duration, SystemTime, UNIX_EPOCH}; use anyhow::{Context, Result, anyhow}; use axum::{ Json, extract::{FromRef, FromRequestParts}, http::{StatusCode, request::Parts}, response::{IntoResponse, Response}, }; use jsonwebtoken::{DecodingKey, EncodingKey, Header, Validation, decode, encode}; use serde::{Deserialize, Serialize}; use uuid::Uuid; use crate::{AppState, db}; const SESSION_COOKIE: &str = "chattz_session"; const OAUTH_STATE_COOKIE: &str = "chattz_oauth_state"; const SESSION_TTL_SECS: u64 = 60 * 60 * 24 * 7; #[derive(Debug, Serialize, Deserialize, Clone)] struct SessionClaims { sub: String, exp: usize, iat: usize, } #[derive(Debug, Clone)] pub struct AuthUser { pub id: Uuid, } #[derive(Debug)] pub struct ApiError { pub status: StatusCode, pub message: String, } impl ApiError { pub fn unauthorized(msg: &str) -> Self { Self { status: StatusCode::UNAUTHORIZED, message: msg.to_string(), } } pub fn bad_request(msg: &str) -> Self { Self { status: StatusCode::BAD_REQUEST, message: msg.to_string(), } } pub fn internal(msg: &str) -> Self { Self { status: StatusCode::INTERNAL_SERVER_ERROR, message: msg.to_string(), } } } #[derive(Serialize)] struct ErrorBody<'a> { error: &'a str, } impl IntoResponse for ApiError { fn into_response(self) -> Response { (self.status, Json(ErrorBody { error: &self.message })).into_response() } } impl From for ApiError { fn from(err: anyhow::Error) -> Self { Self::internal(&err.to_string()) } } impl FromRequestParts for AuthUser where AppState: axum::extract::FromRef, S: Send + Sync, { type Rejection = ApiError; async fn from_request_parts(parts: &mut Parts, state: &S) -> Result { let app = AppState::from_ref(state); let token = read_cookie(parts, SESSION_COOKIE).ok_or_else(|| ApiError::unauthorized("missing session"))?; let user_id = verify_session(&token, &app.settings.session_secret) .map_err(|_| ApiError::unauthorized("invalid session"))?; let exists = db::user_exists(&app.db, user_id) .await .map_err(|_| ApiError::unauthorized("session user not found"))?; if !exists { return Err(ApiError::unauthorized("session user not found")); } Ok(Self { id: user_id }) } } pub fn new_oauth_state() -> String { Uuid::new_v4().to_string() } pub fn make_oauth_state_cookie(value: &str, secure: bool) -> String { format!( "{name}={value}; Path=/; HttpOnly; SameSite=Lax; Max-Age=600{secure_flag}", name = OAUTH_STATE_COOKIE, secure_flag = if secure { "; Secure" } else { "" } ) } pub fn clear_oauth_state_cookie(secure: bool) -> String { format!( "{name}=; Path=/; HttpOnly; SameSite=Lax; Max-Age=0{secure_flag}", name = OAUTH_STATE_COOKIE, secure_flag = if secure { "; Secure" } else { "" } ) } pub fn read_oauth_state_from_headers(headers: &axum::http::HeaderMap) -> Option { let raw = headers.get(axum::http::header::COOKIE)?.to_str().ok()?; raw.split(';').find_map(|pair| { let mut kv = pair.trim().splitn(2, '='); let key = kv.next()?; let value = kv.next()?; (key == OAUTH_STATE_COOKIE).then(|| value.to_string()) }) } pub fn new_session_cookie(user_id: Uuid, secret: &str, secure: bool) -> Result { let now = now_ts(); let claims = SessionClaims { sub: user_id.to_string(), iat: now as usize, exp: (now + SESSION_TTL_SECS) as usize, }; let token = encode( &Header::default(), &claims, &EncodingKey::from_secret(secret.as_bytes()), ) .context("failed to encode session token")?; Ok(format!( "{name}={token}; Path=/; HttpOnly; SameSite=Lax; Max-Age={ttl}{secure_flag}", name = SESSION_COOKIE, ttl = SESSION_TTL_SECS, secure_flag = if secure { "; Secure" } else { "" } )) } pub fn clear_session_cookie(secure: bool) -> String { format!( "{name}=; Path=/; HttpOnly; SameSite=Lax; Max-Age=0{secure_flag}", name = SESSION_COOKIE, secure_flag = if secure { "; Secure" } else { "" } ) } fn verify_session(token: &str, secret: &str) -> Result { let mut validation = Validation::default(); validation.validate_exp = true; let data = decode::( token, &DecodingKey::from_secret(secret.as_bytes()), &validation, ) .context("failed to decode session token")?; let user_id = Uuid::parse_str(&data.claims.sub).context("invalid sub in session token")?; Ok(user_id) } pub fn validate_oauth_state(expected_cookie: Option, query_state: &str) -> Result<()> { let expected = expected_cookie.ok_or_else(|| anyhow!("missing oauth state cookie"))?; if expected != query_state { return Err(anyhow!("invalid oauth state")); } Ok(()) } fn read_cookie(parts: &Parts, name: &str) -> Option { let raw = parts.headers.get(axum::http::header::COOKIE)?.to_str().ok()?; raw.split(';').find_map(|pair| { let mut kv = pair.trim().splitn(2, '='); let key = kv.next()?; let value = kv.next()?; (key == name).then(|| value.to_string()) }) } fn now_ts() -> u64 { SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap_or_else(|_| Duration::from_secs(0)) .as_secs() }