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

use axum::{
    Json,
    extract::{Extension, Query, State},
    http::{HeaderMap, StatusCode},
    response::{IntoResponse, Redirect, Response},
};
use module_auth::{
    AuthLoginMethod, BindWechatPhoneInput, CreateWechatAuthStateInput, WechatAuthError,
};
use platform_auth::WechatAuthScene;
use shared_contracts::auth::{
    WechatBindPhoneRequest, WechatBindPhoneResponse, WechatCallbackQuery, WechatStartQuery,
    WechatStartResponse,
};
use time::OffsetDateTime;
use url::Url;

use crate::{
    api_response::json_success_body,
    auth::AuthenticatedAccessToken,
    auth_payload::map_auth_user_payload,
    auth_session::{
        attach_set_cookie_header, build_refresh_session_cookie_header, create_auth_session,
    },
    http_error::AppError,
    platform_errors::{attach_retry_after, map_wechat_provider_error},
    request_context::RequestContext,
    session_client::resolve_session_client_context,
    state::AppState,
};

pub async fn start_wechat_login(
    State(state): State<AppState>,
    Extension(request_context): Extension<RequestContext>,
    headers: HeaderMap,
    Query(query): Query<WechatStartQuery>,
) -> Result<Json<serde_json::Value>, AppError> {
    if !state.config.wechat_auth_enabled {
        return Err(AppError::from_status(StatusCode::BAD_REQUEST).with_message("微信登录暂未启用"));
    }
    let user_agent = headers
        .get("user-agent")
        .and_then(|value| value.to_str().ok())
        .map(|value| value.to_string());
    let scene = resolve_wechat_scene(user_agent.as_deref())?;
    let state_record = state
        .wechat_auth_state_service()
        .create_state(
            CreateWechatAuthStateInput {
                redirect_path: normalize_redirect_path(
                    query.redirect_path.as_deref(),
                    &state.config.wechat_redirect_path,
                ),
                scene: map_wechat_scene_to_domain(&scene),
                request_user_agent: user_agent.clone(),
            },
            OffsetDateTime::now_utc(),
        )
        .map_err(map_wechat_auth_error)?;
    let authorization_url = state
        .wechat_provider()
        .build_authorization_url(
            &resolve_wechat_callback_url(&state, &headers)?,
            &state_record.state.state_token,
            &scene,
        )
        .map_err(map_wechat_provider_error)?;

    Ok(json_success_body(
        Some(&request_context),
        WechatStartResponse { authorization_url },
    ))
}

pub async fn handle_wechat_callback(
    State(state): State<AppState>,
    headers: HeaderMap,
    Query(query): Query<WechatCallbackQuery>,
) -> Result<impl IntoResponse, AppError> {
    if !state.config.wechat_auth_enabled {
        return Err(AppError::from_status(StatusCode::BAD_REQUEST).with_message("微信登录暂未启用"));
    }
    let fallback_redirect = state.config.wechat_redirect_path.clone();
    let state_token = query
        .state
        .as_deref()
        .unwrap_or_default()
        .trim()
        .to_string();
    if state_token.is_empty() {
        return Ok(Redirect::to(&build_auth_result_redirect_url(
            &fallback_redirect,
            &[
                ("auth_provider", "wechat"),
                ("auth_error", "微信登录状态已失效，请重新发起登录。"),
            ],
        ))
        .into_response());
    }

    let consumed = match state
        .wechat_auth_state_service()
        .consume_state(&state_token, OffsetDateTime::now_utc())
    {
        Ok(value) => value,
        Err(_) => {
            return Ok(Redirect::to(&build_auth_result_redirect_url(
                &fallback_redirect,
                &[
                    ("auth_provider", "wechat"),
                    ("auth_error", "微信登录状态已失效，请重新发起登录。"),
                ],
            ))
            .into_response());
        }
    };

    let redirect_path = consumed.state.redirect_path.clone();
    let session_client = resolve_session_client_context(&headers);

    let result = match state
        .wechat_provider()
        .resolve_callback_profile(query.code.as_deref(), query.mock_code.as_deref())
        .await
    {
        Ok(profile) => state
            .wechat_auth_service()
            .resolve_login(module_auth::ResolveWechatLoginInput {
                profile: map_wechat_profile_to_domain(profile),
            })
            .await
            .map_err(map_wechat_auth_error),
        Err(error) => Err(map_wechat_provider_error(error)),
    };

    match result {
        Ok(result) => {
            let signed_session = create_auth_session(
                &state,
                &result.user,
                &session_client,
                AuthLoginMethod::Wechat,
            )?;
            state
                .sync_auth_store_snapshot_to_spacetime()
                .await
                .map_err(|error| {
                    AppError::from_status(StatusCode::INTERNAL_SERVER_ERROR)
                        .with_message(format!("同步认证快照失败：{error}"))
                })?;
            let mut response = Redirect::to(&build_auth_result_redirect_url(
                &redirect_path,
                &[
                    ("auth_provider", "wechat"),
                    ("auth_token", signed_session.access_token.as_str()),
                    ("auth_binding_status", result.user.binding_status.as_str()),
                ],
            ))
            .into_response();
            attach_set_cookie_header(
                response.headers_mut(),
                build_refresh_session_cookie_header(&state, &signed_session.refresh_token)?,
            );
            Ok(response)
        }
        Err(error) => Ok(Redirect::to(&build_auth_result_redirect_url(
            &redirect_path,
            &[("auth_provider", "wechat"), ("auth_error", error.message())],
        ))
        .into_response()),
    }
}

