持久化认证状态到 SpacetimeDB 正式表

server-rs/crates/spacetime-module/src/auth.rs

@@ -1,4 +1,6 @@
 use crate::*;
+use serde::{Deserialize, Serialize};
+use serde_json::Value;
 
 const AUTH_STORE_SNAPSHOT_ID: &str = "default";
 
@@ -21,6 +23,20 @@ pub struct AuthStoreSnapshotProcedureResult {
     pub error_message: Option<String>,
 }
 
+#[derive(Clone, Debug, PartialEq, Eq, SpacetimeType)]
+pub struct AuthStoreSnapshotImportRecord {
+    pub imported_user_count: u32,
+    pub imported_identity_count: u32,
+    pub imported_refresh_session_count: u32,
+}
+
+#[derive(Clone, Debug, PartialEq, Eq, SpacetimeType)]
+pub struct AuthStoreSnapshotImportProcedureResult {
+    pub ok: bool,
+    pub record: Option<AuthStoreSnapshotImportRecord>,
+    pub error_message: Option<String>,
+}
+
 #[spacetimedb::table(accessor = auth_store_snapshot)]
 pub struct AuthStoreSnapshot {
     #[primary_key]
@@ -29,6 +45,63 @@ pub struct AuthStoreSnapshot {
     pub(crate) updated_at: Timestamp,
 }
 
+#[spacetimedb::table(
+    accessor = user_account,
+    index(accessor = by_user_account_username, btree(columns = [username])),
+    index(accessor = by_user_account_public_code, btree(columns = [public_user_code]))
+)]
+pub struct UserAccount {
+    #[primary_key]
+    pub(crate) user_id: String,
+    pub(crate) public_user_code: String,
+    pub(crate) username: String,
+    pub(crate) display_name: String,
+    pub(crate) phone_number_masked: Option<String>,
+    pub(crate) phone_number_e164: Option<String>,
+    pub(crate) login_method: String,
+    pub(crate) binding_status: String,
+    pub(crate) wechat_bound: bool,
+    pub(crate) password_hash: String,
+    pub(crate) password_login_enabled: bool,
+    pub(crate) token_version: u64,
+}
+
+#[spacetimedb::table(
+    accessor = auth_identity,
+    index(accessor = by_auth_identity_user_id, btree(columns = [user_id])),
+    index(accessor = by_auth_identity_provider_uid, btree(columns = [provider, provider_uid]))
+)]
+pub struct AuthIdentity {
+    #[primary_key]
+    pub(crate) identity_id: String,
+    pub(crate) user_id: String,
+    pub(crate) provider: String,
+    pub(crate) provider_uid: String,
+    pub(crate) provider_union_id: Option<String>,
+    pub(crate) phone_e164: Option<String>,
+    pub(crate) display_name: Option<String>,
+    pub(crate) avatar_url: Option<String>,
+}
+
+#[spacetimedb::table(
+    accessor = refresh_session,
+    index(accessor = by_refresh_session_user_id, btree(columns = [user_id])),
+    index(accessor = by_refresh_session_token_hash, btree(columns = [refresh_token_hash]))
+)]
+pub struct RefreshSession {
+    #[primary_key]
+    pub(crate) session_id: String,
+    pub(crate) user_id: String,
+    pub(crate) refresh_token_hash: String,
+    pub(crate) issued_by_provider: String,
+    pub(crate) client_info_json: String,
+    pub(crate) expires_at: String,
+    pub(crate) revoked_at: Option<String>,
+    pub(crate) created_at: String,
+    pub(crate) updated_at: String,
+    pub(crate) last_seen_at: String,
+}
+
 // Axum 启动恢复认证状态时读取当前快照；记录不存在代表尚未产生登录态。
 #[spacetimedb::procedure]
 pub fn get_auth_store_snapshot(ctx: &mut ProcedureContext) -> AuthStoreSnapshotProcedureResult {
@@ -66,6 +139,43 @@ pub fn upsert_auth_store_snapshot(
     }
 }
 
