| @@ -104,6 +104,15 @@ func GetCustomOAuthProviderBySlug(slug string) (*CustomOAuthProvider, error) { | |||||
| return &provider, nil | 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 | // CreateCustomOAuthProvider creates a new custom OAuth provider | ||||
| func CreateCustomOAuthProvider(provider *CustomOAuthProvider) error { | func CreateCustomOAuthProvider(provider *CustomOAuthProvider) error { | ||||
| if err := validateCustomOAuthProvider(provider); err != nil { | 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 | 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. | // FinalizeOAuthUserCreation performs post-transaction tasks for OAuth user creation. | ||||
| // This should be called after the transaction commits successfully. | // This should be called after the transaction commits successfully. | ||||
| func (user *User) FinalizeOAuthUserCreation(inviterId int) { | 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 获取所有同步用户 | // GetSyncedUsers 获取所有同步用户 | ||||
| func GetSyncedUsers() []User { | func GetSyncedUsers() []User { | ||||
| var users []User | var users []User | ||||
| @@ -10,9 +10,9 @@ import ( | |||||
| // UserOAuthBinding stores the binding relationship between users and custom OAuth providers | // UserOAuthBinding stores the binding relationship between users and custom OAuth providers | ||||
| type UserOAuthBinding struct { | type UserOAuthBinding struct { | ||||
| Id int `json:"id" gorm:"primaryKey"` | 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"` | CreatedAt time.Time `json:"created_at"` | ||||
| } | } | ||||
| @@ -104,6 +104,36 @@ func CreateUserOAuthBindingWithTx(tx *gorm.DB, binding *UserOAuthBinding) error | |||||
| return tx.Create(binding).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) | // UpdateUserOAuthBinding updates an existing OAuth binding (e.g., rebind to different OAuth account) | ||||
| func UpdateUserOAuthBinding(userId, providerId int, newProviderUserId string) error { | func UpdateUserOAuthBinding(userId, providerId int, newProviderUserId string) error { | ||||
| // Check if the new provider user ID is already taken by another user | // Check if the new provider user ID is already taken by another user | ||||