controller/user_manage_test.go
Rama main · réplica de lectura del repositorio original.
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())
})
}
}