papercrate/backend/src/routes/auth.rs

use axum::{
    extract::State,
    http::{header::SET_COOKIE, HeaderMap, HeaderValue, StatusCode},
    response::{IntoResponse, Response},
    Json,
};
use axum_extra::{
    headers::{authorization::Bearer, Authorization, Cookie},
    typed_header::TypedHeader,
};
use chrono::{Duration as ChronoDuration, Utc};
use diesel::{pg::PgConnection, prelude::*};
use rand::{rngs::OsRng, RngCore};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use uuid::Uuid;

use crate::{
    auth::{password, AuthenticatedUser},
    error::{AppError, AppResult},
    models::{NewRefreshToken, RefreshToken, Tenant, User, UserMembership},
    schema::{
        refresh_tokens, tenants::dsl as tenant_dsl, user_memberships::dsl as memberships_dsl,
        users::dsl,
    },
    state::AppState,
};

use crate::schema::refresh_tokens::dsl as refresh_dsl;

const REFRESH_COOKIE_NAME: &str = "refresh_token";

#[derive(Deserialize)]
pub struct LoginRequest {
    pub username: String,
    pub password: String,
    #[serde(default)]
    pub preferred_tenant_slug: Option<String>,
}

#[derive(Serialize)]
pub struct LoginResponse {
    pub access_token: String,
    pub token_type: String,
    pub expires_in: i64,
    pub tenant: TenantSnippet,
}

#[derive(Serialize)]
pub struct TenantSummary {
    pub tenant_id: Uuid,
    pub slug: String,
}

#[derive(Serialize)]
pub struct TenantSnippet {
    pub id: Uuid,
    pub slug: String,
}

#[derive(Serialize)]
pub struct TenantSelectionResponse {
    pub access_token: String,
    pub tenants: Vec<TenantSummary>,
}

#[derive(Serialize)]
pub struct TenantListResponse {
    pub tenants: Vec<TenantSnippet>,
}

#[derive(Deserialize)]
pub struct TenantSelectionRequest {
    pub tenant_id: Uuid,
}

pub async fn login(
    State(state): State<AppState>,
    Json(payload): Json<LoginRequest>,
) -> AppResult<Response> {
    let mut conn = state.db_unscoped()?;

    let user: User = dsl::users
        .filter(dsl::username.eq(&payload.username))
        .first(&mut conn)?;

    let valid = password::verify_password(&payload.password, &user.password_hash)
        .map_err(|_| AppError::unauthorized())?;

    if !valid {
        return Err(AppError::unauthorized());
    }

    let memberships: Vec<(UserMembership, Tenant)> = memberships_dsl::user_memberships
        .inner_join(tenant_dsl::tenants)
        .filter(memberships_dsl::user_id.eq(user.id))
        .load(&mut conn)?;

    if memberships.is_empty() {
        return Err(AppError::unauthorized());
    }

    let preferred_slug = payload
        .preferred_tenant_slug
        .as_ref()
        .map(|slug| slug.trim().to_string())
        .filter(|slug| !slug.is_empty());

    if let Some(tenant) = preferred_slug.as_ref().and_then(|slug| {
        memberships
            .iter()
            .find(|(_, tenant)| tenant.slug.eq_ignore_ascii_case(slug))
    }) {
        return issue_session(&state, &mut conn, &user, tenant.1.id);
    }

    if memberships.len() == 1 {
        let tenant_id = memberships[0].1.id;
        return issue_session(&state, &mut conn, &user, tenant_id);
    }

    let selection_token = state
        .jwt
        .generate_tenant_selector_token(user.id)
        .map_err(AppError::from)?;

    let tenants = memberships
        .into_iter()
        .map(|(_, tenant)| TenantSummary {
            tenant_id: tenant.id,
            slug: tenant.slug,
        })
        .collect();

    let response = Json(TenantSelectionResponse {
        access_token: selection_token,
        tenants,
    })
    .into_response();

    Ok(response)
}

