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 OAUTH_STATE_COOKIE: &str = "chattz_oauth_state";
const SESSION_TTL_SECS: u64 = 60 * 15; // 15 minutes
const REFRESH_TTL_SECS: u64 = 60 * 60 * 24 * 30; // 30 days

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct SessionClaims {
    pub sub: String,
    pub kind: String, // "access" or "refresh"
    pub exp: usize,
    pub 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_bearer_token(parts)
            .or_else(|| read_query_token(parts))
            .ok_or_else(|| ApiError::unauthorized("missing jwt token"))?;

        let user_id = verify_session(&token, &app.settings.session_secret, "access")
            .map_err(|_| ApiError::unauthorized("invalid or expired token"))?;

        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_jwt_tokens(user_id: Uuid, secret: &str) -> Result<(String, String)> {
    let now = now_ts();

    let access_claims = SessionClaims {
        sub: user_id.to_string(),
        kind: "access".to_string(),
        iat: now as usize,
        exp: (now + SESSION_TTL_SECS) as usize,
    };

    let refresh_claims = SessionClaims {
        sub: user_id.to_string(),
        kind: "refresh".to_string(),
        iat: now as usize,
        exp: (now + REFRESH_TTL_SECS) as usize,
    };

    let access_token = encode(
        &Header::default(),
        &access_claims,
        &EncodingKey::from_secret(secret.as_bytes()),
    )
    .context("failed to encode access token")?;

    let refresh_token = encode(
        &Header::default(),
        &refresh_claims,
        &EncodingKey::from_secret(secret.as_bytes()),
    )
    .context("failed to encode refresh token")?;

    Ok((access_token, refresh_token))
}

pub fn verify_session(token: &str, secret: &str, expected_kind: &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")?;

    if data.claims.kind != expected_kind {
        return Err(anyhow!("invalid token kind"));
    }

    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_bearer_token(parts: &Parts) -> Option<String> {
    let raw = parts
        .headers
        .get(axum::http::header::AUTHORIZATION)?
        .to_str()
        .ok()?;
    if raw.starts_with("Bearer ") {
        Some(raw["Bearer ".len()..].trim().to_string())
    } else {
        None
    }
}

fn read_query_token(parts: &Parts) -> Option<String> {
    let query = parts.uri.query()?;

    // Simple query param parsing without pulling in url::Url overhead
    for pair in query.split('&') {
        let mut kv = pair.splitn(2, '=');
        let key = kv.next()?;
        let value = kv.next()?;
        if key == "token" {
            return Some(value.to_string());
        }
    }
    None
}

fn now_ts() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_else(|_| Duration::from_secs(0))
        .as_secs()
}