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 tracing::warn;

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);

    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) => {
            match state.refresh_auth_store_from_spacetime().await {
                Ok(true) => {}
                Ok(false) => {
                    return Err(map_refresh_error_with_clear_cookie(
                        &state,
                        RefreshSessionError::SessionNotFound,
                    ));
                }
                Err(error) => {
                    warn!(
                        request_id = request_context.request_id(),
                        error = %error,
                        "refresh session 本地未命中后刷新认证工作集失败"
                    );
                    return Err(map_refresh_error_with_clear_cookie(
                        &state,
                        RefreshSessionError::SessionNotFound,
                    ));
                }
            }
            state
                .refresh_session_service()
                .rotate_session(
                    RotateRefreshSessionInput {
                        refresh_token_hash,
                        next_refresh_token_hash,
                    },
                    OffsetDateTime::now_utc(),
                )
                .map_err(|error| map_refresh_error_with_clear_cookie(&state, error))?
        }
        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_snapshot_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
}