pub async fn refresh(
    State(state): State<AppState>,
    jar: Option<TypedHeader<Cookie>>,
) -> AppResult<Response> {
    let cookies = jar.ok_or_else(AppError::unauthorized)?;
    let refresh_value = cookies
        .get(REFRESH_COOKIE_NAME)
        .ok_or_else(AppError::unauthorized)?;

    let hashed = hash_refresh_token(refresh_value);
    let mut conn = state.db_unscoped()?;
    let now = Utc::now();
    let now_naive = now.naive_utc();

    let token = match refresh_dsl::refresh_tokens
        .filter(refresh_dsl::token_hash.eq(&hashed))
        .filter(refresh_dsl::revoked_at.is_null())
        .filter(refresh_dsl::expires_at.gt(now_naive))
        .first::<RefreshToken>(&mut conn)
    {
        Ok(token) => token,
        Err(diesel::result::Error::NotFound) => return Err(AppError::unauthorized()),
        Err(err) => return Err(AppError::from(err)),
    };

    diesel::update(refresh_dsl::refresh_tokens.filter(refresh_dsl::id.eq(token.id)))
        .set((
            refresh_dsl::revoked_at.eq(now_naive),
            refresh_dsl::updated_at.eq(now_naive),
        ))
        .execute(&mut conn)?;

    let user: User = dsl::users
        .find(token.user_id)
        .first(&mut conn)
        .map_err(AppError::from)?;

    issue_session(&state, &mut conn, &user, token.tenant_id)
}

pub async fn select_tenant(
    State(state): State<AppState>,
    TypedHeader(Authorization(bearer)): TypedHeader<Authorization<Bearer>>,
    Json(payload): Json<TenantSelectionRequest>,
) -> AppResult<Response> {
    let claims = state
        .jwt
        .verify_tenant_selector_token(bearer.token())
        .map_err(|_| AppError::unauthorized())?;

    let mut conn = state.db_unscoped()?;

    let membership_exists = memberships_dsl::user_memberships
        .filter(memberships_dsl::user_id.eq(claims.sub))
        .filter(memberships_dsl::tenant_id.eq(payload.tenant_id))
        .inner_join(tenant_dsl::tenants)
        .select(memberships_dsl::id)
        .first::<Uuid>(&mut conn)
        .optional()?;

    if membership_exists.is_none() {
        return Err(AppError::unauthorized());
    }

    let user: User = dsl::users
        .find(claims.sub)
        .first(&mut conn)
        .map_err(AppError::from)?;

    issue_session(&state, &mut conn, &user, payload.tenant_id)
}

pub async fn logout(
    State(state): State<AppState>,
    user: AuthenticatedUser,
    jar: Option<TypedHeader<Cookie>>,
) -> AppResult<(HeaderMap, StatusCode)> {
    let mut conn = state.db_unscoped()?;
    let now = Utc::now().naive_utc();
    let mut rows_affected = 0;

    if let Some(cookies) = jar {
        if let Some(value) = cookies.get(REFRESH_COOKIE_NAME) {
            let hashed = hash_refresh_token(value);
            rows_affected = diesel::update(
                refresh_dsl::refresh_tokens
                    .filter(refresh_dsl::token_hash.eq(hashed))
                    .filter(refresh_dsl::user_id.eq(user.user_id))
                    .filter(refresh_dsl::revoked_at.is_null()),
            )
            .set((
                refresh_dsl::revoked_at.eq(now),
                refresh_dsl::updated_at.eq(now),
            ))
            .execute(&mut conn)
            .unwrap_or(0);
        }
    }

    if rows_affected == 0 {
        let _ = diesel::update(
            refresh_dsl::refresh_tokens
                .filter(refresh_dsl::user_id.eq(user.user_id))
                .filter(refresh_dsl::revoked_at.is_null()),
        )
        .set((
            refresh_dsl::revoked_at.eq(now),
            refresh_dsl::updated_at.eq(now),
        ))
        .execute(&mut conn);
    }

    let mut headers = HeaderMap::new();
    headers.insert(SET_COOKIE, build_clear_refresh_cookie(&state));
    Ok((headers, StatusCode::NO_CONTENT))
}

pub async fn me(user: AuthenticatedUser) -> Json<AuthenticatedUser> {
    Json(user)
}

