| @@ -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 { | |||
| @@ -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) | |||
| } | |||
| @@ -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 | |||
| @@ -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 | |||