cneicy
ceddfd3d6b
Sync production source including routing, billing, Astra detection and public uptime. Update custom build instructions and license packaging.
999 lines
40 KiB
Go
999 lines
40 KiB
Go
package controller
|
|
|
|
import (
|
|
"bytes"
|
|
"database/sql"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"os"
|
|
"strconv"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/middleware"
|
|
"github.com/QuantumNous/new-api/model"
|
|
"github.com/QuantumNous/new-api/service"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/glebarez/sqlite"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"gorm.io/driver/mysql"
|
|
"gorm.io/driver/postgres"
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
type tokenAPIResponse struct {
|
|
Success bool `json:"success"`
|
|
Message string `json:"message"`
|
|
Data json.RawMessage `json:"data"`
|
|
}
|
|
|
|
type tokenPageResponse struct {
|
|
Items []tokenResponseItem `json:"items"`
|
|
}
|
|
|
|
type tokenResponseItem struct {
|
|
ID int `json:"id"`
|
|
Name string `json:"name"`
|
|
Key string `json:"key"`
|
|
Status int `json:"status"`
|
|
}
|
|
|
|
type tokenKeyResponse struct {
|
|
Key string `json:"key"`
|
|
}
|
|
|
|
type sqliteColumnInfo struct {
|
|
Name string `gorm:"column:name"`
|
|
Type string `gorm:"column:type"`
|
|
}
|
|
|
|
type legacyToken struct {
|
|
Id int `gorm:"primaryKey"`
|
|
UserId int `gorm:"index"`
|
|
Key string `gorm:"column:key;type:char(48);uniqueIndex"`
|
|
Status int `gorm:"default:1"`
|
|
Name string `gorm:"index"`
|
|
CreatedTime int64 `gorm:"bigint"`
|
|
AccessedTime int64 `gorm:"bigint"`
|
|
ExpiredTime int64 `gorm:"bigint;default:-1"`
|
|
RemainQuota int `gorm:"default:0"`
|
|
UnlimitedQuota bool
|
|
ModelLimitsEnabled bool
|
|
ModelLimits string `gorm:"type:text"`
|
|
AllowIps *string `gorm:"default:''"`
|
|
UsedQuota int `gorm:"default:0"`
|
|
Group string `gorm:"column:group;default:''"`
|
|
CrossGroupRetry bool
|
|
DeletedAt gorm.DeletedAt `gorm:"index"`
|
|
}
|
|
|
|
func (legacyToken) TableName() string {
|
|
return "tokens"
|
|
}
|
|
|
|
func openTokenControllerTestDB(t *testing.T) *gorm.DB {
|
|
t.Helper()
|
|
|
|
gin.SetMode(gin.TestMode)
|
|
common.SetDatabaseTypes(common.DatabaseTypeSQLite, common.DatabaseTypeSQLite)
|
|
common.RedisEnabled = false
|
|
|
|
dsn := fmt.Sprintf("file:%s?mode=memory&cache=shared", strings.ReplaceAll(t.Name(), "/", "_"))
|
|
db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{})
|
|
if err != nil {
|
|
t.Fatalf("failed to open sqlite db: %v", err)
|
|
}
|
|
model.DB = db
|
|
model.LOG_DB = db
|
|
|
|
t.Cleanup(func() {
|
|
sqlDB, err := db.DB()
|
|
if err == nil {
|
|
_ = sqlDB.Close()
|
|
}
|
|
})
|
|
|
|
return db
|
|
}
|
|
|
|
func migrateTokenControllerTestDB(t *testing.T, db *gorm.DB) {
|
|
t.Helper()
|
|
|
|
if err := db.AutoMigrate(&model.Token{}); err != nil {
|
|
t.Fatalf("failed to migrate token table: %v", err)
|
|
}
|
|
}
|
|
|
|
func setupTokenControllerTestDB(t *testing.T) *gorm.DB {
|
|
t.Helper()
|
|
|
|
db := openTokenControllerTestDB(t)
|
|
migrateTokenControllerTestDB(t, db)
|
|
return db
|
|
}
|
|
|
|
func TestTokenAstraMigration(t *testing.T) {
|
|
for _, dialect := range []string{"sqlite", "mysql", "postgres"} {
|
|
for _, upgrade := range []bool{false, true} {
|
|
t.Run(fmt.Sprintf("%s/upgrade=%t", dialect, upgrade), func(t *testing.T) {
|
|
var db *gorm.DB
|
|
if dialect == "sqlite" {
|
|
db = openTokenControllerTestDB(t)
|
|
} else {
|
|
dsn := os.Getenv("TEST_" + strings.ToUpper(dialect) + "_DSN")
|
|
if dsn == "" {
|
|
t.Skip("external test database is not configured")
|
|
}
|
|
var managed *bool
|
|
db, managed = openTokenControllerExternalDB(t, dialect, dsn)
|
|
*managed = true
|
|
if db.Migrator().HasTable("options") {
|
|
t.Fatal("requires an empty test database")
|
|
}
|
|
}
|
|
t.Cleanup(func() { require.NoError(t, db.Migrator().DropTable(&model.Option{})) })
|
|
if upgrade {
|
|
require.NoError(t, db.AutoMigrate(&legacyToken{}))
|
|
require.NoError(t, db.Create(&legacyToken{UserId: 7, Key: "legacy-astra", Name: "preserved", RemainQuota: 123}).Error)
|
|
}
|
|
require.NoError(t, db.AutoMigrate(&model.Token{}, &model.Option{}))
|
|
require.NoError(t, model.MigrateTokenAstraAutoSwitch())
|
|
require.NoError(t, model.MigrateTokenAstraDetection())
|
|
if upgrade {
|
|
var legacy model.Token
|
|
require.NoError(t, db.First(&legacy, 1).Error)
|
|
require.True(t, legacy.AstraAutoSwitch)
|
|
require.True(t, legacy.AstraDetectionEnabled)
|
|
require.Equal(t, 123, legacy.RemainQuota)
|
|
legacy.AstraAutoSwitch = false
|
|
legacy.AstraDetectionEnabled = false
|
|
require.NoError(t, legacy.Update())
|
|
}
|
|
fresh := model.Token{UserId: 7, Key: "new-astra"}
|
|
require.NoError(t, db.Create(&fresh).Error)
|
|
for range 2 {
|
|
require.NoError(t, db.AutoMigrate(&model.Token{}, &model.Option{}))
|
|
require.NoError(t, model.MigrateTokenAstraAutoSwitch())
|
|
require.NoError(t, model.MigrateTokenAstraDetection())
|
|
}
|
|
var tokens []model.Token
|
|
require.NoError(t, db.Find(&tokens).Error)
|
|
for _, token := range tokens {
|
|
require.False(t, token.AstraAutoSwitch, "restart must preserve opt-outs and new-token defaults")
|
|
require.False(t, token.AstraDetectionEnabled)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestTokenAstraSwitchAPIPreservesOmittedAndPersistsFalse(t *testing.T) {
|
|
db := setupTokenControllerTestDB(t)
|
|
token := seedToken(t, db, 7, "astra", "astra-api-key")
|
|
token.AstraAutoSwitch = true
|
|
token.AstraDetectionEnabled = true
|
|
require.NoError(t, token.Update())
|
|
for _, enabled := range []any{nil, false, true} {
|
|
body := map[string]any{"id": token.Id, "name": "astra", "unlimited_quota": true, "group": "default", "expired_time": -1}
|
|
if enabled != nil {
|
|
body["astra_auto_switch"] = enabled
|
|
}
|
|
c, writer := newAuthenticatedContext(t, http.MethodPut, "/api/token/", body, 7)
|
|
UpdateToken(c)
|
|
require.True(t, decodeAPIResponse(t, writer).Success, writer.Body.String())
|
|
require.NoError(t, db.First(token, token.Id).Error)
|
|
want := enabled == nil || enabled == true
|
|
require.Equal(t, want, token.AstraAutoSwitch)
|
|
var response struct {
|
|
AstraAutoSwitch bool `json:"astra_auto_switch"`
|
|
}
|
|
require.NoError(t, common.Unmarshal(decodeAPIResponse(t, writer).Data, &response))
|
|
require.Equal(t, want, response.AstraAutoSwitch)
|
|
}
|
|
c, writer := newAuthenticatedContext(t, http.MethodPut, "/api/token/", map[string]any{"id": token.Id, "name": "foreign", "astra_auto_switch": false}, 8)
|
|
UpdateToken(c)
|
|
require.False(t, decodeAPIResponse(t, writer).Success)
|
|
require.NoError(t, db.First(token, token.Id).Error)
|
|
require.True(t, token.AstraAutoSwitch)
|
|
}
|
|
|
|
func openTokenControllerExternalDB(t *testing.T, dialect string, dsn string) (*gorm.DB, *bool) {
|
|
t.Helper()
|
|
|
|
gin.SetMode(gin.TestMode)
|
|
common.RedisEnabled = false
|
|
|
|
var (
|
|
db *gorm.DB
|
|
dbType common.DatabaseType
|
|
err error
|
|
)
|
|
switch dialect {
|
|
case "mysql":
|
|
dbType = common.DatabaseTypeMySQL
|
|
db, err = gorm.Open(mysql.Open(dsn), &gorm.Config{})
|
|
case "postgres":
|
|
dbType = common.DatabaseTypePostgreSQL
|
|
db, err = gorm.Open(postgres.Open(dsn), &gorm.Config{})
|
|
default:
|
|
t.Fatalf("unsupported dialect %q", dialect)
|
|
}
|
|
common.SetDatabaseTypes(dbType, dbType)
|
|
if err != nil {
|
|
t.Fatalf("failed to open %s db: %v", dialect, err)
|
|
}
|
|
|
|
model.DB = db
|
|
model.LOG_DB = db
|
|
|
|
if db.Migrator().HasTable("tokens") {
|
|
t.Skipf("refusing to run %s migration compatibility test against external database because tokens table already exists", dialect)
|
|
}
|
|
|
|
managedTokensTable := new(bool)
|
|
|
|
t.Cleanup(func() {
|
|
if *managedTokensTable && db.Migrator().HasTable("tokens") {
|
|
_ = db.Migrator().DropTable("tokens")
|
|
}
|
|
sqlDB, err := db.DB()
|
|
if err == nil {
|
|
_ = sqlDB.Close()
|
|
}
|
|
})
|
|
|
|
return db, managedTokensTable
|
|
}
|
|
|
|
func seedToken(t *testing.T, db *gorm.DB, userID int, name string, rawKey string) *model.Token {
|
|
t.Helper()
|
|
|
|
token := &model.Token{
|
|
UserId: userID,
|
|
Name: name,
|
|
Key: rawKey,
|
|
Status: common.TokenStatusEnabled,
|
|
CreatedTime: 1,
|
|
AccessedTime: 1,
|
|
ExpiredTime: -1,
|
|
RemainQuota: 100,
|
|
UnlimitedQuota: true,
|
|
Group: "default",
|
|
}
|
|
if err := db.Create(token).Error; err != nil {
|
|
t.Fatalf("failed to create token: %v", err)
|
|
}
|
|
return token
|
|
}
|
|
|
|
func newAuthenticatedContext(t *testing.T, method string, target string, body any, userID int) (*gin.Context, *httptest.ResponseRecorder) {
|
|
t.Helper()
|
|
|
|
var requestBody *bytes.Reader
|
|
if body != nil {
|
|
payload, err := common.Marshal(body)
|
|
if err != nil {
|
|
t.Fatalf("failed to marshal request body: %v", err)
|
|
}
|
|
requestBody = bytes.NewReader(payload)
|
|
} else {
|
|
requestBody = bytes.NewReader(nil)
|
|
}
|
|
|
|
recorder := httptest.NewRecorder()
|
|
ctx, _ := gin.CreateTestContext(recorder)
|
|
ctx.Request = httptest.NewRequest(method, target, requestBody)
|
|
if body != nil {
|
|
ctx.Request.Header.Set("Content-Type", "application/json")
|
|
}
|
|
ctx.Set("id", userID)
|
|
return ctx, recorder
|
|
}
|
|
|
|
func decodeAPIResponse(t *testing.T, recorder *httptest.ResponseRecorder) tokenAPIResponse {
|
|
t.Helper()
|
|
|
|
var response tokenAPIResponse
|
|
if err := common.Unmarshal(recorder.Body.Bytes(), &response); err != nil {
|
|
t.Fatalf("failed to decode api response: %v", err)
|
|
}
|
|
return response
|
|
}
|
|
|
|
func getSQLiteColumnType(t *testing.T, db *gorm.DB, tableName string, columnName string) string {
|
|
t.Helper()
|
|
|
|
var columns []sqliteColumnInfo
|
|
if err := db.Raw("PRAGMA table_info(" + tableName + ")").Scan(&columns).Error; err != nil {
|
|
t.Fatalf("failed to inspect %s schema: %v", tableName, err)
|
|
}
|
|
|
|
for _, column := range columns {
|
|
if column.Name == columnName {
|
|
return strings.ToLower(column.Type)
|
|
}
|
|
}
|
|
|
|
t.Fatalf("column %s not found in %s schema", columnName, tableName)
|
|
return ""
|
|
}
|
|
|
|
func getTokenKeyColumnType(t *testing.T, db *gorm.DB, dialect string) string {
|
|
t.Helper()
|
|
|
|
switch dialect {
|
|
case "sqlite":
|
|
return getSQLiteColumnType(t, db, "tokens", "key")
|
|
case "mysql":
|
|
var columnType string
|
|
if err := db.Raw(`SELECT COLUMN_TYPE FROM information_schema.columns
|
|
WHERE table_schema = DATABASE() AND table_name = ? AND column_name = ?`,
|
|
"tokens", "key").Scan(&columnType).Error; err != nil {
|
|
t.Fatalf("failed to inspect mysql token key column: %v", err)
|
|
}
|
|
return strings.ToLower(columnType)
|
|
case "postgres":
|
|
var dataType string
|
|
var maxLength sql.NullInt64
|
|
if err := db.Raw(`SELECT data_type, character_maximum_length
|
|
FROM information_schema.columns
|
|
WHERE table_schema = current_schema() AND table_name = ? AND column_name = ?`,
|
|
"tokens", "key").Row().Scan(&dataType, &maxLength); err != nil {
|
|
t.Fatalf("failed to inspect postgres token key column: %v", err)
|
|
}
|
|
switch strings.ToLower(dataType) {
|
|
case "character varying":
|
|
return fmt.Sprintf("varchar(%d)", maxLength.Int64)
|
|
case "character":
|
|
return fmt.Sprintf("char(%d)", maxLength.Int64)
|
|
default:
|
|
if maxLength.Valid {
|
|
return fmt.Sprintf("%s(%d)", strings.ToLower(dataType), maxLength.Int64)
|
|
}
|
|
return strings.ToLower(dataType)
|
|
}
|
|
default:
|
|
t.Fatalf("unsupported dialect %q", dialect)
|
|
return ""
|
|
}
|
|
}
|
|
|
|
func getTokenAutoGroupsColumnType(t *testing.T, db *gorm.DB, dialect string) string {
|
|
t.Helper()
|
|
|
|
switch dialect {
|
|
case "sqlite":
|
|
return getSQLiteColumnType(t, db, "tokens", "auto_groups")
|
|
case "mysql":
|
|
var columnType string
|
|
if err := db.Raw(`SELECT DATA_TYPE FROM information_schema.columns
|
|
WHERE table_schema = DATABASE() AND table_name = ? AND column_name = ?`,
|
|
"tokens", "auto_groups").Scan(&columnType).Error; err != nil {
|
|
t.Fatalf("failed to inspect mysql token auto_groups column: %v", err)
|
|
}
|
|
return strings.ToLower(columnType)
|
|
case "postgres":
|
|
var dataType string
|
|
if err := db.Raw(`SELECT data_type FROM information_schema.columns
|
|
WHERE table_schema = current_schema() AND table_name = ? AND column_name = ?`,
|
|
"tokens", "auto_groups").Scan(&dataType).Error; err != nil {
|
|
t.Fatalf("failed to inspect postgres token auto_groups column: %v", err)
|
|
}
|
|
return strings.ToLower(dataType)
|
|
default:
|
|
t.Fatalf("unsupported dialect %q", dialect)
|
|
return ""
|
|
}
|
|
}
|
|
|
|
func runTokenMigrationCompatibilityTest(t *testing.T, db *gorm.DB, dialect string, managedTokensTable *bool) {
|
|
t.Helper()
|
|
|
|
legacyKey := strings.Repeat("a", 48)
|
|
longKey := strings.Repeat("b", 64)
|
|
|
|
if err := db.AutoMigrate(&legacyToken{}); err != nil {
|
|
t.Fatalf("failed to create legacy token schema: %v", err)
|
|
}
|
|
if managedTokensTable != nil {
|
|
*managedTokensTable = true
|
|
}
|
|
if err := db.Create(&legacyToken{
|
|
UserId: 7,
|
|
Key: legacyKey,
|
|
Status: common.TokenStatusEnabled,
|
|
Name: "legacy-token",
|
|
CreatedTime: 1,
|
|
AccessedTime: 1,
|
|
ExpiredTime: -1,
|
|
RemainQuota: 100,
|
|
UnlimitedQuota: true,
|
|
ModelLimitsEnabled: false,
|
|
ModelLimits: "",
|
|
AllowIps: common.GetPointer(""),
|
|
UsedQuota: 0,
|
|
Group: "default",
|
|
CrossGroupRetry: false,
|
|
}).Error; err != nil {
|
|
t.Fatalf("failed to seed legacy token row: %v", err)
|
|
}
|
|
|
|
if got := getTokenKeyColumnType(t, db, dialect); got != "char(48)" {
|
|
t.Fatalf("expected legacy key column type char(48), got %q", got)
|
|
}
|
|
|
|
migrateTokenControllerTestDB(t, db)
|
|
|
|
if got := getTokenKeyColumnType(t, db, dialect); got != "varchar(128)" {
|
|
t.Fatalf("expected migrated key column type varchar(128), got %q", got)
|
|
}
|
|
if !db.Migrator().HasColumn(&model.Token{}, "auto_groups") {
|
|
t.Fatal("expected migration to add auto_groups column")
|
|
}
|
|
if got := getTokenAutoGroupsColumnType(t, db, dialect); got != "text" {
|
|
t.Fatalf("expected migrated auto_groups column type text, got %q", got)
|
|
}
|
|
|
|
var migratedToken model.Token
|
|
if err := db.First(&migratedToken, "name = ?", "legacy-token").Error; err != nil {
|
|
t.Fatalf("failed to load migrated token row: %v", err)
|
|
}
|
|
if migratedToken.Key != legacyKey {
|
|
t.Fatalf("expected migrated token key %q, got %q", legacyKey, migratedToken.Key)
|
|
}
|
|
if migratedToken.Name != "legacy-token" {
|
|
t.Fatalf("expected migrated token name to be preserved, got %q", migratedToken.Name)
|
|
}
|
|
if migratedToken.AutoGroups != "" {
|
|
t.Fatalf("expected legacy token to inherit global Auto groups, got %q", migratedToken.AutoGroups)
|
|
}
|
|
|
|
inserted := model.Token{
|
|
UserId: 8,
|
|
Name: "long-token",
|
|
Key: longKey,
|
|
Status: common.TokenStatusEnabled,
|
|
CreatedTime: 1,
|
|
AccessedTime: 1,
|
|
ExpiredTime: -1,
|
|
RemainQuota: 200,
|
|
UnlimitedQuota: true,
|
|
ModelLimitsEnabled: false,
|
|
ModelLimits: "",
|
|
AllowIps: common.GetPointer(""),
|
|
UsedQuota: 0,
|
|
Group: "default",
|
|
CrossGroupRetry: false,
|
|
}
|
|
if err := db.Create(&inserted).Error; err != nil {
|
|
t.Fatalf("failed to insert long token after migration: %v", err)
|
|
}
|
|
|
|
var fetched model.Token
|
|
if err := db.First(&fetched, "id = ?", inserted.Id).Error; err != nil {
|
|
t.Fatalf("failed to fetch long token after migration: %v", err)
|
|
}
|
|
if fetched.Key != longKey {
|
|
t.Fatalf("expected long token key %q, got %q", longKey, fetched.Key)
|
|
}
|
|
}
|
|
|
|
func TestTokenAutoMigrateUsesVarchar128KeyColumn(t *testing.T) {
|
|
db := setupTokenControllerTestDB(t)
|
|
|
|
if got := getTokenKeyColumnType(t, db, "sqlite"); got != "varchar(128)" {
|
|
t.Fatalf("expected key column type varchar(128), got %q", got)
|
|
}
|
|
if got := getSQLiteColumnType(t, db, "tokens", "auto_groups"); got != "text" {
|
|
t.Fatalf("expected auto_groups column type text, got %q", got)
|
|
}
|
|
}
|
|
|
|
func TestTokenMigrationFromChar48ToVarchar128(t *testing.T) {
|
|
db := openTokenControllerTestDB(t)
|
|
runTokenMigrationCompatibilityTest(t, db, "sqlite", nil)
|
|
}
|
|
|
|
func TestTokenMigrationFromChar48ToVarchar128MySQL(t *testing.T) {
|
|
dsn := os.Getenv("TEST_MYSQL_DSN")
|
|
if dsn == "" {
|
|
t.Skip("set TEST_MYSQL_DSN to run mysql migration compatibility test")
|
|
}
|
|
|
|
db, managedTokensTable := openTokenControllerExternalDB(t, "mysql", dsn)
|
|
runTokenMigrationCompatibilityTest(t, db, "mysql", managedTokensTable)
|
|
}
|
|
|
|
func TestTokenMigrationFromChar48ToVarchar128Postgres(t *testing.T) {
|
|
dsn := os.Getenv("TEST_POSTGRES_DSN")
|
|
if dsn == "" {
|
|
t.Skip("set TEST_POSTGRES_DSN to run postgres migration compatibility test")
|
|
}
|
|
|
|
db, managedTokensTable := openTokenControllerExternalDB(t, "postgres", dsn)
|
|
runTokenMigrationCompatibilityTest(t, db, "postgres", managedTokensTable)
|
|
}
|
|
|
|
func TestGetAllTokensMasksKeyInResponse(t *testing.T) {
|
|
db := setupTokenControllerTestDB(t)
|
|
token := seedToken(t, db, 1, "list-token", "abcd1234efgh5678")
|
|
seedToken(t, db, 2, "other-user-token", "zzzz1234yyyy5678")
|
|
|
|
ctx, recorder := newAuthenticatedContext(t, http.MethodGet, "/api/token/?p=1&size=10", nil, 1)
|
|
GetAllTokens(ctx)
|
|
|
|
response := decodeAPIResponse(t, recorder)
|
|
if !response.Success {
|
|
t.Fatalf("expected success response, got message: %s", response.Message)
|
|
}
|
|
|
|
var page tokenPageResponse
|
|
if err := common.Unmarshal(response.Data, &page); err != nil {
|
|
t.Fatalf("failed to decode token page response: %v", err)
|
|
}
|
|
if len(page.Items) != 1 {
|
|
t.Fatalf("expected exactly one token, got %d", len(page.Items))
|
|
}
|
|
if page.Items[0].Key != token.GetMaskedKey() {
|
|
t.Fatalf("expected masked key %q, got %q", token.GetMaskedKey(), page.Items[0].Key)
|
|
}
|
|
if strings.Contains(recorder.Body.String(), token.Key) {
|
|
t.Fatalf("list response leaked raw token key: %s", recorder.Body.String())
|
|
}
|
|
}
|
|
|
|
func TestSearchTokensMasksKeyInResponse(t *testing.T) {
|
|
db := setupTokenControllerTestDB(t)
|
|
token := seedToken(t, db, 1, "searchable-token", "ijkl1234mnop5678")
|
|
|
|
ctx, recorder := newAuthenticatedContext(t, http.MethodGet, "/api/token/search?keyword=searchable-token&p=1&size=10", nil, 1)
|
|
SearchTokens(ctx)
|
|
|
|
response := decodeAPIResponse(t, recorder)
|
|
if !response.Success {
|
|
t.Fatalf("expected success response, got message: %s", response.Message)
|
|
}
|
|
|
|
var page tokenPageResponse
|
|
if err := common.Unmarshal(response.Data, &page); err != nil {
|
|
t.Fatalf("failed to decode search response: %v", err)
|
|
}
|
|
if len(page.Items) != 1 {
|
|
t.Fatalf("expected exactly one search result, got %d", len(page.Items))
|
|
}
|
|
if page.Items[0].Key != token.GetMaskedKey() {
|
|
t.Fatalf("expected masked search key %q, got %q", token.GetMaskedKey(), page.Items[0].Key)
|
|
}
|
|
if strings.Contains(recorder.Body.String(), token.Key) {
|
|
t.Fatalf("search response leaked raw token key: %s", recorder.Body.String())
|
|
}
|
|
}
|
|
|
|
func TestGetTokenMasksKeyInResponse(t *testing.T) {
|
|
db := setupTokenControllerTestDB(t)
|
|
token := seedToken(t, db, 1, "detail-token", "qrst1234uvwx5678")
|
|
|
|
ctx, recorder := newAuthenticatedContext(t, http.MethodGet, "/api/token/"+strconv.Itoa(token.Id), nil, 1)
|
|
ctx.Params = gin.Params{{Key: "id", Value: strconv.Itoa(token.Id)}}
|
|
GetToken(ctx)
|
|
|
|
response := decodeAPIResponse(t, recorder)
|
|
if !response.Success {
|
|
t.Fatalf("expected success response, got message: %s", response.Message)
|
|
}
|
|
|
|
var detail tokenResponseItem
|
|
if err := common.Unmarshal(response.Data, &detail); err != nil {
|
|
t.Fatalf("failed to decode token detail response: %v", err)
|
|
}
|
|
if detail.Key != token.GetMaskedKey() {
|
|
t.Fatalf("expected masked detail key %q, got %q", token.GetMaskedKey(), detail.Key)
|
|
}
|
|
if strings.Contains(recorder.Body.String(), token.Key) {
|
|
t.Fatalf("detail response leaked raw token key: %s", recorder.Body.String())
|
|
}
|
|
}
|
|
|
|
func TestUpdateTokenMasksKeyInResponse(t *testing.T) {
|
|
db := setupTokenControllerTestDB(t)
|
|
token := seedToken(t, db, 1, "editable-token", "yzab1234cdef5678")
|
|
|
|
body := map[string]any{
|
|
"id": token.Id,
|
|
"name": "updated-token",
|
|
"expired_time": -1,
|
|
"remain_quota": 100,
|
|
"unlimited_quota": true,
|
|
"model_limits_enabled": false,
|
|
"model_limits": "",
|
|
"group": "default",
|
|
"cross_group_retry": false,
|
|
}
|
|
|
|
ctx, recorder := newAuthenticatedContext(t, http.MethodPut, "/api/token/", body, 1)
|
|
UpdateToken(ctx)
|
|
|
|
response := decodeAPIResponse(t, recorder)
|
|
if !response.Success {
|
|
t.Fatalf("expected success response, got message: %s", response.Message)
|
|
}
|
|
|
|
var detail tokenResponseItem
|
|
if err := common.Unmarshal(response.Data, &detail); err != nil {
|
|
t.Fatalf("failed to decode token update response: %v", err)
|
|
}
|
|
if detail.Key != token.GetMaskedKey() {
|
|
t.Fatalf("expected masked update key %q, got %q", token.GetMaskedKey(), detail.Key)
|
|
}
|
|
if strings.Contains(recorder.Body.String(), token.Key) {
|
|
t.Fatalf("update response leaked raw token key: %s", recorder.Body.String())
|
|
}
|
|
}
|
|
|
|
func TestGetTokenKeyRequiresOwnershipAndReturnsFullKey(t *testing.T) {
|
|
db := setupTokenControllerTestDB(t)
|
|
token := seedToken(t, db, 1, "owned-token", "owner1234token5678")
|
|
|
|
authorizedCtx, authorizedRecorder := newAuthenticatedContext(t, http.MethodPost, "/api/token/"+strconv.Itoa(token.Id)+"/key", nil, 1)
|
|
authorizedCtx.Params = gin.Params{{Key: "id", Value: strconv.Itoa(token.Id)}}
|
|
GetTokenKey(authorizedCtx)
|
|
|
|
authorizedResponse := decodeAPIResponse(t, authorizedRecorder)
|
|
if !authorizedResponse.Success {
|
|
t.Fatalf("expected authorized key fetch to succeed, got message: %s", authorizedResponse.Message)
|
|
}
|
|
|
|
var keyData tokenKeyResponse
|
|
if err := common.Unmarshal(authorizedResponse.Data, &keyData); err != nil {
|
|
t.Fatalf("failed to decode token key response: %v", err)
|
|
}
|
|
if keyData.Key != token.GetFullKey() {
|
|
t.Fatalf("expected full key %q, got %q", token.GetFullKey(), keyData.Key)
|
|
}
|
|
|
|
unauthorizedCtx, unauthorizedRecorder := newAuthenticatedContext(t, http.MethodPost, "/api/token/"+strconv.Itoa(token.Id)+"/key", nil, 2)
|
|
unauthorizedCtx.Params = gin.Params{{Key: "id", Value: strconv.Itoa(token.Id)}}
|
|
GetTokenKey(unauthorizedCtx)
|
|
|
|
unauthorizedResponse := decodeAPIResponse(t, unauthorizedRecorder)
|
|
if unauthorizedResponse.Success {
|
|
t.Fatalf("expected unauthorized key fetch to fail")
|
|
}
|
|
if strings.Contains(unauthorizedRecorder.Body.String(), token.Key) {
|
|
t.Fatalf("unauthorized key response leaked raw token key: %s", unauthorizedRecorder.Body.String())
|
|
}
|
|
}
|
|
|
|
func TestAPITokenAuditDatabaseMatrix(t *testing.T) {
|
|
for _, database := range []struct {
|
|
name, env string
|
|
typ common.DatabaseType
|
|
}{
|
|
{"sqlite", "", common.DatabaseTypeSQLite},
|
|
{"mysql", "AUDIT_MYSQL_DSN", common.DatabaseTypeMySQL},
|
|
{"postgres", "AUDIT_POSTGRES_DSN", common.DatabaseTypePostgreSQL},
|
|
} {
|
|
for _, separateLog := range []bool{false, true} {
|
|
t.Run(fmt.Sprintf("%s/separate_log=%v", database.name, separateLog), func(t *testing.T) {
|
|
dsn := os.Getenv(database.env)
|
|
if database.env != "" && dsn == "" {
|
|
t.Skip(database.env + " is not configured")
|
|
}
|
|
previousDB, previousLogDB := model.DB, model.LOG_DB
|
|
previousMain, previousLog := common.MainDatabaseType(), common.LogDatabaseType()
|
|
previousRedis, previousMaster, previousSecret := common.RedisEnabled, common.IsMasterNode, common.SessionSecret
|
|
t.Cleanup(func() {
|
|
model.DB, model.LOG_DB = previousDB, previousLogDB
|
|
common.SetDatabaseTypes(previousMain, previousLog)
|
|
common.RedisEnabled, common.IsMasterNode, common.SessionSecret = previousRedis, previousMaster, previousSecret
|
|
})
|
|
common.RedisEnabled, common.IsMasterNode = false, true
|
|
common.SessionSecret = "api-token-audit-test-secret"
|
|
t.Setenv("LOG_SQL_DSN", "")
|
|
db, _ := newAuditTestDatabase(t, database.name, dsn)
|
|
model.DB = db
|
|
common.SetDatabaseTypes(database.typ, database.typ)
|
|
require.NoError(t, db.AutoMigrate(&model.User{}, &model.UserSession{}, &model.Token{}))
|
|
// Initialize production column quoting as well as the existing audit table.
|
|
require.NoError(t, model.InitLogDB())
|
|
if separateLog {
|
|
logDB, _ := newAuditTestDatabase(t, database.name, dsn)
|
|
model.LOG_DB = logDB
|
|
require.NoError(t, model.MigrateAuditLogs())
|
|
}
|
|
versionSQL := "SELECT version()"
|
|
if database.name == "sqlite" {
|
|
versionSQL = "SELECT sqlite_version()"
|
|
}
|
|
var version string
|
|
require.NoError(t, db.Raw(versionSQL).Scan(&version).Error)
|
|
t.Logf("database version: %s", version)
|
|
verifyAPITokenAudit(t)
|
|
if separateLog {
|
|
var count int64
|
|
require.NoError(t, db.Model(&model.AuditLog{}).Count(&count).Error)
|
|
assert.Zero(t, count, "all audit events must use the configured log database")
|
|
}
|
|
})
|
|
}
|
|
}
|
|
}
|
|
|
|
func verifyAPITokenAudit(t *testing.T) {
|
|
t.Helper()
|
|
pat := "api-token-audit-pat-secret"
|
|
user := &model.User{Username: "token-owner", Password: "placeholder", Role: common.RoleCommonUser, Status: common.UserStatusEnabled, Group: "default", AuthVersion: 1, AccessToken: &pat, AffCode: "token-owner"}
|
|
require.NoError(t, model.DB.Create(user).Error)
|
|
other := &model.User{Username: "other-owner", Password: "placeholder", Role: common.RoleCommonUser, Status: common.UserStatusEnabled, Group: "default", AuthVersion: 1, AffCode: "other-owner"}
|
|
require.NoError(t, model.DB.Create(other).Error)
|
|
session := &model.UserSession{SID: "token-audit-session", UserID: user.Id, Version: 1, UserAuthVersion: 1, Status: model.UserSessionStatusActive, RefreshHash: "placeholder", LoginMethod: "password", LastActiveAt: time.Now().Unix(), ExpiresAt: time.Now().Add(time.Hour).Unix()}
|
|
require.NoError(t, model.CreateUserSession(session))
|
|
jwt, _, err := service.IssueAccessToken(service.AuthIdentity{UserID: user.Id, SessionID: session.SID, UserAuthVersion: 1, SessionVersion: 1})
|
|
require.NoError(t, err)
|
|
router := gin.New()
|
|
router.Use(middleware.RequestId(), middleware.AccessTokenAudit())
|
|
tokenRoutes := router.Group("/api/token", middleware.UserAuth(), middleware.TokenOperationAudit())
|
|
tokenRoutes.POST("/", AddToken)
|
|
tokenRoutes.PUT("/", UpdateToken)
|
|
tokenRoutes.DELETE("/:id", DeleteToken)
|
|
tokenRoutes.POST("/batch", DeleteTokenBatch)
|
|
tokenRoutes.POST("/batch/keys", GetTokenKeysBatch)
|
|
tokenRoutes.POST("/:id/key", func(c *gin.Context) {
|
|
if c.GetHeader("X-Test-Limit") != "" {
|
|
c.AbortWithStatusJSON(429, gin.H{"success": false})
|
|
return
|
|
}
|
|
c.Next()
|
|
}, GetTokenKey)
|
|
tokenRoutes.GET("/", GetAllTokens)
|
|
tokenRoutes.GET("/:id", GetToken)
|
|
tokenRoutes.GET("/search", SearchTokens)
|
|
router.GET("/api/audit/self", middleware.UserAuth(), GetAuditLogs)
|
|
|
|
for index, tc := range []struct {
|
|
name, method, path, body, action string
|
|
success bool
|
|
params string
|
|
initialStatus int
|
|
failWrite, rateLimit, usePAT bool
|
|
}{
|
|
{name: "create", method: "POST", path: "/", body: `{"name":"created","expired_time":-1,"unlimited_quota":true}`, action: "token.create", success: true},
|
|
{name: "invalid create", method: "POST", path: "/", body: `{"name":"attempt","remain_quota":-1}`, action: "token.create", params: `{"name":"attempt"}`},
|
|
{name: "malformed body", method: "POST", path: "/", body: `{"key":"raw-body-secret"`, action: "token.create", params: `{}`},
|
|
{name: "validation error exceeds audit buffer", method: "POST", path: "/", body: `{"remain_quota":` + strings.Repeat("9", 64*1024) + `}`, action: "token.create", params: `{}`},
|
|
{name: "create storage failure", method: "POST", path: "/", body: `{"name":"attempt","unlimited_quota":true}`, action: "token.create", params: `{"name":"attempt"}`, failWrite: true},
|
|
{name: "normalized update", method: "PUT", path: "/", body: `{"id":$id,"name":"renamed","expired_time":-1,"remain_quota":200,"unlimited_quota":true,"group":"default","cross_group_retry":true,"allow_ips":""}`, action: "token.update", success: true, params: `{"id":$id,"name":"renamed","changed_fields":["name","remain_quota","group","cross_group_retry","auto_groups"]}`},
|
|
{name: "configuration values stay private", method: "PUT", path: "/", body: `{"id":$id,"name":"owned","expired_time":42,"remain_quota":100,"unlimited_quota":false,"model_limits_enabled":true,"model_limits":"private-model-configuration","allow_ips":"203.0.113.57","group":"auto","cross_group_retry":true}`, action: "token.update", success: true, params: `{"id":$id,"name":"owned","changed_fields":["expired_time","unlimited_quota","model_limits_enabled","model_limits","allow_ips"]}`},
|
|
{name: "unchanged update", method: "PUT", path: "/", body: `{"id":$id,"name":"owned","expired_time":-1,"remain_quota":100,"unlimited_quota":true,"group":"auto","cross_group_retry":true,"allow_ips":""}`, action: "token.update", success: true, params: `{"id":$id,"name":"owned","changed_fields":[]}`},
|
|
{name: "successful response exceeds audit buffer", method: "PUT", path: "/", body: `{"id":$id,"name":"owned","expired_time":-1,"remain_quota":100,"unlimited_quota":true,"group":"auto","cross_group_retry":true,"allow_ips":"","model_limits":"` + strings.Repeat("m", 64*1024-128) + `"}`, action: "token.update", success: true, params: `{"id":$id,"name":"owned","changed_fields":["model_limits"]}`},
|
|
{name: "update storage failure", method: "PUT", path: "/", body: `{"id":$id,"name":"failed-rename","unlimited_quota":true}`, action: "token.update", params: `{"id":$id,"name":"owned"}`, failWrite: true},
|
|
{name: "foreign update", method: "PUT", path: "/", body: `{"id":$other,"name":"forged-name","unlimited_quota":true}`, action: "token.update", params: `{"id":$other}`},
|
|
{name: "disable", method: "PUT", path: "/?status_only=true", body: `{"id":$id,"status":2}`, action: "token.status_update", success: true, params: `{"id":$id,"name":"owned","from":1,"to":2}`},
|
|
{name: "enable", method: "PUT", path: "/?status_only=true", body: `{"id":$id,"status":1}`, action: "token.status_update", success: true, params: `{"id":$id,"name":"owned","from":2,"to":1}`, initialStatus: common.TokenStatusDisabled},
|
|
{name: "expired enable", method: "PUT", path: "/?status_only=true", body: `{"id":$id,"status":1}`, action: "token.status_update", params: `{"id":$id,"name":"owned"}`, initialStatus: common.TokenStatusExpired},
|
|
{name: "delete", method: "DELETE", path: "/$id", action: "token.delete", success: true, params: `{"id":$id,"name":"owned"}`},
|
|
{name: "foreign delete", method: "DELETE", path: "/$other", action: "token.delete", params: `{"id":$other}`},
|
|
{name: "missing delete", method: "DELETE", path: "/999999", action: "token.delete", params: `{"id":999999}`},
|
|
{name: "key view", method: "POST", path: "/$id/key", action: "token.key_view", success: true, params: `{"id":$id,"name":"owned"}`},
|
|
{name: "PAT key view", method: "POST", path: "/$id/key", action: "token.key_view", success: true, params: `{"id":$id,"name":"owned"}`, usePAT: true},
|
|
{name: "foreign key view", method: "POST", path: "/$other/key", action: "token.key_view", params: `{"id":$other}`},
|
|
{name: "rate limited key view", method: "POST", path: "/$id/key", action: "token.key_view", params: `{"id":$id}`, rateLimit: true},
|
|
{name: "batch delete partial and duplicate", method: "POST", path: "/batch", body: `{"ids":[$id,$id,$other,999999]}`, action: "token.delete_batch", success: true, params: `{"requested_ids":[$id,$id,$other,999999],"total":4,"count":1}`},
|
|
{name: "empty batch delete", method: "POST", path: "/batch", body: `{"ids":[]}`, action: "token.delete_batch", params: `{"requested_ids":[],"total":0}`},
|
|
{name: "batch keys partial and duplicate", method: "POST", path: "/batch/keys", body: `{"ids":[$id,$id,$other,999999]}`, action: "token.key_view_batch", success: true, params: `{"requested_ids":[$id,$id,$other,999999],"total":4,"count":1,"returned_ids":[$id]}`},
|
|
{name: "batch keys no matches", method: "POST", path: "/batch/keys", body: `{"ids":[$other,999999]}`, action: "token.key_view_batch", success: true, params: `{"requested_ids":[$other,999999],"total":2,"count":0,"returned_ids":[]}`},
|
|
{name: "empty batch keys", method: "POST", path: "/batch/keys", body: `{"ids":[]}`, action: "token.key_view_batch", params: `{"requested_ids":[],"total":0}`},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
empty := ""
|
|
owned := &model.Token{UserId: user.Id, Name: "owned", Key: fmt.Sprintf("owned-key-secret-%d", index), Status: common.TokenStatusEnabled, ExpiredTime: -1, RemainQuota: 100, UnlimitedQuota: true, Group: "auto", CrossGroupRetry: true, AutoGroups: `["default"]`, AllowIps: &empty}
|
|
if tc.initialStatus != 0 {
|
|
owned.Status = tc.initialStatus
|
|
}
|
|
if owned.Status == common.TokenStatusExpired {
|
|
owned.ExpiredTime = 1
|
|
}
|
|
foreign := &model.Token{UserId: other.Id, Name: "private-foreign-name", Key: fmt.Sprintf("foreign-key-secret-%d", index)}
|
|
require.NoError(t, model.DB.Create(owned).Error)
|
|
require.NoError(t, model.DB.Create(foreign).Error)
|
|
replace := strings.NewReplacer("$id", strconv.Itoa(owned.Id), "$other", strconv.Itoa(foreign.Id))
|
|
if tc.failWrite {
|
|
fail := func(tx *gorm.DB) {
|
|
if tx.Statement.Table == "tokens" {
|
|
_ = tx.AddError(errors.New("raw-storage-error-secret"))
|
|
}
|
|
}
|
|
if tc.method == "POST" {
|
|
require.NoError(t, model.DB.Callback().Create().Before("gorm:create").Register("token-audit:fail", fail))
|
|
t.Cleanup(func() { require.NoError(t, model.DB.Callback().Create().Remove("token-audit:fail")) })
|
|
} else {
|
|
require.NoError(t, model.DB.Callback().Update().Before("gorm:update").Register("token-audit:fail", fail))
|
|
t.Cleanup(func() { require.NoError(t, model.DB.Callback().Update().Remove("token-audit:fail")) })
|
|
}
|
|
}
|
|
request := httptest.NewRequest(tc.method, "/api/token"+replace.Replace(tc.path), strings.NewReader(replace.Replace(tc.body)))
|
|
credential := jwt
|
|
if tc.usePAT {
|
|
credential = pat
|
|
}
|
|
request.Header.Set("Authorization", "Bearer "+credential)
|
|
request.Header.Set("Content-Type", "application/json")
|
|
request.Header.Set("User-Agent", "api-token-audit-client")
|
|
request.RemoteAddr = "192.0.2.12:4321"
|
|
if tc.rateLimit {
|
|
request.Header.Set("X-Test-Limit", "1")
|
|
}
|
|
response := httptest.NewRecorder()
|
|
router.ServeHTTP(response, request)
|
|
if strings.Contains(tc.name, "exceeds audit buffer") {
|
|
assert.Greater(t, response.Body.Len(), 64*1024)
|
|
}
|
|
var events []model.AuditLog
|
|
require.NoError(t, model.LOG_DB.Where("request_id = ?", response.Header().Get(common.RequestIdKey)).Find(&events).Error)
|
|
expectedEvents := 1
|
|
if tc.usePAT {
|
|
expectedEvents = 2
|
|
}
|
|
require.Len(t, events, expectedEvents)
|
|
var operation *model.AuditLog
|
|
for i := range events {
|
|
if events[i].Category == model.AuditCategorySecurity {
|
|
require.Nil(t, operation, "one operation event per request")
|
|
operation = &events[i]
|
|
} else {
|
|
assert.Equal(t, model.AuditCategoryAccessToken, events[i].Category)
|
|
assert.Equal(t, model.AccessTokenFingerprint(pat), events[i].TokenRef)
|
|
}
|
|
}
|
|
require.NotNil(t, operation)
|
|
assert.Equal(t, tc.action, operation.Action)
|
|
assert.Equal(t, tc.success, operation.Success)
|
|
assert.Equal(t, response.Code, operation.Status)
|
|
if tc.rateLimit {
|
|
assert.Equal(t, 429, response.Code)
|
|
} else {
|
|
assert.Equal(t, 200, response.Code)
|
|
assert.Equal(t, tc.success, decodeAPIResponse(t, response).Success)
|
|
}
|
|
assert.Equal(t, user.Id, operation.UserId)
|
|
assert.Equal(t, user.Username, operation.Username)
|
|
assert.Equal(t, common.RoleCommonUser, operation.ActorRole)
|
|
assert.Equal(t, "192.0.2.12", operation.Ip)
|
|
assert.Equal(t, "api-token-audit-client", operation.UserAgent)
|
|
assert.Equal(t, tc.method, operation.Method)
|
|
expectedRoute := strings.NewReplacer("$id", ":id", "$other", ":id", "999999", ":id").Replace(strings.Split(tc.path, "?")[0])
|
|
assert.Equal(t, "/api/token"+expectedRoute, operation.Route)
|
|
assert.NotEmpty(t, operation.RequestId)
|
|
assert.Empty(t, operation.TokenRef)
|
|
assert.Nil(t, operation.Other.AdminInfo)
|
|
authMethod := "session"
|
|
if tc.usePAT {
|
|
authMethod = "access_token"
|
|
}
|
|
assert.Equal(t, authMethod, operation.AuthMethod)
|
|
require.NotNil(t, operation.Other.Op)
|
|
assert.Equal(t, tc.action, operation.Other.Op.Action)
|
|
params, err := common.Marshal(operation.Other.Op.Params)
|
|
require.NoError(t, err)
|
|
if tc.name == "create" {
|
|
var created model.Token
|
|
require.NoError(t, model.DB.Where("user_id = ? AND name = ?", user.Id, "created").First(&created).Error)
|
|
assert.JSONEq(t, fmt.Sprintf(`{"id":%d,"name":"created"}`, created.Id), string(params))
|
|
assert.NotContains(t, string(params), created.Key)
|
|
} else if tc.params == `{}` {
|
|
assert.Empty(t, operation.Other.Op.Params)
|
|
} else {
|
|
assert.JSONEq(t, replace.Replace(tc.params), string(params))
|
|
}
|
|
encoded, err := common.Marshal(events)
|
|
require.NoError(t, err)
|
|
for _, secret := range []string{pat, jwt, owned.Key, foreign.Key, foreign.Name, "raw-body-secret", "raw-storage-error-secret", "private-model-configuration", "203.0.113.57", "Authorization"} {
|
|
assert.NotContains(t, string(encoded), secret)
|
|
}
|
|
if tc.action == "token.key_view" && tc.success {
|
|
assert.Contains(t, response.Body.String(), owned.GetFullKey())
|
|
}
|
|
})
|
|
}
|
|
|
|
t.Run("bounded batch metadata", func(t *testing.T) {
|
|
ids := make([]int, 101)
|
|
for i := range ids {
|
|
ids[i] = 10000 + i
|
|
}
|
|
body, err := common.Marshal(TokenBatch{Ids: ids})
|
|
require.NoError(t, err)
|
|
for _, path := range []string{"/api/token/batch", "/api/token/batch/keys"} {
|
|
request := httptest.NewRequest("POST", path, bytes.NewReader(body))
|
|
request.Header.Set("Authorization", "Bearer "+jwt)
|
|
request.Header.Set("Content-Type", "application/json")
|
|
response := httptest.NewRecorder()
|
|
router.ServeHTTP(response, request)
|
|
var event model.AuditLog
|
|
require.NoError(t, model.LOG_DB.Where("request_id = ?", response.Header().Get(common.RequestIdKey)).First(&event).Error)
|
|
var params struct {
|
|
IDs []int `json:"requested_ids"`
|
|
Truncated bool `json:"requested_ids_truncated"`
|
|
Total int `json:"total"`
|
|
Count *int `json:"count"`
|
|
}
|
|
encoded, err := common.Marshal(event.Other.Op.Params)
|
|
require.NoError(t, err)
|
|
require.NoError(t, common.Unmarshal(encoded, ¶ms))
|
|
assert.Equal(t, ids[:100], params.IDs)
|
|
assert.True(t, params.Truncated)
|
|
assert.Equal(t, 101, params.Total)
|
|
assert.Equal(t, path == "/api/token/batch", event.Success)
|
|
if event.Success {
|
|
require.NotNil(t, params.Count)
|
|
assert.Zero(t, *params.Count)
|
|
} else {
|
|
assert.Nil(t, params.Count)
|
|
}
|
|
}
|
|
})
|
|
|
|
t.Run("reads and unauthenticated writes add no operation audit", func(t *testing.T) {
|
|
for _, tc := range []struct{ method, path, credential string }{
|
|
{"GET", "/api/token/", jwt}, {"GET", "/api/token/search?keyword=private-search", jwt},
|
|
{"GET", "/api/token/999999", jwt}, {"POST", "/api/token/", ""},
|
|
} {
|
|
response := auditRequest(router, tc.method, tc.path, tc.credential)
|
|
var count int64
|
|
require.NoError(t, model.LOG_DB.Model(&model.AuditLog{}).Where("request_id = ?", response.Header().Get(common.RequestIdKey)).Count(&count).Error)
|
|
assert.Zero(t, count)
|
|
}
|
|
})
|
|
|
|
t.Run("self audit excludes other owners", func(t *testing.T) {
|
|
model.RecordAuditLog(nil, model.AuditLog{UserId: other.Id, Username: other.Username, ActorRole: common.RoleCommonUser, Category: model.AuditCategorySecurity, Action: "token.delete", Content: "other-user-audit", Success: true})
|
|
response := auditRequest(router, "GET", "/api/audit/self?category=security&page_size=100", jwt)
|
|
assert.Equal(t, 200, response.Code)
|
|
assert.Contains(t, response.Body.String(), "token.create")
|
|
assert.NotContains(t, response.Body.String(), "other-user-audit")
|
|
})
|
|
|
|
t.Run("audit storage failure preserves operation result", func(t *testing.T) {
|
|
require.NoError(t, model.LOG_DB.Callback().Create().Before("gorm:create").Register("token-audit:log-fail", func(tx *gorm.DB) {
|
|
if tx.Statement.Table == "audit_logs" {
|
|
_ = tx.AddError(errors.New("audit unavailable"))
|
|
}
|
|
}))
|
|
t.Cleanup(func() { require.NoError(t, model.LOG_DB.Callback().Create().Remove("token-audit:log-fail")) })
|
|
request := httptest.NewRequest("POST", "/api/token/", strings.NewReader(`{"name":"audit-down","unlimited_quota":true}`))
|
|
request.Header.Set("Authorization", "Bearer "+jwt)
|
|
request.Header.Set("Content-Type", "application/json")
|
|
response := httptest.NewRecorder()
|
|
router.ServeHTTP(response, request)
|
|
require.True(t, decodeAPIResponse(t, response).Success)
|
|
var count int64
|
|
require.NoError(t, model.DB.Model(&model.Token{}).Where("name = ?", "audit-down").Count(&count).Error)
|
|
assert.EqualValues(t, 1, count)
|
|
})
|
|
}
|
|
|
|
func TestTokenDetectionOffDisablesSwitchAndSurvivesOmittedUpdate(t *testing.T) {
|
|
db := setupTokenControllerTestDB(t)
|
|
token := seedToken(t, db, 7, "detection", "detection-api-key")
|
|
token.AstraDetectionEnabled, token.AstraAutoSwitch = true, true
|
|
require.NoError(t, token.Update())
|
|
for _, detection := range []any{false, nil, true} {
|
|
body := map[string]any{"id": token.Id, "name": "detection", "unlimited_quota": true, "group": "default", "expired_time": -1}
|
|
if detection != nil {
|
|
body["astra_detection_enabled"] = detection
|
|
}
|
|
c, writer := newAuthenticatedContext(t, http.MethodPut, "/api/token/", body, 7)
|
|
UpdateToken(c)
|
|
require.True(t, decodeAPIResponse(t, writer).Success, writer.Body.String())
|
|
require.NoError(t, db.First(token, token.Id).Error)
|
|
require.Equal(t, detection == true, token.AstraDetectionEnabled)
|
|
require.False(t, token.AstraAutoSwitch)
|
|
}
|
|
}
|