Genarrative/server-rs/crates/api-server/src/refresh_session.rs

use axum::{
    extract::{Extension, State},
    http::HeaderMap,
    response::IntoResponse,
};
use module_auth::{RefreshSessionError, RotateRefreshSessionInput};
use platform_auth::hash_refresh_session_token;
use shared_contracts::auth::RefreshSessionResponse;
use time::OffsetDateTime;

use crate::{
    api_response::json_success_body,
    auth::RefreshSessionToken,
    auth_session::{
        attach_set_cookie_header, build_clear_refresh_session_cookie_header,
        build_refresh_session_cookie_header, map_refresh_session_error,
        record_daily_login_tracking_event_after_auth_success, sign_access_token_for_user,
    },
    http_error::AppError,
    request_context::RequestContext,
    state::AppState,
};

pub async fn refresh_session(
    State(state): State<AppState>,
    Extension(request_context): Extension<RequestContext>,
    maybe_refresh_token: Option<Extension<RefreshSessionToken>>,
) -> Result<impl IntoResponse, AppError> {
    let raw_refresh_token = maybe_refresh_token
        .map(|token| token.0.token().to_string())
        .unwrap_or_default();
    if raw_refresh_token.trim().is_empty() {
        return Err(map_refresh_error_with_clear_cookie(
            &state,
            RefreshSessionError::MissingToken,
        ));
    }
    let refresh_token_hash = hash_refresh_session_token(&raw_refresh_token);
    let next_refresh_token = platform_auth::create_refresh_session_token();
    let next_refresh_token_hash = hash_refresh_session_token(&next_refresh_token);

    // refresh_session 是跨节点的正式认证入口；先加载最新投影，再在本机工作集执行领域轮换，
    // 避免请求落到旧节点时把合法 refresh cookie 误判为不存在。
    state
        .refresh_auth_store_from_spacetime()
        .await
        .map_err(|error| {
            AppError::from_status(axum::http::StatusCode::INTERNAL_SERVER_ERROR)
                .with_message(format!("刷新认证状态失败：{error}"))
        })?;

    let rotated = match state.refresh_session_service().rotate_session(
        RotateRefreshSessionInput {
            refresh_token_hash: refresh_token_hash.clone(),
            next_refresh_token_hash: next_refresh_token_hash.clone(),
        },
        OffsetDateTime::now_utc(),
    ) {
        Ok(rotated) => rotated,
        Err(RefreshSessionError::SessionNotFound) => {
            return Err(map_refresh_error_with_clear_cookie(
                &state,
                RefreshSessionError::SessionNotFound,
            ));
        }
        Err(error) => return Err(map_refresh_error_with_clear_cookie(&state, error)),
    };
    let access_token = sign_access_token_for_user(
        &state,
        &rotated.user,
        &rotated.session.session_id,
        Some(&rotated.session.issued_by_provider),
        Some(&rotated.session.client_info),
    )?;
    state
        .sync_auth_store_tables_to_spacetime()
        .await
        .map_err(|error| {
            AppError::from_status(axum::http::StatusCode::INTERNAL_SERVER_ERROR)
                .with_message(format!("同步认证状态失败：{error}"))
        })?;
    record_daily_login_tracking_event_after_auth_success(
        &state,
        &request_context,
        &rotated.user.id,
        rotated.session.issued_by_provider.clone(),
    )
    .await;

    let mut headers = HeaderMap::new();
    attach_set_cookie_header(
        &mut headers,
        build_refresh_session_cookie_header(&state, &next_refresh_token)?,
    );

    Ok((
        headers,
        json_success_body(
            Some(&request_context),
            RefreshSessionResponse {
                token: access_token,
            },
        ),
    ))
}

fn map_refresh_error_with_clear_cookie(state: &AppState, error: RefreshSessionError) -> AppError {
    let response_error = map_refresh_session_error(error);
    if let Ok(set_cookie) = build_clear_refresh_session_cookie_header(state) {
        return response_error.with_header("set-cookie", set_cookie);
    }

    response_error
}