diff --git a/model/custom_oauth_provider.go b/model/custom_oauth_provider.go index 12b4d11..5f3cf04 100644 --- a/model/custom_oauth_provider.go +++ b/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 { diff --git a/model/region_migration_user_test.go b/model/region_migration_user_test.go new file mode 100644 index 0000000..cc29080 --- /dev/null +++ b/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) +} diff --git a/model/user.go b/model/user.go index e9b0da7..142e1de 100644 --- a/model/user.go +++ b/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 diff --git a/model/user_oauth_binding.go b/model/user_oauth_binding.go index 4921662..c6ce868 100644 --- a/model/user_oauth_binding.go +++ b/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