Просмотр исходного кода

feat(user-migration): add imported user and oauth migration helpers

master
fengsilin 1 неделю назад
Родитель
Сommit
38ba0dd24b
4 измененных файлов: 309 добавлений и 3 удалений
  1. +9
    -0
      model/custom_oauth_provider.go
  2. +187
    -0
      model/region_migration_user_test.go
  3. +80
    -0
      model/user.go
  4. +33
    -3
      model/user_oauth_binding.go

+ 9
- 0
model/custom_oauth_provider.go Просмотреть файл

@@ -104,6 +104,15 @@ func GetCustomOAuthProviderBySlug(slug string) (*CustomOAuthProvider, error) {
return &provider, nil
}

func GetCustomOAuthProvidersBySlugs(slugs []string) ([]*CustomOAuthProvider, error) {
var providers []*CustomOAuthProvider
if len(slugs) == 0 {
return providers, nil
}
err := DB.Where("slug IN ?", slugs).Order("id asc").Find(&providers).Error
return providers, err
}

// CreateCustomOAuthProvider creates a new custom OAuth provider
func CreateCustomOAuthProvider(provider *CustomOAuthProvider) error {
if err := validateCustomOAuthProvider(provider); err != nil {


+ 187
- 0
model/region_migration_user_test.go Просмотреть файл

@@ -0,0 +1,187 @@
package model

import (
"testing"

"github.com/QuantumNous/new-api/common"
"github.com/glebarez/sqlite"
"github.com/stretchr/testify/require"
"gorm.io/gorm"
)

func setupRegionMigrationUserDB(t *testing.T) *gorm.DB {
t.Helper()

db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
require.NoError(t, err)

sqlDB, err := db.DB()
require.NoError(t, err)
sqlDB.SetMaxOpenConns(1)

require.NoError(t, db.AutoMigrate(&User{}, &CustomOAuthProvider{}, &UserOAuthBinding{}))

return db
}

func TestInsertImportedUserWithHashedPassword_PreservesHash(t *testing.T) {
db := setupRegionMigrationUserDB(t)
orig := DB
DB = db
defer func() {
DB = orig
}()

user := &User{
Username: "ov-alice",
Password: "already-hashed-value",
Email: "alice@example.com",
DisplayName: "Alice",
Status: common.UserStatusEnabled,
Role: common.RoleCommonUser,
Group: "default",
Quota: 123,
AffCode: "A001",
CreatedAt: 1710000000,
}

require.NoError(t, InsertImportedUserWithHashedPassword(user))

var saved User
require.NoError(t, db.Where("username = ?", "ov-alice").First(&saved).Error)
require.Equal(t, "already-hashed-value", saved.Password)
require.Greater(t, saved.Id, 0)
}

func TestInsertImportedUserWithHashedPassword_RejectsExplicitId(t *testing.T) {
db := setupRegionMigrationUserDB(t)
orig := DB
DB = db
defer func() {
DB = orig
}()

user := &User{
Id: 10000001,
Username: "bad-import",
Password: "hash",
AffCode: "X001",
}

err := InsertImportedUserWithHashedPassword(user)
require.Error(t, err)
require.Contains(t, err.Error(), "Id must not be set")
}

func TestCopyOAuthBindingsIfMissing_Idempotent(t *testing.T) {
db := setupRegionMigrationUserDB(t)
orig := DB
DB = db
defer func() {
DB = orig
}()

require.NoError(t, db.Create(&CustomOAuthProvider{
Id: 10,
Name: "GitHub Enterprise",
Slug: "gh-enterprise",
Enabled: true,
ClientId: "x",
AuthorizationEndpoint: "a",
TokenEndpoint: "t",
UserInfoEndpoint: "u",
}).Error)
require.NoError(t, db.Create(&User{Id: 1, Username: "cn-user", Password: "hash", AffCode: "AFF1"}).Error)

bindings := []*UserOAuthBinding{{
ProviderId: 10,
ProviderUserId: "u-100",
}}

require.NoError(t, CopyOAuthBindingsIfMissing(1, bindings))
require.NoError(t, CopyOAuthBindingsIfMissing(1, bindings))

var count int64
require.NoError(t, db.Model(&UserOAuthBinding{}).Where("user_id = ?", 1).Count(&count).Error)
require.EqualValues(t, 1, count)
}

func TestConvertUserToSynced_Idempotent(t *testing.T) {
db := setupRegionMigrationUserDB(t)
orig := DB
DB = db
defer func() {
DB = orig
}()

require.NoError(t, db.Create(&User{
Id: 10000001,
Username: "ov-user",
Password: "hash",
AffCode: "AFF2",
Quota: 88,
Source: common.UserSourceLocal,
}).Error)

require.NoError(t, ConvertUserToSynced(10000001, 2001, 88))
require.NoError(t, ConvertUserToSynced(10000001, 2001, 88))

var saved User
require.NoError(t, db.First(&saved, 10000001).Error)
require.Equal(t, common.UserSourceSynced, saved.Source)
require.Equal(t, 2001, saved.RemoteUserId)
}

func TestConvertUserToSynced_DifferentRemoteUserIdFails(t *testing.T) {
db := setupRegionMigrationUserDB(t)
orig := DB
DB = db
defer func() {
DB = orig
}()

require.NoError(t, db.Create(&User{
Id: 10000002,
Username: "ov-user2",
Password: "hash",
AffCode: "AFF3",
Source: common.UserSourceLocal,
}).Error)

require.NoError(t, ConvertUserToSynced(10000002, 3001, 100))

err := ConvertUserToSynced(10000002, 9999, 100)
require.Error(t, err)
require.Contains(t, err.Error(), "already synced")
}

func TestListLocalUsersForMigration_FiltersLocalOnly(t *testing.T) {
db := setupRegionMigrationUserDB(t)
orig := DB
DB = db
defer func() {
DB = orig
}()

require.NoError(t, db.Create(&User{
Id: 10000001,
Username: "local-1",
Password: "h",
AffCode: "L1",
Source: common.UserSourceLocal,
}).Error)
require.NoError(t, db.Create(&User{
Id: 10000002,
Username: "synced-1",
Password: "h",
AffCode: "S1",
Source: common.UserSourceSynced,
RemoteUserId: 1,
}).Error)

users, total, err := ListLocalUsersForMigration(0, 20)
require.NoError(t, err)
require.EqualValues(t, 1, total)
require.Len(t, users, 1)
require.Equal(t, "local-1", users[0].Username)
}

+ 80
- 0
model/user.go Просмотреть файл

@@ -521,6 +521,31 @@ func (user *User) InsertWithTx(tx *gorm.DB, inviterId int) error {
return nil
}

// InsertImportedUserWithHashedPassword 专用于迁移导入,直接写入已有密码哈希。
func InsertImportedUserWithHashedPassword(user *User) error {
return InsertImportedUserWithHashedPasswordTx(DB, user)
}

func InsertImportedUserWithHashedPasswordTx(tx *gorm.DB, user *User) error {
if user == nil {
return errors.New("user is nil")
}
if user.Id != 0 {
return errors.New("InsertImportedUserWithHashedPassword: Id must not be set, let DB auto-assign")
}
if user.AffCode == "" {
user.AffCode = common.GetRandomString(4)
}
if user.CreatedAt == 0 {
user.CreatedAt = common.GetTimestamp()
}
if user.Setting == "" {
defaultSetting := dto.UserSetting{}
user.SetSetting(defaultSetting)
}
return tx.Create(user).Error
}

// FinalizeOAuthUserCreation performs post-transaction tasks for OAuth user creation.
// This should be called after the transaction commits successfully.
func (user *User) FinalizeOAuthUserCreation(inviterId int) {
@@ -1020,6 +1045,61 @@ func DeltaUpdateUserQuota(id int, delta int) (err error) {
}
}

func ListLocalUsersForMigration(page, pageSize int) ([]*User, int64, error) {
if page < 0 {
page = 0
}
if pageSize <= 0 {
pageSize = 20
}

var users []*User
var total int64
query := DB.Model(&User{}).Where("source = ?", common.UserSourceLocal)
if err := query.Count(&total).Error; err != nil {
return nil, 0, err
}
err := query.Order("id asc").Offset(page * pageSize).Limit(pageSize).Find(&users).Error
return users, total, err
}

func HasOVSyncedCopyByRemoteId(cnUserId int) (bool, error) {
var count int64
err := DB.Model(&User{}).
Where("remote_user_id = ? AND source = ?", cnUserId, common.UserSourceSynced).
Count(&count).Error
return count > 0, err
}

func ConvertUserToSynced(userId, remoteUserId, syncedQuota int) error {
var user User
if err := DB.Where("id = ?", userId).First(&user).Error; err != nil {
return err
}

now := common.GetTimestamp()
if user.Source == common.UserSourceSynced {
if user.RemoteUserId != remoteUserId {
return fmt.Errorf("user %d is already synced to remote_user_id=%d, cannot change to %d",
userId, user.RemoteUserId, remoteUserId)
}
if user.SyncedQuota == syncedQuota {
return nil
}
return DB.Model(&User{}).Where("id = ?", userId).Updates(map[string]any{
"synced_quota": syncedQuota,
"last_sync_at": now,
}).Error
}

return DB.Model(&User{}).Where("id = ?", userId).Updates(map[string]any{
"source": common.UserSourceSynced,
"remote_user_id": remoteUserId,
"synced_quota": syncedQuota,
"last_sync_at": now,
}).Error
}

// GetSyncedUsers 获取所有同步用户
func GetSyncedUsers() []User {
var users []User


+ 33
- 3
model/user_oauth_binding.go Просмотреть файл

@@ -10,9 +10,9 @@ import (
// UserOAuthBinding stores the binding relationship between users and custom OAuth providers
type UserOAuthBinding struct {
Id int `json:"id" gorm:"primaryKey"`
UserId int `json:"user_id" gorm:"not null;uniqueIndex:ux_user_provider"` // User ID - one binding per user per provider
ProviderId int `json:"provider_id" gorm:"not null;uniqueIndex:ux_user_provider;uniqueIndex:ux_provider_userid"` // Custom OAuth provider ID
ProviderUserId string `json:"provider_user_id" gorm:"type:varchar(256);not null;uniqueIndex:ux_provider_userid"` // User ID from OAuth provider - one OAuth account per provider
UserId int `json:"user_id" gorm:"not null;uniqueIndex:ux_user_provider"` // User ID - one binding per user per provider
ProviderId int `json:"provider_id" gorm:"not null;uniqueIndex:ux_user_provider;uniqueIndex:ux_provider_userid"` // Custom OAuth provider ID
ProviderUserId string `json:"provider_user_id" gorm:"type:varchar(256);not null;uniqueIndex:ux_provider_userid"` // User ID from OAuth provider - one OAuth account per provider
CreatedAt time.Time `json:"created_at"`
}

@@ -104,6 +104,36 @@ func CreateUserOAuthBindingWithTx(tx *gorm.DB, binding *UserOAuthBinding) error
return tx.Create(binding).Error
}

func CopyOAuthBindingsIfMissing(userId int, bindings []*UserOAuthBinding) error {
return CopyOAuthBindingsIfMissingTx(DB, userId, bindings)
}

func CopyOAuthBindingsIfMissingTx(tx *gorm.DB, userId int, bindings []*UserOAuthBinding) error {
for _, binding := range bindings {
if binding == nil {
continue
}
var existing UserOAuthBinding
result := tx.Where("user_id = ? AND provider_id = ?", userId, binding.ProviderId).Limit(1).Find(&existing)
if result.Error != nil {
return result.Error
}
if result.RowsAffected > 0 {
continue
}

newBinding := &UserOAuthBinding{
UserId: userId,
ProviderId: binding.ProviderId,
ProviderUserId: binding.ProviderUserId,
}
if err := CreateUserOAuthBindingWithTx(tx, newBinding); err != nil {
return err
}
}
return nil
}

// UpdateUserOAuthBinding updates an existing OAuth binding (e.g., rebind to different OAuth account)
func UpdateUserOAuthBinding(userId, providerId int, newProviderUserId string) error {
// Check if the new provider user ID is already taken by another user


Загрузка…
Отмена
Сохранить