+#[spacetimedb::procedure]
+pub fn import_auth_store_snapshot(
+    ctx: &mut ProcedureContext,
+) -> AuthStoreSnapshotImportProcedureResult {
+    match ctx.try_with_tx(|tx| import_auth_store_snapshot_tx(tx)) {
+        Ok(record) => AuthStoreSnapshotImportProcedureResult {
+            ok: true,
+            record: Some(record),
+            error_message: None,
+        },
+        Err(message) => AuthStoreSnapshotImportProcedureResult {
+            ok: false,
+            record: None,
+            error_message: Some(message),
+        },
+    }
+}
+
+// Axum ??????????????? module-auth ????????????????
+#[spacetimedb::procedure]
+pub fn export_auth_store_snapshot_from_tables(
+    ctx: &mut ProcedureContext,
+) -> AuthStoreSnapshotProcedureResult {
+    match ctx.try_with_tx(|tx| export_auth_store_snapshot_from_tables_tx(tx)) {
+        Ok(record) => AuthStoreSnapshotProcedureResult {
+            ok: true,
+            record: Some(record),
+            error_message: None,
+        },
+        Err(message) => AuthStoreSnapshotProcedureResult {
+            ok: false,
+            record: None,
+            error_message: Some(message),
+        },
+    }
+}
+
 fn get_auth_store_snapshot_tx(ctx: &ReducerContext) -> Result<AuthStoreSnapshotRecord, String> {
     Ok(
         match ctx
@@ -120,3 +230,326 @@ fn upsert_auth_store_snapshot_tx(
         updated_at_micros: Some(input.updated_at_micros),
     })
 }