pub async fn list_tenants(
    State(state): State<AppState>,
    auth: Option<TypedHeader<Authorization<Bearer>>>,
) -> AppResult<Json<TenantListResponse>> {
    let bearer = auth.ok_or_else(AppError::unauthorized)?;
    let token = bearer.token();

    let user_id = match state.jwt.verify_token(token) {
        Ok(claims) => claims.sub,
        Err(_) => {
            let claims = state
                .jwt
                .verify_tenant_selector_token(token)
                .map_err(|_| AppError::unauthorized())?;
            claims.sub
        }
    };

    let mut conn = state.db_unscoped()?;

    let tenants = memberships_dsl::user_memberships
        .inner_join(tenant_dsl::tenants)
        .filter(memberships_dsl::user_id.eq(user_id))
        .select((tenant_dsl::id, tenant_dsl::slug))
        .load::<(Uuid, String)>(&mut conn)?
        .into_iter()
        .map(|(id, slug)| TenantSnippet { id, slug })
        .collect();

    Ok(Json(TenantListResponse { tenants }))
}

fn issue_session(
    state: &AppState,
    conn: &mut PgConnection,
    user: &User,
    tenant_id: Uuid,
) -> AppResult<Response> {
    let now = Utc::now();
    let access_token = state
        .jwt
        .generate_token(user.id, tenant_id, &user.username)
        .map_err(AppError::from)?;

    let tenant_slug: String = tenant_dsl::tenants
        .find(tenant_id)
        .select(tenant_dsl::slug)
        .first(conn)
        .map_err(AppError::from)?;

    let refresh_value = generate_refresh_token();
    let refresh_hash = hash_refresh_token(&refresh_value);
    let refresh_expires_at = now + ChronoDuration::days(state.config.refresh_token_expiry_days);

    let new_refresh = NewRefreshToken {
        id: Uuid::new_v4(),
        user_id: user.id,
        token_hash: refresh_hash,
        issued_at: now.naive_utc(),
        expires_at: refresh_expires_at.naive_utc(),
        tenant_id,
    };

    diesel::insert_into(refresh_tokens::table)
        .values(&new_refresh)
        .execute(conn)?;

    let mut response = Json(LoginResponse {
        access_token,
        token_type: "Bearer".to_string(),
        expires_in: state.config.jwt_expiry_minutes * 60,
        tenant: TenantSnippet {
            id: tenant_id,
            slug: tenant_slug,
        },
    })
    .into_response();

    response.headers_mut().insert(
        SET_COOKIE,
        build_refresh_cookie(state, &refresh_value, refresh_expires_at),
    );

    Ok(response)
}

fn hash_refresh_token(token: &str) -> String {
    let mut hasher = Sha256::new();
    hasher.update(token.as_bytes());
    hex::encode(hasher.finalize())
}

fn generate_refresh_token() -> String {
    let mut bytes = [0u8; 32];
    OsRng.fill_bytes(&mut bytes);
    hex::encode(bytes)
}

fn build_refresh_cookie(
    state: &AppState,
    token: &str,
    expires_at: chrono::DateTime<Utc>,
) -> HeaderValue {
    let max_age = ChronoDuration::days(state.config.refresh_token_expiry_days).num_seconds();

    let mut parts = vec![format!("{}={}", REFRESH_COOKIE_NAME, token)];
    parts.push("Path=/".into());
    parts.push("HttpOnly".into());
    parts.push("SameSite=Strict".into());
    parts.push(format!("Max-Age={}", max_age));
    parts.push(format!("Expires={}", expires_at.to_rfc2822()));
    if state.config.refresh_cookie_secure {
        parts.push("Secure".into());
    }
    if let Some(domain) = &state.config.refresh_cookie_domain {
        parts.push(format!("Domain={}", domain));
    }

    HeaderValue::from_str(&parts.join("; ")).expect("valid refresh cookie")
}

fn build_clear_refresh_cookie(state: &AppState) -> HeaderValue {
    let mut parts = vec![format!("{}=", REFRESH_COOKIE_NAME)];
    parts.push("Path=/".into());
    parts.push("HttpOnly".into());
    parts.push("SameSite=Strict".into());
    parts.push("Max-Age=0".into());
    parts.push("Expires=Thu, 01 Jan 1970 00:00:00 GMT".into());
    if state.config.refresh_cookie_secure {
        parts.push("Secure".into());
    }
    if let Some(domain) = &state.config.refresh_cookie_domain {
        parts.push(format!("Domain={}", domain));
    }

    HeaderValue::from_str(&parts.join("; ")).expect("valid refresh cookie")
}