model/auth_cache_integrity_test.go
Rama main · réplica de lectura del repositorio original.
package model
import (
"errors"
"os"
"sync/atomic"
"testing"
"time"
"github.com/QuantumNous/new-api/common"
"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"
)
func TestAuthCacheIntegrityDatabaseMatrix(t *testing.T) {
for _, dialect := range []string{"sqlite", "mysql", "postgres"} {
t.Run(dialect, func(t *testing.T) {
var driver gorm.Dialector
switch dialect {
case "sqlite":
driver = sqlite.Open(":memory:")
case "mysql":
dsn := os.Getenv("TEST_MYSQL_DSN")
if dsn == "" {
t.Skip("TEST_MYSQL_DSN is not configured")
}
driver = mysql.Open(dsn)
case "postgres":
dsn := os.Getenv("TEST_POSTGRES_DSN")
if dsn == "" {
t.Skip("TEST_POSTGRES_DSN is not configured")
}
driver = postgres.Open(dsn)
}
db, err := gorm.Open(driver, &gorm.Config{})
require.NoError(t, err)
sqlDB, err := db.DB()
require.NoError(t, err)
sqlDB.SetMaxOpenConns(1)
previousDB, previousType := DB, common.MainDatabaseType()
DB = db
common.SetMainDatabaseType(common.DatabaseType(dialect))
initCol()
t.Cleanup(func() {
DB = previousDB
common.SetMainDatabaseType(previousType)
initCol()
assert.NoError(t, sqlDB.Close())
})
require.NoError(t, db.AutoMigrate(&User{}, &Token{}, &UserSession{}))
var version string
query := "SELECT version()"
if dialect == "sqlite" {
query = "SELECT sqlite_version()"
}
require.NoError(t, db.Raw(query).Scan(&version).Error)
t.Logf("database version: %s", version)
useUserCacheMiniRedis(t)
user := User{Username: "cache-integrity-user", AffCode: "cache-int", Role: common.RoleCommonUser,
Status: common.UserStatusEnabled, AuthVersion: 1, Group: "default", Quota: 100}
require.NoError(t, DB.Create(&user).Error)
t.Cleanup(func() { assert.NoError(t, db.Unscoped().Delete(&user).Error) })
require.NoError(t, populateUserCache(user))
quotaResult, err := cacheApplyUserQuotaDelta(user.Id, -70)
require.NoError(t, err)
require.Equal(t, cacheQuotaOK, quotaResult)
t.Run("user_metadata_and_partial_refresh", func(t *testing.T) {
for _, fields := range []map[string]any{
{"Role": common.RoleRootUser},
{"AuthVersion": 99, "Status": common.UserStatusDisabled},
{"Group": "admin", "Setting": `{"language":"evil"}`},
} {
require.NoError(t, common.RDB.HSet(t.Context(), getUserCacheKey(user.Id), fields).Err())
_, err := cacheGetUserBase(user.Id)
require.Error(t, err)
// The version fence intentionally rejects snapshots below a cached version.
if _, ok := fields["AuthVersion"]; ok {
require.NoError(t, common.RDB.HSet(t.Context(), getUserCacheKey(user.Id), "AuthVersion", user.AuthVersion).Err())
}
loaded, err := GetUserCache(user.Id)
require.NoError(t, err)
assert.Equal(t, user.Role, loaded.Role)
assert.Equal(t, user.Group, loaded.Group)
assert.Equal(t, user.Status, loaded.Status)
cached, err := cacheGetUserBase(user.Id)
require.NoError(t, err)
assert.Equal(t, 30, cached.Quota)
}
require.NoError(t, common.RDB.HSet(t.Context(), getUserCacheKey(user.Id), "Role", common.RoleRootUser).Err())
require.NoError(t, DB.Model(&User{}).Where("id = ?", user.Id).Update("group", "paid").Error)
require.NoError(t, RefreshUserGroupCache(user.Id))
cached, err := cacheGetUserBase(user.Id)
require.NoError(t, err)
assert.Equal(t, common.RoleCommonUser, cached.Role)
assert.Equal(t, "paid", cached.Group)
assert.Equal(t, 30, cached.Quota)
})
allowIPs := "192.168.0.0/16"
token := Token{UserId: user.Id, Key: "cache-integrity-token", Status: common.TokenStatusEnabled,
ExpiredTime: -1, RemainQuota: 100, ModelLimitsEnabled: true, ModelLimits: "allowed-model",
AllowIps: &allowIPs, Group: "auto", AutoGroups: `["default"]`}
require.NoError(t, DB.Create(&token).Error)
t.Cleanup(func() { assert.NoError(t, db.Unscoped().Delete(&token).Error) })
_, err = GetTokenByKey(token.Key, false)
require.NoError(t, err)
quotaResult, err = cacheTryReserveTokenQuota(token.Id, token.Key, 70)
require.NoError(t, err)
require.Equal(t, cacheQuotaOK, quotaResult)
t.Run("token_identity_and_permissions", func(t *testing.T) {
for _, fields := range []map[string]any{
{"UserId": user.Id + 1},
{"UnlimitedQuota": "true", "AllowUserQuotaQuery": "true"},
{"ModelLimitsEnabled": "false", "ModelLimits": "", "AllowIps": ""},
{"Group": "privileged", "AutoGroups": `["privileged"]`, "CrossGroupRetry": "true"},
{"Status": common.TokenStatusDisabled, "ExpiredTime": 1},
{"UnlimitedQuota": "garbage"},
{"DeletedAt": time.Now().Format(time.RFC3339)},
} {
require.NoError(t, common.RDB.HSet(t.Context(), getTokenCacheKey(token.Key), fields).Err())
_, err := cacheGetTokenByKey(token.Key)
require.Error(t, err)
_, err = GetTokenByKey(token.Key, false)
require.NoError(t, err)
cached, err := cacheGetTokenByKey(token.Key)
require.NoError(t, err)
assert.Equal(t, token.UserId, cached.UserId)
assert.False(t, cached.UnlimitedQuota)
assert.False(t, cached.AllowUserQuotaQuery)
assert.True(t, cached.ModelLimitsEnabled)
assert.Equal(t, token.ModelLimits, cached.ModelLimits)
assert.Equal(t, token.AllowIps, cached.AllowIps)
assert.Equal(t, token.Group, cached.Group)
assert.Equal(t, token.AutoGroups, cached.AutoGroups)
assert.False(t, cached.CrossGroupRetry)
assert.Equal(t, token.Status, cached.Status)
assert.Equal(t, token.ExpiredTime, cached.ExpiredTime)
assert.Equal(t, 30, cached.RemainQuota)
assert.Equal(t, 70, cached.UsedQuota)
assert.Positive(t, cached.AccessedTime)
}
})
session := newTestUserSession("cache-integrity-session", user.Id, time.Now().Unix())
require.NoError(t, CreateUserSession(session))
t.Cleanup(func() { assert.NoError(t, db.Delete(session).Error) })
t.Run("session_owner_and_versions", func(t *testing.T) {
require.NoError(t, common.RDB.HSet(t.Context(), userSessionCacheKey(session.SID),
"UserID", user.Id+1, "Version", 2, "UserAuthVersion", 2, "ExpiresAt", session.ExpiresAt+3600).Err())
_, err := getUserSessionCache(session.SID)
require.Error(t, err)
// A higher cached version is a deny fence for delayed active snapshots.
require.NoError(t, common.RDB.HSet(t.Context(), userSessionCacheKey(session.SID), "Version", 1).Err())
loaded, err := GetUserSessionCached(session.SID)
require.NoError(t, err)
assert.Equal(t, session.UserID, loaded.UserID)
assert.Equal(t, session.UserAuthVersion, loaded.UserAuthVersion)
assert.Equal(t, session.ExpiresAt, loaded.ExpiresAt)
})
t.Run("unsigned_and_expired_cache", func(t *testing.T) {
for _, key := range []string{getUserCacheKey(user.Id), getTokenCacheKey(token.Key), userSessionCacheKey(session.SID)} {
require.NoError(t, common.RDB.HDel(t.Context(), key, "CacheMAC", "CacheExpiresAt").Err())
}
_, err := GetUserCache(user.Id)
require.NoError(t, err)
_, err = GetTokenByKey(token.Key, false)
require.NoError(t, err)
_, err = GetUserSessionCached(session.SID)
require.NoError(t, err)
cachedUser, err := cacheGetUserBase(user.Id)
require.NoError(t, err)
cachedToken, err := cacheGetTokenByKey(token.Key)
require.NoError(t, err)
cachedSession, err := getUserSessionCache(session.SID)
require.NoError(t, err)
expired := time.Now().Add(-time.Second).UnixMilli()
cachedUser.CacheExpiresAt = expired
cachedToken.CacheExpiresAt = expired
cachedSession.CacheExpiresAt = expired
userMAC, err := cachedUser.cacheMAC()
require.NoError(t, err)
tokenMAC, err := cachedToken.cacheMAC()
require.NoError(t, err)
sessionMAC, err := authCacheMAC(userSessionCacheKey(session.SID), cachedSession)
require.NoError(t, err)
for key, mac := range map[string]string{getUserCacheKey(user.Id): userMAC, getTokenCacheKey(token.Key): tokenMAC, userSessionCacheKey(session.SID): sessionMAC} {
require.NoError(t, common.RDB.HSet(t.Context(), key, "CacheExpiresAt", expired, "CacheMAC", mac).Err())
require.NoError(t, common.RDB.Persist(t.Context(), key).Err())
}
_, err = cacheGetUserBase(user.Id)
assert.Error(t, err)
_, err = cacheGetTokenByKey(token.Key)
assert.Error(t, err)
_, err = getUserSessionCache(session.SID)
assert.Error(t, err)
_, err = GetUserCache(user.Id)
require.NoError(t, err)
_, err = GetTokenByKey(token.Key, false)
require.NoError(t, err)
_, err = GetUserSessionCached(session.SID)
require.NoError(t, err)
})
t.Run("repair_during_quota_reservation", func(t *testing.T) {
require.NoError(t, common.RDB.HSet(t.Context(), getTokenCacheKey(token.Key), "CacheMAC", "invalid").Err())
snapshotRead := make(chan struct{})
releaseRepair := make(chan struct{}, 1)
defer close(releaseRepair)
var intercepted atomic.Bool
require.NoError(t, DB.Callback().Query().After("gorm:query").Register("auth_cache_quota_race", func(tx *gorm.DB) {
if tx.Statement.Table == "tokens" && intercepted.CompareAndSwap(false, true) {
close(snapshotRead)
<-releaseRepair
}
}))
t.Cleanup(func() { assert.NoError(t, DB.Callback().Query().Remove("auth_cache_quota_race")) })
done := make(chan error, 1)
go func() {
_, err := GetTokenByKey(token.Key, false)
done <- err
}()
<-snapshotRead
result, err := cacheTryReserveTokenQuota(token.Id, token.Key, 10)
require.NoError(t, err)
require.Equal(t, cacheQuotaOK, result)
releaseRepair <- struct{}{}
require.NoError(t, <-done)
cached, err := cacheGetTokenByKey(token.Key)
require.NoError(t, err)
assert.Equal(t, 20, cached.RemainQuota)
assert.Equal(t, 80, cached.UsedQuota)
})
t.Run("cross_identity_copy_and_database_failure", func(t *testing.T) {
for source, destination := range map[string]string{
getUserCacheKey(user.Id): getUserCacheKey(user.Id + 1000),
getTokenCacheKey(token.Key): getTokenCacheKey("unknown-token"),
userSessionCacheKey(session.SID): userSessionCacheKey("unknown-session"),
} {
fields, err := common.RDB.HGetAll(t.Context(), source).Result()
require.NoError(t, err)
require.NoError(t, common.RDB.HSet(t.Context(), destination, fields).Err())
}
require.NoError(t, common.RDB.HSet(t.Context(), getUserCacheKey(user.Id+1000), "Id", user.Id+1000).Err())
require.NoError(t, common.RDB.HSet(t.Context(), userSessionCacheKey("unknown-session"), "SID", "unknown-session").Err())
_, err := GetUserCache(user.Id + 1000)
assert.ErrorIs(t, err, gorm.ErrRecordNotFound)
_, err = GetTokenByKey("unknown-token", false)
assert.ErrorIs(t, err, gorm.ErrRecordNotFound)
_, err = GetUserSessionCached("unknown-session")
assert.ErrorIs(t, err, gorm.ErrRecordNotFound)
require.NoError(t, DB.Callback().Query().Before("gorm:query").Register("auth_cache_db_failure", func(tx *gorm.DB) {
tx.AddError(errors.New("database unavailable"))
}))
defer func() { require.NoError(t, DB.Callback().Query().Remove("auth_cache_db_failure")) }()
_, err = GetTokenByKey(token.Key, false)
require.NoError(t, err, "valid token cache must not query the database")
_, err = GetUserSessionCached(session.SID)
require.NoError(t, err, "valid session cache must not query the database")
require.NoError(t, common.RDB.HSet(t.Context(), getTokenCacheKey(token.Key), "CacheMAC", "invalid").Err())
require.NoError(t, common.RDB.HSet(t.Context(), userSessionCacheKey(session.SID), "CacheMAC", "invalid").Err())
_, err = GetTokenByKey(token.Key, false)
assert.Error(t, err)
_, err = GetUserSessionCached(session.SID)
assert.Error(t, err)
})
t.Run("session_revocation_cannot_be_cleared", func(t *testing.T) {
revoked, err := RevokeUserSession(user.Id, session.SID, "test")
require.NoError(t, err)
require.True(t, revoked)
require.NoError(t, common.RDB.HSet(t.Context(), userSessionCacheKey(session.SID),
"Status", UserSessionStatusActive, "RevokedAt", 0, "RevokedReason", "").Err())
_, err = GetUserSessionCached(session.SID)
assert.ErrorIs(t, err, ErrUserSessionInactive)
})
t.Run("unreadable_quota_counters", func(t *testing.T) {
require.NoError(t, common.RDB.HSet(t.Context(), getUserCacheKey(user.Id), "Quota", "garbage").Err())
require.NoError(t, common.RDB.HSet(t.Context(), getTokenCacheKey(token.Key), "RemainQuota", "garbage").Err())
_, err := GetUserCache(user.Id)
require.NoError(t, err)
_, err = GetTokenByKey(token.Key, false)
require.NoError(t, err)
cachedUser, err := cacheGetUserBase(user.Id)
require.NoError(t, err)
assert.Equal(t, user.Quota, cachedUser.Quota)
cachedToken, err := cacheGetTokenByKey(token.Key)
require.NoError(t, err)
assert.Equal(t, token.RemainQuota, cachedToken.RemainQuota)
assert.Equal(t, 80, cachedToken.UsedQuota)
})
})
}
}