pub async fn bind_wechat_phone(
    State(state): State<AppState>,
    Extension(request_context): Extension<RequestContext>,
    Extension(authenticated): Extension<AuthenticatedAccessToken>,
    headers: HeaderMap,
    Json(payload): Json<WechatBindPhoneRequest>,
) -> Result<impl IntoResponse, AppError> {
    if !state.config.wechat_auth_enabled {
        return Err(AppError::from_status(StatusCode::BAD_REQUEST).with_message("微信登录暂未启用"));
    }
    let result = state
        .phone_auth_service()
        .bind_wechat_phone(
            BindWechatPhoneInput {
                user_id: authenticated.claims().user_id().to_string(),
                phone_number: payload.phone,
                verify_code: payload.code,
            },
            OffsetDateTime::now_utc(),
        )
        .await
        .map_err(map_wechat_bind_phone_error)?;
    if result.activated_new_user {
        crate::registration_reward::grant_new_user_registration_wallet_reward(
            &state,
            &request_context,
            &result.user.id,
        )
        .await;
    }
    let session_client = resolve_session_client_context(&headers);
    let signed_session = create_auth_session(
        &state,
        &result.user,
        &session_client,
        AuthLoginMethod::Wechat,
    )?;
    state
        .sync_auth_store_snapshot_to_spacetime()
        .await
        .map_err(|error| {
            AppError::from_status(StatusCode::INTERNAL_SERVER_ERROR)
                .with_message(format!("同步认证快照失败：{error}"))
        })?;

    let mut response_headers = HeaderMap::new();
    attach_set_cookie_header(
        &mut response_headers,
        build_refresh_session_cookie_header(&state, &signed_session.refresh_token)?,
    );

    Ok((
        response_headers,
        json_success_body(
            Some(&request_context),
            WechatBindPhoneResponse {
                token: signed_session.access_token,
                user: map_auth_user_payload(result.user),
            },
        ),
    ))
}

fn resolve_wechat_scene(user_agent: Option<&str>) -> Result<WechatAuthScene, AppError> {
    let user_agent = user_agent.unwrap_or_default();
    let is_wechat = user_agent.contains("MicroMessenger");
    let is_mobile = user_agent.contains("Android")
        || user_agent.contains("iPhone")
        || user_agent.contains("iPad")
        || user_agent.contains("Mobile");

    if is_wechat {
        return Ok(WechatAuthScene::WechatInApp);
    }
    if is_mobile {
        return Err(AppError::from_status(StatusCode::BAD_REQUEST)
            .with_message("当前浏览器请使用手机号登录，或在微信内打开后再使用微信登录"));
    }

    Ok(WechatAuthScene::Desktop)
}

fn map_wechat_scene_to_domain(scene: &WechatAuthScene) -> module_auth::WechatAuthScene {
    match scene {
        WechatAuthScene::Desktop => module_auth::WechatAuthScene::Desktop,
        WechatAuthScene::WechatInApp => module_auth::WechatAuthScene::WechatInApp,
    }
}

fn map_wechat_profile_to_domain(
    profile: platform_auth::WechatIdentityProfile,
) -> module_auth::WechatIdentityProfile {
    module_auth::WechatIdentityProfile {
        provider_uid: profile.provider_uid,
        provider_union_id: profile.provider_union_id,
        display_name: profile.display_name,
        avatar_url: profile.avatar_url,
    }
}