+
+fn import_auth_store_snapshot_tx(
+    ctx: &ReducerContext,
+) -> Result<AuthStoreSnapshotImportRecord, String> {
+    let snapshot = ctx
+        .db
+        .auth_store_snapshot()
+        .snapshot_id()
+        .find(&AUTH_STORE_SNAPSHOT_ID.to_string())
+        .ok_or_else(|| "认证快照不存在，无法导入正式表".to_string())?;
+    let parsed = serde_json::from_str::<PersistentAuthStoreSnapshot>(&snapshot.snapshot_json)
+        .map_err(|error| format!("认证快照 JSON 解析失败：{error}"))?;
+
+    clear_auth_target_tables(ctx);
+
+    let mut imported_user_count = 0_u32;
+    let mut imported_identity_count = 0_u32;
+    let mut imported_refresh_session_count = 0_u32;
+
+    for stored_user in parsed.users_by_username.into_values() {
+        let user = stored_user.user;
+        ctx.db.user_account().insert(UserAccount {
+            user_id: user.id.clone(),
+            public_user_code: user.public_user_code,
+            username: user.username,
+            display_name: user.display_name,
+            phone_number_masked: user.phone_number_masked,
+            phone_number_e164: stored_user.phone_number.clone(),
+            login_method: user.login_method,
+            binding_status: user.binding_status,
+            wechat_bound: user.wechat_bound,
+            password_hash: stored_user.password_hash,
+            password_login_enabled: stored_user.password_login_enabled,
+            token_version: user.token_version,
+        });
+        imported_user_count += 1;
+
+        if let Some(phone_number) = stored_user.phone_number {
+            ctx.db.auth_identity().insert(AuthIdentity {
+                identity_id: format!("authi_phone_{}", sanitize_identity_component(&phone_number)),
+                user_id: user.id,
+                provider: "phone".to_string(),
+                provider_uid: phone_number.clone(),
+                provider_union_id: None,
+                phone_e164: Some(phone_number),
+                display_name: None,
+                avatar_url: None,
+            });
+            imported_identity_count += 1;
+        }
+    }
+
+    for identity in parsed.wechat_identity_by_provider_uid.into_values() {
+        ctx.db.auth_identity().insert(AuthIdentity {
+            identity_id: format!(
+                "authi_wechat_{}",
+                sanitize_identity_component(&identity.provider_uid)
+            ),
+            user_id: identity.user_id,
+            provider: "wechat".to_string(),
+            provider_uid: identity.provider_uid,
+            provider_union_id: identity.provider_union_id,
+            phone_e164: None,
+            display_name: identity.display_name,
+            avatar_url: identity.avatar_url,
+        });
+        imported_identity_count += 1;
+    }
+
+    for stored_session in parsed.sessions_by_id.into_values() {
+        let session = stored_session.session;
+        let client_info_json = serde_json::to_string(&session.client_info)
+            .map_err(|error| format!("客户端身份序列化失败：{error}"))?;
+        ctx.db.refresh_session().insert(RefreshSession {
+            session_id: session.session_id,
+            user_id: session.user_id,
+            refresh_token_hash: session.refresh_token_hash,
+            issued_by_provider: session.issued_by_provider,
+            client_info_json,
+            expires_at: session.expires_at,
+            revoked_at: session.revoked_at,
+            created_at: session.created_at,
+            updated_at: session.updated_at,
+            last_seen_at: session.last_seen_at,
+        });
+        imported_refresh_session_count += 1;
+    }
+
+    Ok(AuthStoreSnapshotImportRecord {
+        imported_user_count,
+        imported_identity_count,
+        imported_refresh_session_count,
+    })
+}
+
+fn export_auth_store_snapshot_from_tables_tx(
+    ctx: &ReducerContext,
+) -> Result<AuthStoreSnapshotRecord, String> {
+    let users = ctx.db.user_account().iter().collect::<Vec<_>>();
+    let identities = ctx.db.auth_identity().iter().collect::<Vec<_>>();
+    let sessions = ctx.db.refresh_session().iter().collect::<Vec<_>>();
+    if users.is_empty() && identities.is_empty() && sessions.is_empty() {
+        return Ok(AuthStoreSnapshotRecord {
+            snapshot_json: None,
+            updated_at_micros: None,
+        });
+    }
+
+    let mut phone_identity_by_user_id = std::collections::HashMap::new();
+    let mut phone_to_user_id = std::collections::HashMap::new();
+    let mut wechat_identity_by_provider_uid = std::collections::HashMap::new();
+    let mut user_id_by_provider_union_id = std::collections::HashMap::new();
+
+    for identity in identities {
+        match identity.provider.as_str() {
+            "phone" => {
+                let phone_number = identity
+                    .phone_e164
+                    .clone()
+                    .unwrap_or_else(|| identity.provider_uid.clone());
+                phone_to_user_id.insert(phone_number.clone(), identity.user_id.clone());
+                phone_identity_by_user_id.insert(identity.user_id, phone_number);
+            }
+            "wechat" => {
+                if let Some(union_id) = identity.provider_union_id.clone() {
+                    user_id_by_provider_union_id.insert(union_id, identity.user_id.clone());
+                }
+                wechat_identity_by_provider_uid.insert(
+                    identity.provider_uid.clone(),
+                    StoredWechatIdentitySnapshot {
+                        user_id: identity.user_id,
+                        provider_uid: identity.provider_uid,
+                        provider_union_id: identity.provider_union_id,
+                        display_name: identity.display_name,
+                        avatar_url: identity.avatar_url,
+                    },
+                );
+            }
+            _ => {}
+        }
+    }
+
+    let mut next_user_id = 1_u64;
+    let mut users_by_username = std::collections::HashMap::new();
+    for user in users {
+        if let Some(numeric_id) = user
+            .user_id
+            .strip_prefix("user_")
+            .and_then(|value| value.parse::<u64>().ok())
+        {
+            next_user_id = next_user_id.max(numeric_id.saturating_add(1));
+        }
+        let auth_user = AuthUserSnapshot {
+            id: user.user_id.clone(),
+            public_user_code: user.public_user_code,
+            username: user.username.clone(),
+            display_name: user.display_name,
+            phone_number_masked: user.phone_number_masked,
+            login_method: user.login_method,
+            binding_status: user.binding_status,
+            wechat_bound: user.wechat_bound,
+            token_version: user.token_version,
+        };
+        users_by_username.insert(
+            user.username,
+            StoredPasswordUserSnapshot {
+                user: auth_user,
+                password_hash: user.password_hash,
+                password_login_enabled: user.password_login_enabled,
+                phone_number: user
+                    .phone_number_e164
+                    .or_else(|| phone_identity_by_user_id.remove(&user.user_id)),
+            },
+        );
+    }
+
+    let mut sessions_by_id = std::collections::HashMap::new();
+    let mut session_id_by_refresh_token_hash = std::collections::HashMap::new();
+    for session in sessions {
+        let client_info = serde_json::from_str::<Value>(&session.client_info_json)
+            .map_err(|error| format!("refresh session ????? JSON ?????{error}"))?;
+        session_id_by_refresh_token_hash.insert(
+            session.refresh_token_hash.clone(),
+            session.session_id.clone(),
+        );
+        sessions_by_id.insert(
+            session.session_id.clone(),
+            StoredRefreshSessionSnapshot {
+                session: RefreshSessionSnapshot {
+                    session_id: session.session_id,
+                    user_id: session.user_id,
+                    refresh_token_hash: session.refresh_token_hash,
+                    issued_by_provider: session.issued_by_provider,
+                    client_info,
+                    expires_at: session.expires_at,
+                    revoked_at: session.revoked_at,
+                    created_at: session.created_at,
+                    updated_at: session.updated_at,
+                    last_seen_at: session.last_seen_at,
+                },
+            },
+        );
+    }
+
+    let snapshot = PersistentAuthStoreSnapshot {
+        next_user_id,
+        users_by_username,
+        phone_to_user_id,
+        sessions_by_id,
+        session_id_by_refresh_token_hash,
+        wechat_identity_by_provider_uid,
+        user_id_by_provider_union_id,
+    };
+    let snapshot_json =
+        serde_json::to_string_pretty(&snapshot).map_err(|error| format!("?????????????{error}"))?;
+
+    Ok(AuthStoreSnapshotRecord {
+        snapshot_json: Some(snapshot_json),
+        updated_at_micros: None,
+    })
+}
+
+fn clear_auth_target_tables(ctx: &ReducerContext) {
+    for row in ctx.db.refresh_session().iter().collect::<Vec<_>>() {
+        ctx.db
+            .refresh_session()
+            .session_id()
+            .delete(&row.session_id);
+    }
+    for row in ctx.db.auth_identity().iter().collect::<Vec<_>>() {
+        ctx.db
+            .auth_identity()
+            .identity_id()
+            .delete(&row.identity_id);
+    }
+    for row in ctx.db.user_account().iter().collect::<Vec<_>>() {
+        ctx.db.user_account().user_id().delete(&row.user_id);
+    }
+}
+
+fn sanitize_identity_component(value: &str) -> String {
+    let sanitized = value
+        .chars()
+        .map(|character| {
+            if character.is_ascii_alphanumeric() {
+                character
+            } else {
+                '_'
+            }
+        })
+        .collect::<String>();
+    sanitized.trim_matches('_').to_string()
+}
+
+#[derive(Deserialize, Serialize)]
+struct PersistentAuthStoreSnapshot {
+    #[serde(default = "default_next_user_id")]
+    next_user_id: u64,
+    users_by_username: std::collections::HashMap<String, StoredPasswordUserSnapshot>,
+    #[serde(default)]
+    phone_to_user_id: std::collections::HashMap<String, String>,
+    sessions_by_id: std::collections::HashMap<String, StoredRefreshSessionSnapshot>,
+    #[serde(default)]
+    session_id_by_refresh_token_hash: std::collections::HashMap<String, String>,
+    wechat_identity_by_provider_uid:
+        std::collections::HashMap<String, StoredWechatIdentitySnapshot>,
+    #[serde(default)]
+    user_id_by_provider_union_id: std::collections::HashMap<String, String>,
+}
+
+fn default_next_user_id() -> u64 {
+    1
+}
+
+#[derive(Deserialize, Serialize)]
+struct StoredPasswordUserSnapshot {
+    user: AuthUserSnapshot,
+    password_hash: String,
+    #[serde(default)]
+    password_login_enabled: bool,
+    phone_number: Option<String>,
+}
+
+#[derive(Deserialize, Serialize)]
+struct AuthUserSnapshot {
+    id: String,
+    public_user_code: String,
+    username: String,
+    display_name: String,
+    phone_number_masked: Option<String>,
+    login_method: String,
+    binding_status: String,
+    wechat_bound: bool,
+    token_version: u64,
+}
+
+#[derive(Deserialize, Serialize)]
+struct StoredWechatIdentitySnapshot {
+    user_id: String,
+    provider_uid: String,
+    provider_union_id: Option<String>,
+    display_name: Option<String>,
+    avatar_url: Option<String>,
+}
+
+#[derive(Deserialize, Serialize)]
+struct StoredRefreshSessionSnapshot {
+    session: RefreshSessionSnapshot,
+}
+
+#[derive(Deserialize, Serialize)]
+struct RefreshSessionSnapshot {
+    session_id: String,
+    user_id: String,
+    refresh_token_hash: String,
+    issued_by_provider: String,
+    client_info: Value,
+    expires_at: String,
+    revoked_at: Option<String>,
+    created_at: String,
+    updated_at: String,
+    last_seen_at: String,
+}