mirror of
https://github.com/QuantumNous/new-api.git
synced 2026-10-28 16:46:57 +08:00
Calcium-Ion
6461f084ff
SMTP settings now only decide whether mail can be sent. Verification codes, password reset, email binding and account security mails depend on SMTP alone, and no notification switch affects them. Admins get a site-wide low balance reminder switch, on by default, next to the SMTP settings. The default threshold is entered as an amount. Users choose which notifications they receive. Low balance reminders are on for everyone by default. Admins can also turn channel notifications and upstream model update notices on or off. The root user receives channel notifications until it turns them off. Other admins receive them after they turn them on. A low balance reminder is sent once, when a charge takes the balance below the threshold. A top-up above the threshold arms it again. Channel notifications reach every admin who has them on, in that admin's language. Email binding codes and account security mails follow the account's language. Saving notification settings no longer clears the language, the sidebar layout or the billing preference. Existing settings keep their behavior, and old clients that do not send the new fields keep the stored switches. Webhook, Bark and Gotify payloads are unchanged.
727 lines
32 KiB
Go
727 lines
32 KiB
Go
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())
|
|
})
|
|
}
|
|
}
|