fn normalize_redirect_path(raw_value: Option<&str>, fallback: &str) -> String {
    let Some(raw_value) = raw_value.map(str::trim).filter(|value| !value.is_empty()) else {
        return fallback.to_string();
    };
    if raw_value.starts_with('/') {
        return raw_value.to_string();
    }
    Url::parse(raw_value)
        .map(|url| {
            format!(
                "{}{}{}",
                url.path(),
                url.query().map(|v| format!("?{v}")).unwrap_or_default(),
                url.fragment().map(|v| format!("#{v}")).unwrap_or_default()
            )
        })
        .unwrap_or_else(|_| fallback.to_string())
}

fn resolve_wechat_callback_url(state: &AppState, headers: &HeaderMap) -> Result<String, AppError> {
    let proto = headers
        .get("x-forwarded-proto")
        .and_then(|value| value.to_str().ok())
        .and_then(|value| value.split(',').next())
        .map(str::trim)
        .filter(|value| !value.is_empty())
        .unwrap_or("http");
    let host = headers
        .get("x-forwarded-host")
        .or_else(|| headers.get("host"))
        .and_then(|value| value.to_str().ok())
        .and_then(|value| value.split(',').next())
        .map(str::trim)
        .filter(|value| !value.is_empty())
        .unwrap_or("127.0.0.1:3000");
    Ok(format!(
        "{proto}://{host}{}",
        state.config.wechat_callback_path
    ))
}

fn build_auth_result_redirect_url(redirect_path: &str, params: &[(&str, &str)]) -> String {
    let hash = params
        .iter()
        .map(|(key, value)| {
            format!(
                "{}={}",
                urlencoding::encode(key),
                urlencoding::encode(value)
            )
        })
        .collect::<Vec<_>>()
        .join("&");
    let path_without_hash = redirect_path.split('#').next().unwrap_or("/");
    format!(
        "{}#{}",
        if path_without_hash.is_empty() {
            "/"
        } else {
            path_without_hash
        },
        hash
    )
}

#[allow(dead_code)]
fn _assert_response_type(_: Response) {}

fn map_wechat_auth_error(error: WechatAuthError) -> AppError {
    match error {
        WechatAuthError::MissingProfile
        | WechatAuthError::StateNotFound
        | WechatAuthError::StateExpired
        | WechatAuthError::StateConsumed
        | WechatAuthError::MissingWechatIdentity => {
            AppError::from_status(StatusCode::BAD_REQUEST).with_message(error.to_string())
        }
        WechatAuthError::UserNotFound => {
            AppError::from_status(StatusCode::UNAUTHORIZED).with_message(error.to_string())
        }
        WechatAuthError::Store(_) | WechatAuthError::PasswordHash(_) => {
            AppError::from_status(StatusCode::INTERNAL_SERVER_ERROR).with_message(error.to_string())
        }
    }
}

fn map_wechat_bind_phone_error(error: module_auth::PhoneAuthError) -> AppError {
    match error {
        module_auth::PhoneAuthError::InvalidPhoneNumber
        | module_auth::PhoneAuthError::InvalidVerifyCode
        | module_auth::PhoneAuthError::VerifyCodeNotFound
        | module_auth::PhoneAuthError::VerifyCodeExpired
        | module_auth::PhoneAuthError::UserStateMismatch => {
            AppError::from_status(StatusCode::BAD_REQUEST).with_message(error.to_string())
        }
        module_auth::PhoneAuthError::SendCoolingDown {
            retry_after_seconds,
        } => {
            let app_error = AppError::from_status(StatusCode::TOO_MANY_REQUESTS)
                .with_message(error.to_string())
                .with_details(serde_json::json!({ "retryAfterSeconds": retry_after_seconds }));
            attach_retry_after(app_error, retry_after_seconds)
        }
        module_auth::PhoneAuthError::VerifyAttemptsExceeded => {
            AppError::from_status(StatusCode::TOO_MANY_REQUESTS).with_message(error.to_string())
        }
        module_auth::PhoneAuthError::UserNotFound => {
            AppError::from_status(StatusCode::UNAUTHORIZED).with_message(error.to_string())
        }
        module_auth::PhoneAuthError::Store(_) | module_auth::PhoneAuthError::PasswordHash(_) => {
            AppError::from_status(StatusCode::INTERNAL_SERVER_ERROR).with_message(error.to_string())
        }
    }
}