package controller import ( "errors" "fmt" "net/http" "net/http/httptest" "os" "sort" "strconv" "strings" "sync" "testing" "time" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/i18n" "github.com/QuantumNous/new-api/middleware" "github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/relaykit/dto" "github.com/QuantumNous/new-api/service" "github.com/QuantumNous/new-api/service/authz" "github.com/alicebob/miniredis/v2" "github.com/go-redis/redis/v8" "github.com/gin-gonic/gin" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "gorm.io/gorm" ) func setupManageUserTestDB(t *testing.T) *gorm.DB { t.Helper() require.NoError(t, i18n.Init()) previousDB, previousLogDB := model.DB, model.LOG_DB previousRedisEnabled, previousSecret := common.RedisEnabled, common.SessionSecret previousMainDatabaseType, previousLogDatabaseType := common.MainDatabaseType(), common.LogDatabaseType() dialect := os.Getenv("TEST_MANAGE_USER_DIALECT") if dialect == "" { dialect = "sqlite" } databaseTypes := map[string]common.DatabaseType{ "sqlite": common.DatabaseTypeSQLite, "mysql": common.DatabaseTypeMySQL, "postgres": common.DatabaseTypePostgreSQL, } require.Contains(t, databaseTypes, dialect) dsn := os.Getenv("TEST_" + strings.ToUpper(dialect) + "_DSN") db, _ := newAuditTestDatabase(t, dialect, dsn) logDB := db if os.Getenv("TEST_MANAGE_USER_SEPARATE_LOG_DB") == "1" { logDB, _ = newAuditTestDatabase(t, dialect, dsn) } model.DB, model.LOG_DB = db, logDB common.RedisEnabled = false common.SessionSecret = "manage-user-test-secret" common.SetDatabaseTypes(databaseTypes[dialect], databaseTypes[dialect]) t.Cleanup(func() { model.DB, model.LOG_DB = previousDB, previousLogDB common.RedisEnabled, common.SessionSecret = previousRedisEnabled, previousSecret common.SetDatabaseTypes(previousMainDatabaseType, previousLogDatabaseType) sqlDB, err := db.DB() if err == nil { _ = sqlDB.Close() } if logDB != db { sqlLogDB, err := logDB.DB() if err == nil { _ = sqlLogDB.Close() } } }) require.NoError(t, db.AutoMigrate(&model.User{}, &model.UserSession{}, &model.CasbinRule{}, &model.AuthzRole{}, &model.AuthFlow{}, &model.TwoFA{}, &model.PasskeyCredential{}, &model.UserAccessToken{})) require.NoError(t, logDB.AutoMigrate(&model.Log{}, &model.AuditLog{})) versionQuery := "SELECT version()" if dialect == "sqlite" { versionQuery = "SELECT sqlite_version()" } var version string require.NoError(t, db.Raw(versionQuery).Scan(&version).Error) t.Logf("database: %s %s, separate log database: %v", dialect, version, logDB != db) return db } func performManageUserRequest(t *testing.T, body string) *httptest.ResponseRecorder { t.Helper() return performVerifiedManageUserRequest(t, body, service.AuthIdentity{UserID: 9999}, "") } // performVerifiedManageUserRequest drives ManageUser as the root operator from // a dashboard login session, carrying the step-up proof status or role changes need. func performVerifiedManageUserRequest(t *testing.T, body string, identity service.AuthIdentity, proof string) *httptest.ResponseRecorder { t.Helper() gin.SetMode(gin.TestMode) recorder := httptest.NewRecorder() c, _ := gin.CreateTestContext(recorder) c.Request = httptest.NewRequest(http.MethodPost, "/api/user/manage", strings.NewReader(body)) c.Request.Header.Set("Content-Type", "application/json") if proof != "" { c.Request.Header.Set("X-Security-Proof", proof) } c.Set("id", identity.UserID) c.Set("role", common.RoleRootUser) c.Set("username", "root-operator") c.Set("session_id", identity.SessionID) c.Set("auth_version", identity.UserAuthVersion) c.Set("session_version", identity.SessionVersion) c.Set(common.RequestIdKey, "quota-test-request") ManageUser(c) return recorder } // manageUserProof signs the root operator in and issues a password-method proof // bound to one ManageUser operation. func manageUserProof(t *testing.T, db *gorm.DB, operation service.VerificationOperation) (service.AuthIdentity, string) { t.Helper() operator := createQuotaTestOperator(t, db, common.RoleRootUser) require.NoError(t, model.PublishUserAuthCache(operator.Id)) bundle, err := service.CreateLoginSession(operator.Id, "password", "127.0.0.1", "manage-user-test") require.NoError(t, err) identity, err := service.ParseAccessToken(bundle.AccessToken) require.NoError(t, err) binding, err := service.BindVerificationOperation(operation) require.NoError(t, err) proof, _, err := service.IssueSecurityProof(identity, service.VerificationMethodPassword, binding) require.NoError(t, err) return identity, proof } func TestManageUserDisableAdvancesAuthVersionOnceAndRevokesSession(t *testing.T) { db := setupManageUserTestDB(t) now := time.Now().Unix() user := model.User{ Username: "managed-disable-user", Password: "password", Role: common.RoleCommonUser, Status: common.UserStatusEnabled, Group: "default", AuthVersion: 1, } require.NoError(t, db.Create(&user).Error) require.NoError(t, db.Create(&model.UserSession{ SID: "managed-disable-session", UserID: user.Id, Version: 1, UserAuthVersion: 1, Status: model.UserSessionStatusActive, RefreshHash: "refresh-hash", LoginMethod: "password", LastActiveAt: now, ExpiresAt: now + 3600, }).Error) identity, proof := manageUserProof(t, db, service.VerificationOperation{Scope: service.VerificationScopeAdminUserManage, Context: []byte(fmt.Sprintf(`{"user_id":%d,"action":"disable"}`, user.Id))}) recorder := performVerifiedManageUserRequest(t, fmt.Sprintf(`{"id":%d,"action":"disable"}`, user.Id), identity, proof) assert.Equal(t, http.StatusOK, recorder.Code) assert.Contains(t, recorder.Body.String(), `"success":true`) var updated model.User require.NoError(t, db.First(&updated, user.Id).Error) assert.Equal(t, common.UserStatusDisabled, updated.Status) assert.EqualValues(t, 2, updated.AuthVersion) var session model.UserSession require.NoError(t, db.First(&session, "sid = ?", "managed-disable-session").Error) assert.Equal(t, model.UserSessionStatusRevoked, session.Status) } func TestManageUserDemoteAdvancesAuthVersionAndRevokesSessionsOnce(t *testing.T) { db := setupManageUserTestDB(t) previousMaster := common.IsMasterNode common.IsMasterNode = false t.Cleanup(func() { common.IsMasterNode = previousMaster }) require.NoError(t, authz.Init(db)) now := time.Now().Unix() user := model.User{ Username: "managed-demote-user", Password: "password", Role: common.RoleAdminUser, Status: common.UserStatusEnabled, Group: "default", AuthVersion: 1, } require.NoError(t, db.Create(&user).Error) for _, sid := range []string{"managed-demote-session-one", "managed-demote-session-two"} { require.NoError(t, db.Create(&model.UserSession{ SID: sid, UserID: user.Id, Version: 1, UserAuthVersion: 1, Status: model.UserSessionStatusActive, RefreshHash: "refresh-" + sid, LoginMethod: "password", LastActiveAt: now, ExpiresAt: now + 3600, }).Error) } sessionUpdateCount := 0 require.NoError(t, db.Callback().Update().Before("gorm:update").Register("test:count_demote_session_updates", func(tx *gorm.DB) { if tx.Statement != nil && tx.Statement.Table == "user_sessions" { sessionUpdateCount++ } })) identity, proof := manageUserProof(t, db, service.VerificationOperation{Scope: service.VerificationScopeAdminUserManage, Context: []byte(fmt.Sprintf(`{"user_id":%d,"action":"demote"}`, user.Id))}) recorder := performVerifiedManageUserRequest(t, fmt.Sprintf(`{"id":%d,"action":"demote"}`, user.Id), identity, proof) assert.Equal(t, http.StatusOK, recorder.Code) assert.Contains(t, recorder.Body.String(), `"success":true`) var updated model.User require.NoError(t, db.First(&updated, user.Id).Error) assert.Equal(t, common.RoleCommonUser, updated.Role) assert.EqualValues(t, 2, updated.AuthVersion) var sessions []model.UserSession require.NoError(t, db.Where("user_id = ?", user.Id).Order("sid asc").Find(&sessions).Error) require.Len(t, sessions, 2) for _, session := range sessions { assert.Equal(t, model.UserSessionStatusRevoked, session.Status) assert.Equal(t, "admin_demote", session.RevokedReason) } assert.Equal(t, 1, sessionUpdateCount) } func TestManageUserDeleteReturnsImmediatelyAndUnknownActionFails(t *testing.T) { db := setupManageUserTestDB(t) deleted := model.User{ Username: "managed-delete-user", Password: "password", Role: common.RoleCommonUser, Status: common.UserStatusEnabled, Group: "default", AuthVersion: 1, AffCode: "delete-aff", } require.NoError(t, db.Create(&deleted).Error) identity, proof := manageUserProof(t, db, service.VerificationOperation{Scope: service.VerificationScopeAdminUserDelete, Context: []byte(fmt.Sprintf(`{"user_id":%d}`, deleted.Id))}) recorder := performVerifiedManageUserRequest(t, fmt.Sprintf(`{"id":%d,"action":"delete"}`, deleted.Id), identity, proof) assert.Contains(t, recorder.Body.String(), `"success":true`) var deletedCount int64 require.NoError(t, db.Unscoped().Model(&model.User{}).Where("id = ? AND deleted_at IS NOT NULL", deleted.Id).Count(&deletedCount).Error) assert.EqualValues(t, 1, deletedCount) unchanged := model.User{ Username: "managed-unknown-user", Password: "password", Role: common.RoleCommonUser, Status: common.UserStatusEnabled, Group: "default", AuthVersion: 1, AffCode: "unknown-aff", } require.NoError(t, db.Create(&unchanged).Error) recorder = performManageUserRequest(t, fmt.Sprintf(`{"id":%d,"action":"unknown"}`, unchanged.Id)) assert.Contains(t, recorder.Body.String(), `"success":false`) require.NoError(t, db.First(&unchanged, unchanged.Id).Error) assert.EqualValues(t, 1, unchanged.AuthVersion) assert.Equal(t, common.UserStatusEnabled, unchanged.Status) } func createQuotaTestOperator(t *testing.T, db *gorm.DB, role int) model.User { t.Helper() if role == 0 { role = common.RoleRootUser } operator := model.User{Id: 9999, Username: "root-operator", Password: "root-operator-hash", Role: role, Status: common.UserStatusEnabled, AuthVersion: 1, AffCode: "root-operator-aff"} require.NoError(t, db.Create(&operator).Error) return operator } func TestManageUserQuotaRecordsTopupAndAudit(t *testing.T) { for _, tc := range []struct { name, mode, action, content string value, wantQuota int }{ {"add", "add", "user.quota_add", "Increased user quota by 500", 500, 1500}, {"subtract", "subtract", "user.quota_subtract", "Decreased user quota by 500", 500, 500}, {"override_up", "override", "user.quota_override", "Overrode user quota from 1000 to 1500", 1500, 1500}, {"override_down", "override", "user.quota_override", "Overrode user quota from 1000 to 500", 500, 500}, {"override_unchanged", "override", "user.quota_override", "Overrode user quota from 1000 to 1000", 1000, 1000}, {"override_zero", "override", "user.quota_override", "Overrode user quota from 1000 to 0", 0, 0}, {"override_negative", "override", "user.quota_override", "Overrode user quota from 1000 to -1", -1, -1}, } { t.Run(tc.name, func(t *testing.T) { db := setupManageUserTestDB(t) user := model.User{Username: "quota-owner", Role: common.RoleCommonUser, Quota: 1000, AffCode: "quota-owner-aff"} require.NoError(t, db.Create(&user).Error) createQuotaTestOperator(t, db, common.RoleRootUser) recorder := performManageUserRequest(t, fmt.Sprintf(`{"id":%d,"action":"add_quota","mode":%q,"value":%d}`, user.Id, tc.mode, tc.value)) assert.Equal(t, http.StatusOK, recorder.Code) require.Contains(t, recorder.Body.String(), `"success":true`) require.NoError(t, db.First(&user, user.Id).Error) assert.Equal(t, tc.wantQuota, user.Quota) logs, total, err := model.GetAllLogs(model.LogFilter{Type: model.LogTypeTopup}, 0, 20) require.NoError(t, err) assert.EqualValues(t, 1, total) require.Len(t, logs, 1) assert.Equal(t, user.Id, logs[0].UserId) assert.Equal(t, user.Username, logs[0].Username) assert.Equal(t, tc.content, logs[0].Content) model.FormatAdminLogs(logs) var other model.AuditOther require.NoError(t, common.UnmarshalJsonStr(logs[0].Other, &other)) assert.Equal(t, &model.AuditAdminInfo{AdminID: 9999, AdminUsername: "root-operator", AdminRole: common.RoleRootUser, AuthMethod: "session"}, other.AdminInfo) logs, total, err = model.GetUserLogs(user.Id, model.LogFilter{Type: model.LogTypeTopup}, 0, 20) require.NoError(t, err) assert.EqualValues(t, 1, total) require.Len(t, logs, 1) other = model.AuditOther{} require.NoError(t, common.UnmarshalJsonStr(logs[0].Other, &other)) assert.Nil(t, other.AdminInfo) require.NotNil(t, other.Op) assert.Equal(t, tc.action, other.Op.Action) params, err := common.Marshal(other.Op.Params) require.NoError(t, err) expectedParams := model.AuditFields{"target_user_id": user.Id, "target_username": user.Username, "mode": tc.mode, "requested_quota": tc.value, "from": 1000, "to": tc.wantQuota} if tc.mode != "override" { expectedParams["quota"] = tc.value } expected, err := common.Marshal(expectedParams) require.NoError(t, err) assert.JSONEq(t, string(expected), string(params)) assert.Equal(t, "quota-test-request", logs[0].RequestId) assert.Empty(t, logs[0].Ip, "recipient logs must not disclose the administrator IP") logs, total, err = model.GetUserLogs(9999, model.LogFilter{Type: model.LogTypeTopup}, 0, 20) require.NoError(t, err) assert.Zero(t, total) assert.Empty(t, logs) var audits []model.AuditLog require.NoError(t, model.LOG_DB.Find(&audits).Error) require.Len(t, audits, 1) assert.Equal(t, 9999, audits[0].UserId) assert.Equal(t, "root-operator", audits[0].Username) assert.Equal(t, tc.action, audits[0].Action) assert.True(t, audits[0].Success) params, err = common.Marshal(audits[0].Other.Op.Params) require.NoError(t, err) assert.JSONEq(t, string(expected), string(params)) assert.Equal(t, "quota-test-request", audits[0].RequestId) }) } } func TestManageUserQuotaFailuresDoNotRecordTopup(t *testing.T) { for _, tc := range []struct { name, mode string value int failUpdate bool }{ {"zero_add", "add", 0, false}, {"negative_subtract", "subtract", -1, false}, {"invalid_mode", "invalid", 500, false}, {"add_update_error", "add", 500, true}, {"subtract_update_error", "subtract", 500, true}, {"override_update_error", "override", 500, true}, } { t.Run(tc.name, func(t *testing.T) { db := setupManageUserTestDB(t) createQuotaTestOperator(t, db, common.RoleRootUser) user := model.User{Username: "quota-owner", Role: common.RoleCommonUser, Quota: 1000} require.NoError(t, db.Create(&user).Error) if tc.failUpdate { require.NoError(t, db.Callback().Update().Before("gorm:update").Register("test:fail_quota_update", func(tx *gorm.DB) { if tx.Statement.Table == "users" { tx.AddError(errors.New("quota update unavailable")) } })) } recorder := performManageUserRequest(t, fmt.Sprintf(`{"id":%d,"action":"add_quota","mode":%q,"value":%d}`, user.Id, tc.mode, tc.value)) assert.Contains(t, recorder.Body.String(), `"success":false`) require.NoError(t, db.First(&user, user.Id).Error) assert.Equal(t, 1000, user.Quota) var logCount, auditCount int64 require.NoError(t, model.LOG_DB.Model(&model.Log{}).Count(&logCount).Error) require.NoError(t, model.LOG_DB.Model(&model.AuditLog{}).Count(&auditCount).Error) assert.Zero(t, logCount) assert.EqualValues(t, 1, auditCount) var audit model.AuditLog require.NoError(t, model.LOG_DB.First(&audit).Error) assert.False(t, audit.Success) params, err := common.Marshal(audit.Other.Op.Params) require.NoError(t, err) assert.NotContains(t, string(params), `"from"`) assert.NotContains(t, string(params), `"to"`) assert.NotContains(t, string(params), `"target_username"`) assert.NotContains(t, string(params), "quota update unavailable") reason := "invalid_parameters" if tc.failUpdate { reason = "database_error" } assert.Contains(t, string(params), `"failure_reason":"`+reason+`"`) assert.Contains(t, string(params), fmt.Sprintf(`"requested_quota":%d`, tc.value)) }) } } func TestManageUserQuotaLogFailureKeepsSuccessfulAdjustment(t *testing.T) { for _, failedTable := range []string{"logs", "audit_logs"} { t.Run(failedTable, func(t *testing.T) { db := setupManageUserTestDB(t) createQuotaTestOperator(t, db, common.RoleRootUser) user := model.User{Username: "quota-owner", Role: common.RoleCommonUser, Quota: 1000} require.NoError(t, db.Create(&user).Error) require.NoError(t, model.LOG_DB.Callback().Create().Before("gorm:create").Register("test:fail_quota_log", func(tx *gorm.DB) { if tx.Statement.Table == failedTable { tx.AddError(errors.New("quota log unavailable")) } })) recorder := performManageUserRequest(t, fmt.Sprintf(`{"id":%d,"action":"add_quota","mode":"add","value":500}`, user.Id)) assert.Contains(t, recorder.Body.String(), `"success":true`) require.NoError(t, db.First(&user, user.Id).Error) assert.Equal(t, 1500, user.Quota) var logCount, auditCount int64 require.NoError(t, model.LOG_DB.Model(&model.Log{}).Count(&logCount).Error) require.NoError(t, model.LOG_DB.Model(&model.AuditLog{}).Count(&auditCount).Error) if failedTable == "logs" { assert.Zero(t, logCount) assert.EqualValues(t, 1, auditCount) } else { assert.EqualValues(t, 1, logCount) assert.Zero(t, auditCount) } }) } } func TestManageUserQuotaTargetsAndWalletBounds(t *testing.T) { for _, tc := range []struct { name, mode, reason string before, value, targetID, targetRole, operatorRole int deleted, failRead bool }{ {name: "zero_id", mode: "add", value: 1, reason: "invalid_parameters"}, {name: "negative_id", mode: "subtract", value: 1, targetID: -1, reason: "invalid_parameters"}, {name: "missing", mode: "override", value: 1, targetID: 12345, reason: "target_not_found"}, {name: "deleted", mode: "add", value: 1, targetID: 1, deleted: true, reason: "target_not_found"}, {name: "peer_admin", mode: "add", value: 1, targetID: 1, targetRole: common.RoleAdminUser, operatorRole: common.RoleAdminUser, reason: "permission_denied"}, {name: "higher_role", mode: "override", value: 1, targetID: 1, targetRole: common.RoleRootUser, operatorRole: common.RoleAdminUser, reason: "permission_denied"}, {name: "read_error", mode: "add", value: 1, targetID: 1, failRead: true, reason: "database_error"}, {name: "add_overflow", mode: "add", before: common.MaxWalletQuota, value: 1, targetID: 1, reason: "quota_limit_exceeded"}, {name: "subtract_underflow", mode: "subtract", before: -common.MaxWalletQuota, value: 1, targetID: 1, reason: "quota_limit_exceeded"}, {name: "oversized_add", mode: "add", value: common.MaxWalletQuota + 1, targetID: 1, reason: "quota_limit_exceeded"}, {name: "oversized_subtract", mode: "subtract", value: common.MaxWalletQuota + 1, targetID: 1, reason: "quota_limit_exceeded"}, {name: "oversized_override", mode: "override", value: -common.MaxWalletQuota - 1, targetID: 1, reason: "quota_limit_exceeded"}, } { t.Run(tc.name, func(t *testing.T) { db := setupManageUserTestDB(t) createQuotaTestOperator(t, db, tc.operatorRole) user := model.User{Id: 1, Username: "private-target-name", Role: tc.targetRole, Quota: tc.before} require.NoError(t, db.Create(&user).Error) if tc.deleted { require.NoError(t, db.Delete(&user).Error) } if tc.failRead { require.NoError(t, db.Callback().Query().Before("gorm:query").Register("test:quota_read_error", func(tx *gorm.DB) { if tx.Statement.Table == "users" && len(tx.Statement.Selects) == 0 { tx.AddError(errors.New("private database error")) } })) } recorder := httptest.NewRecorder() c, _ := gin.CreateTestContext(recorder) c.Request = httptest.NewRequest(http.MethodPost, "/api/user/manage", strings.NewReader(fmt.Sprintf(`{"id":%d,"action":"add_quota","mode":%q,"value":%d}`, tc.targetID, tc.mode, tc.value))) role := tc.operatorRole if role == 0 { role = common.RoleRootUser } c.Set("id", 9999) c.Set("role", role) ManageUser(c) require.Contains(t, recorder.Body.String(), `"success":false`) if tc.failRead { require.NoError(t, db.Callback().Query().Remove("test:quota_read_error")) } require.NoError(t, db.Unscoped().First(&user, user.Id).Error) assert.Equal(t, tc.before, user.Quota) var audits []model.AuditLog require.NoError(t, model.LOG_DB.Find(&audits).Error) require.Len(t, audits, 1) assert.False(t, audits[0].Success) params, err := common.Marshal(audits[0].Other.Op.Params) require.NoError(t, err) expected, err := common.Marshal(model.AuditFields{"target_user_id": tc.targetID, "mode": tc.mode, "requested_quota": tc.value, "failure_reason": tc.reason}) require.NoError(t, err) assert.JSONEq(t, string(expected), string(params)) assert.NotContains(t, audits[0].Content, user.Username) var count int64 require.NoError(t, model.LOG_DB.Model(&model.Log{}).Count(&count).Error) assert.Zero(t, count) }) } } func TestManageUserQuotaMiddlewareKeepsOneOperationPerRequest(t *testing.T) { db := setupManageUserTestDB(t) pat := "quota-middleware-test-token" operator := model.User{Id: 9999, Username: "root-operator", Role: common.RoleRootUser, Status: common.UserStatusEnabled, AuthVersion: 1, AccessToken: &pat, Quota: 1000} require.NoError(t, db.Create(&operator).Error) router := gin.New() router.Use(middleware.RequestId(), middleware.AccessTokenAudit()) router.POST("/api/user/manage", middleware.AdminAuth(), ManageUser) for _, tc := range []struct { body, action string success bool }{ {`{"id":9999,"action":"add_quota","mode":"add","value":100}`, "user.quota_add", true}, {`{"id":9999,"action":"add_quota","mode":"subtract","value":0}`, "user.quota_subtract", false}, {`{"id":9999,"action":"add_quota","mode":"invalid","value":1}`, "generic", false}, {`{"id":`, "generic", false}, } { recorder := httptest.NewRecorder() request := httptest.NewRequest(http.MethodPost, "/api/user/manage", strings.NewReader(tc.body)) request.Header.Set("Authorization", "Bearer "+pat) request.Header.Set("Content-Type", "application/json") router.ServeHTTP(recorder, request) assert.Equal(t, http.StatusOK, recorder.Code) assert.Contains(t, recorder.Body.String(), fmt.Sprintf(`"success":%t`, tc.success)) requestID := recorder.Header().Get(common.RequestIdKey) require.NotEmpty(t, requestID) var audits []model.AuditLog require.NoError(t, model.LOG_DB.Where("request_id = ?", requestID).Find(&audits).Error) require.Len(t, audits, 2, "one operation audit and one PAT request audit") for _, audit := range audits { assert.Equal(t, tc.success, audit.Success) if audit.Category != model.AuditCategoryOperation { continue } assert.Equal(t, tc.action, audit.Action) assert.Equal(t, operator.Id, audit.UserId) assert.Equal(t, "/api/user/manage", audit.Route) if tc.success { params, err := common.Marshal(audit.Other.Op.Params) require.NoError(t, err) assert.Contains(t, string(params), `"target_user_id":9999`) assert.Contains(t, string(params), `"target_username":"root-operator"`) } } var count int64 require.NoError(t, model.LOG_DB.Model(&model.Log{}).Where("request_id = ?", requestID).Count(&count).Error) if tc.success { assert.EqualValues(t, 1, count) } else { assert.Zero(t, count) } } require.NoError(t, db.First(&operator, operator.Id).Error) assert.Equal(t, 1100, operator.Quota) } func TestManageUserQuotaConcurrentSnapshots(t *testing.T) { db := setupManageUserTestDB(t) user := model.User{Username: "concurrent-quota", Quota: 1000} require.NoError(t, db.Create(&user).Error) var ready sync.WaitGroup ready.Add(2) release := make(chan struct{}) require.NoError(t, db.Callback().Query().Before("gorm:query").Register("test:concurrent_quota_start", func(tx *gorm.DB) { if tx.Statement.Table == "users" { ready.Done() <-release } })) type result struct { adjustment *model.UserQuotaAdjustment err error value int } results := make(chan result, 2) for _, value := range []int{10, 20} { go func(value int) { adjustment, err := model.AdjustUserQuota(user.Id, common.RoleRootUser, "add", value) results <- result{adjustment, err, value} }(value) } ready.Wait() close(release) var committed []model.UserQuotaAdjustment for range 2 { result := <-results if result.err != nil { require.True(t, common.UsingMainDatabase(common.DatabaseTypeSQLite), "row-locking databases must serialize both adjustments: %v", result.err) assert.Contains(t, strings.ToLower(result.err.Error()), "locked") assert.Nil(t, result.adjustment) continue } require.NotNil(t, result.adjustment) assert.Equal(t, result.value, result.adjustment.After-result.adjustment.Before) committed = append(committed, *result.adjustment) } require.NoError(t, db.Callback().Query().Remove("test:concurrent_quota_start")) require.NotEmpty(t, committed) sort.Slice(committed, func(i, j int) bool { return committed[i].Before < committed[j].Before }) balance := 1000 for _, adjustment := range committed { assert.Equal(t, balance, adjustment.Before) balance = adjustment.After } require.NoError(t, db.First(&user, user.Id).Error) assert.Equal(t, balance, user.Quota) } func TestManageUserQuotaCacheUsesCommittedIntegerDifference(t *testing.T) { for _, tc := range []struct { name, mode string before, cached, value, after, wantCached int failUpdate, failCache, missingCache bool }{ {name: "add_preserves_reservations", mode: "add", before: 1000, cached: 900, value: 500, after: 1500, wantCached: 1400}, {name: "subtract_preserves_reservations", mode: "subtract", before: 1000, cached: 900, value: 500, after: 500, wantCached: 400}, {name: "override_preserves_reservations", mode: "override", before: 1000, cached: 900, value: 2000, after: 2000, wantCached: 1900}, {name: "large_odd_difference", mode: "override", before: common.MaxWalletQuota - 1, cached: common.MaxWalletQuota - 1, value: -common.MaxWalletQuota, after: -common.MaxWalletQuota, wantCached: -common.MaxWalletQuota}, {name: "rollback_does_not_change_cache", mode: "subtract", before: 1000, cached: 900, value: 500, after: 1000, wantCached: 900, failUpdate: true}, {name: "cache_error_keeps_committed_change", mode: "add", before: 1000, value: 500, after: 1500, failCache: true}, {name: "missing_cache_is_not_partially_created", mode: "add", before: 1000, value: 500, after: 1500, missingCache: true}, } { t.Run(tc.name, func(t *testing.T) { db := setupManageUserTestDB(t) operator := createQuotaTestOperator(t, db, common.RoleRootUser) server := miniredis.RunT(t) oldRDB := common.RDB common.RDB = redis.NewClient(&redis.Options{Addr: server.Addr(), MaxRetries: -1}) common.RedisEnabled = true t.Cleanup(func() { _ = common.RDB.Close(); common.RDB = oldRDB }) user := model.User{Username: "cached-quota", Quota: tc.before, AuthVersion: 1} require.NoError(t, db.Create(&user).Error) cache, err := model.GetUserCache(user.Id) require.NoError(t, err) assert.Equal(t, tc.before, cache.Quota) _, err = model.GetUserCache(operator.Id) require.NoError(t, err) keys := server.Keys() var quotaKey string for _, key := range keys { if server.HGet(key, "Id") == strconv.Itoa(user.Id) && server.HGet(key, "Quota") != "" { quotaKey = key break } } require.NotEmpty(t, quotaKey) server.HSet(quotaKey, "Quota", strconv.Itoa(tc.cached)) if tc.missingCache { server.Del(quotaKey) } if tc.failCache { server.SetError("ERR quota cache unavailable") } if tc.failUpdate { require.NoError(t, db.Callback().Update().After("gorm:update").Register("test:cache_quota_rollback", func(tx *gorm.DB) { if tx.Statement.Table == "users" { tx.AddError(errors.New("quota rollback")) } })) } recorder := performManageUserRequest(t, fmt.Sprintf(`{"id":%d,"action":"add_quota","mode":%q,"value":%d}`, user.Id, tc.mode, tc.value)) assert.Contains(t, recorder.Body.String(), fmt.Sprintf(`"success":%t`, !tc.failUpdate)) require.NoError(t, db.First(&user, user.Id).Error) assert.Equal(t, tc.after, user.Quota) if tc.missingCache { // Log username lookup may hydrate the whole user after commit. if server.Exists(quotaKey) { assert.Equal(t, strconv.Itoa(user.Id), server.HGet(quotaKey, "Id")) assert.NotEmpty(t, server.HGet(quotaKey, "CacheSchema")) assert.Equal(t, strconv.Itoa(tc.after), server.HGet(quotaKey, "Quota")) } } else if !tc.failCache { assert.Equal(t, strconv.Itoa(tc.wantCached), server.HGet(quotaKey, "Quota")) } }) } } // Saving notification settings starts from what the user stored: fields the // form does not manage stay, and switches an older client leaves out keep // their stored values. Only administrators change the channel and upstream // switches, and a threshold of 0 means the site default. func TestUpdateUserSettingKeepsWhatTheRequestLeavesOut(t *testing.T) { on, off := true, false stored := dto.UserSetting{ NotifyType: dto.NotifyTypeBark, BarkUrl: "https://api.day.app/key/{{title}}/{{content}}", QuotaWarningThreshold: 500000, QuotaWarningEnabled: &off, ChannelNotifyEnabled: &on, UpstreamModelUpdateNotifyEnabled: true, Language: "zhCN", SidebarModules: `{"chat":{"enabled":true}}`, BillingPreference: "wallet_only", } kept := func(setting dto.UserSetting) dto.UserSetting { setting.Language, setting.SidebarModules, setting.BillingPreference = stored.Language, stored.SidebarModules, stored.BillingPreference return setting } switchesChanged := `{"notify_type":"email","quota_warning_threshold":0,"quota_warning_enabled":true,"channel_notify_enabled":false,"upstream_model_update_notify_enabled":false}` for _, tc := range []struct { name string role int body string wantSuccess bool want dto.UserSetting }{ { name: "older client without the new switches", role: common.RoleAdminUser, body: `{"notify_type":"webhook","quota_warning_threshold":2000,"webhook_url":"https://example.com/hook","webhook_secret":"hook-secret","accept_unset_model_ratio_model":true,"record_ip_log":true}`, wantSuccess: true, want: kept(dto.UserSetting{ NotifyType: dto.NotifyTypeWebhook, WebhookUrl: "https://example.com/hook", WebhookSecret: "hook-secret", QuotaWarningThreshold: 2000, QuotaWarningEnabled: &off, ChannelNotifyEnabled: &on, UpstreamModelUpdateNotifyEnabled: true, AcceptUnsetRatioModel: true, RecordIpLog: true, }), }, { name: "administrator changes the switches and uses the site default threshold", role: common.RoleAdminUser, body: switchesChanged, wantSuccess: true, want: kept(dto.UserSetting{NotifyType: dto.NotifyTypeEmail, QuotaWarningEnabled: &on, ChannelNotifyEnabled: &off}), }, { name: "user cannot change the administrator switches", role: common.RoleCommonUser, body: switchesChanged, wantSuccess: true, want: kept(dto.UserSetting{NotifyType: dto.NotifyTypeEmail, QuotaWarningEnabled: &on, ChannelNotifyEnabled: &on, UpstreamModelUpdateNotifyEnabled: true}), }, { name: "negative threshold", role: common.RoleCommonUser, body: `{"notify_type":"email","quota_warning_threshold":-1}`, want: stored, }, } { t.Run(tc.name, func(t *testing.T) { db := setupManageUserTestDB(t) setting, err := common.Marshal(stored) require.NoError(t, err) user := model.User{Username: "setting-owner", Role: tc.role, Status: common.UserStatusEnabled, AffCode: "setting-owner-aff", Setting: string(setting)} require.NoError(t, db.Create(&user).Error) recorder := httptest.NewRecorder() c, _ := gin.CreateTestContext(recorder) c.Request = httptest.NewRequest(http.MethodPut, "/api/user/setting", strings.NewReader(tc.body)) c.Request.Header.Set("Content-Type", "application/json") c.Set("id", user.Id) UpdateUserSetting(c) assert.Contains(t, recorder.Body.String(), fmt.Sprintf(`"success":%t`, tc.wantSuccess)) require.NoError(t, db.First(&user, user.Id).Error) assert.Equal(t, tc.want, user.GetSetting()) }) } }