discord/src/auth.rs

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<anyhow::Error> for ApiError {
    fn from(err: anyhow::Error) -> Self {
        Self::internal(&err.to_string())
    }
}

impl<S> FromRequestParts<S> for AuthUser
where
    AppState: axum::extract::FromRef<S>,
    S: Send + Sync,
{
    type Rejection = ApiError;

    async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
        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<String> {
    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<String> {
    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<Uuid> {
    let mut validation = Validation::default();
    validation.validate_exp = true;

    let data = decode::<SessionClaims>(
        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<String>, 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<String> {
    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()
}