feat: publish sanitized custom.97 source
This commit is contained in:
1022 files changed
+69903
-19207
No files matched your search
+239
-14
@@ -2,7 +2,9 @@ package model
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"strconv"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
@@ -144,19 +146,22 @@ func TestUsageAccountingSupportsSignedDirectAndBatchDeltas(t *testing.T) {
|
||||
assert.Equal(t, int64(1150), gotChannel.UsedQuota)
|
||||
}
|
||||
|
||||
func TestUpdateUserAccessTokenOnlyUpdatesAccessToken(t *testing.T) {
|
||||
func TestRevokeUserAccessTokenOnlyClearsLegacyToken(t *testing.T) {
|
||||
setupUserUpdateTestState(t)
|
||||
|
||||
createdAt := int64(100)
|
||||
user := User{
|
||||
Id: 2,
|
||||
Username: "token-rotation-user",
|
||||
Password: "password",
|
||||
DisplayName: "before",
|
||||
Status: common.UserStatusEnabled,
|
||||
Quota: 1000,
|
||||
AffQuota: 800,
|
||||
AffHistoryQuota: 1200,
|
||||
Id: 2,
|
||||
Username: "token-revoke-user",
|
||||
Password: "password",
|
||||
DisplayName: "before",
|
||||
Status: common.UserStatusEnabled,
|
||||
Quota: 1000,
|
||||
AffQuota: 800,
|
||||
AffHistoryQuota: 1200,
|
||||
AccessTokenCreatedAt: &createdAt,
|
||||
}
|
||||
user.SetAccessToken("legacy-token")
|
||||
require.NoError(t, DB.Create(&user).Error)
|
||||
|
||||
require.NoError(t, DB.Model(&User{}).Where("id = ?", user.Id).Updates(map[string]any{
|
||||
@@ -165,23 +170,26 @@ func TestUpdateUserAccessTokenOnlyUpdatesAccessToken(t *testing.T) {
|
||||
"display_name": "concurrent-update",
|
||||
}).Error)
|
||||
|
||||
require.NoError(t, UpdateUserAccessToken(user.Id, "rotated-token"))
|
||||
tokenRef, err := RevokeUserAccessToken(user.Id)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, AccessTokenFingerprint("legacy-token"), tokenRef)
|
||||
|
||||
var got User
|
||||
require.NoError(t, DB.First(&got, user.Id).Error)
|
||||
assert.Equal(t, "rotated-token", got.GetAccessToken())
|
||||
assert.Empty(t, got.GetAccessToken())
|
||||
assert.Nil(t, got.AccessTokenCreatedAt)
|
||||
assert.Equal(t, "concurrent-update", got.DisplayName)
|
||||
assert.Equal(t, 1500, got.Quota)
|
||||
assert.Equal(t, 300, got.AffQuota)
|
||||
assert.Equal(t, 1200, got.AffHistoryQuota)
|
||||
}
|
||||
|
||||
func TestUpdateUserAccessTokenRejectsSoftDeletedUser(t *testing.T) {
|
||||
func TestRevokeUserAccessTokenRejectsSoftDeletedUser(t *testing.T) {
|
||||
setupUserUpdateTestState(t)
|
||||
|
||||
user := User{
|
||||
Id: 3,
|
||||
Username: "deleted-token-rotation-user",
|
||||
Username: "deleted-token-revoke-user",
|
||||
Password: "password",
|
||||
Status: common.UserStatusEnabled,
|
||||
}
|
||||
@@ -189,7 +197,7 @@ func TestUpdateUserAccessTokenRejectsSoftDeletedUser(t *testing.T) {
|
||||
require.NoError(t, DB.Create(&user).Error)
|
||||
require.NoError(t, DB.Delete(&user).Error)
|
||||
|
||||
err := UpdateUserAccessToken(user.Id, "orphaned-token")
|
||||
_, err := RevokeUserAccessToken(user.Id)
|
||||
require.ErrorIs(t, err, gorm.ErrRecordNotFound)
|
||||
|
||||
var got User
|
||||
@@ -398,3 +406,220 @@ func TestResetUserPasswordByEmailRequiresSingleActiveMatch(t *testing.T) {
|
||||
err = ResetUserPasswordByEmail("missing@example.com", "NewPassword123")
|
||||
require.True(t, errors.Is(err, ErrEmailNotFound))
|
||||
}
|
||||
|
||||
func createAccessTokenTestUser(t *testing.T, username string) User {
|
||||
t.Helper()
|
||||
user := User{
|
||||
Username: username,
|
||||
Password: "unused-password-hash",
|
||||
Role: common.RoleCommonUser,
|
||||
Status: common.UserStatusEnabled,
|
||||
Group: "default",
|
||||
AuthVersion: 1,
|
||||
AffCode: username + "-aff",
|
||||
}
|
||||
require.NoError(t, DB.Create(&user).Error)
|
||||
return user
|
||||
}
|
||||
|
||||
func createTestUserAccessToken(t *testing.T, userID int, raw string, expiresAt int64) UserAccessToken {
|
||||
t.Helper()
|
||||
token := UserAccessToken{Name: raw, TokenHash: AccessTokenFingerprint(raw), TokenHint: AccessTokenHint(raw), ExpiresAt: expiresAt}
|
||||
require.NoError(t, token.SetScopes([]string{"profile:read"}))
|
||||
require.NoError(t, CreateUserAccessToken(userID, &token, 20))
|
||||
return token
|
||||
}
|
||||
|
||||
func TestLegacyAccessTokenRetireAtIsWrittenOnceAndReadOnly(t *testing.T) {
|
||||
require.NoError(t, DB.AutoMigrate(&Option{}))
|
||||
previous := legacyAccessTokenRetireAt.Load()
|
||||
originalMap := common.OptionMap
|
||||
common.OptionMap = map[string]string{}
|
||||
require.NoError(t, DB.Where(&Option{Key: legacyAccessTokenRetireAtKey}).Delete(&Option{}).Error)
|
||||
t.Cleanup(func() {
|
||||
legacyAccessTokenRetireAt.Store(previous)
|
||||
common.OptionMap = originalMap
|
||||
DB.Where(&Option{Key: legacyAccessTokenRetireAtKey}).Delete(&Option{})
|
||||
})
|
||||
|
||||
const firstStart = int64(1_800_000_000)
|
||||
want := firstStart + 30*24*60*60
|
||||
require.NoError(t, EnsureLegacyAccessTokenRetireAt(firstStart))
|
||||
assert.Equal(t, want, LegacyAccessTokenRetireAt())
|
||||
|
||||
require.NoError(t, EnsureLegacyAccessTokenRetireAt(firstStart+24*60*60))
|
||||
assert.Equal(t, want, LegacyAccessTokenRetireAt())
|
||||
var stored Option
|
||||
require.NoError(t, DB.Where(&Option{Key: legacyAccessTokenRetireAtKey}).First(&stored).Error)
|
||||
assert.Equal(t, strconv.FormatInt(want, 10), stored.Value)
|
||||
|
||||
require.ErrorIs(t, UpdateOption(legacyAccessTokenRetireAtKey, "1"), errLegacyRetireAtReadOnly)
|
||||
require.ErrorIs(t, UpdateOptionsBulk(map[string]string{legacyAccessTokenRetireAtKey: "1"}), errLegacyRetireAtReadOnly)
|
||||
require.NoError(t, DB.Where(&Option{Key: legacyAccessTokenRetireAtKey}).First(&stored).Error)
|
||||
assert.Equal(t, strconv.FormatInt(want, 10), stored.Value)
|
||||
|
||||
require.NoError(t, updateOptionMap(legacyAccessTokenRetireAtKey, stored.Value))
|
||||
assert.NotContains(t, common.OptionMap, legacyAccessTokenRetireAtKey)
|
||||
}
|
||||
|
||||
func TestValidateAccessTokenStopsAtLegacyDeadline(t *testing.T) {
|
||||
setupUserUpdateTestState(t)
|
||||
previous := legacyAccessTokenRetireAt.Load()
|
||||
t.Cleanup(func() { legacyAccessTokenRetireAt.Store(previous) })
|
||||
|
||||
user := createAccessTokenTestUser(t, "legacy-deadline-user")
|
||||
require.NoError(t, DB.Model(&User{}).Where("id = ?", user.Id).Update("access_token", "legacy-deadline-token").Error)
|
||||
|
||||
legacyAccessTokenRetireAt.Store(time.Now().Add(time.Hour).Unix())
|
||||
found, err := ValidateAccessToken("legacy-deadline-token")
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, found)
|
||||
assert.Equal(t, user.Id, found.Id)
|
||||
|
||||
var userQueries int
|
||||
const callbackName = "test:count_legacy_token_queries"
|
||||
require.NoError(t, DB.Callback().Query().Before("gorm:query").Register(callbackName, func(tx *gorm.DB) {
|
||||
if tx.Statement.Table == "users" {
|
||||
userQueries++
|
||||
}
|
||||
}))
|
||||
t.Cleanup(func() { _ = DB.Callback().Query().Remove(callbackName) })
|
||||
|
||||
for _, retireAt := range []int64{time.Now().Add(-time.Second).Unix(), 0} {
|
||||
legacyAccessTokenRetireAt.Store(retireAt)
|
||||
found, err = ValidateAccessToken("legacy-deadline-token")
|
||||
require.ErrorIs(t, err, ErrLegacyAccessTokenRetired)
|
||||
assert.Nil(t, found)
|
||||
}
|
||||
assert.Zero(t, userQueries)
|
||||
}
|
||||
|
||||
func TestUserAccessTokenLimitAndOwnership(t *testing.T) {
|
||||
setupUserUpdateTestState(t)
|
||||
owner := createAccessTokenTestUser(t, "token-owner")
|
||||
other := createAccessTokenTestUser(t, "token-other")
|
||||
|
||||
var first UserAccessToken
|
||||
for i := range 20 {
|
||||
token := createTestUserAccessToken(t, owner.Id, "nap_limit_"+strconv.Itoa(i), 0)
|
||||
if i == 0 {
|
||||
first = token
|
||||
}
|
||||
}
|
||||
extra := UserAccessToken{TokenHash: AccessTokenFingerprint("nap_limit_extra")}
|
||||
require.ErrorIs(t, CreateUserAccessToken(owner.Id, &extra, 20), ErrAccessTokenLimit)
|
||||
var count int64
|
||||
require.NoError(t, DB.Model(&UserAccessToken{}).Where("user_id = ?", owner.Id).Count(&count).Error)
|
||||
assert.EqualValues(t, 20, count)
|
||||
|
||||
_, err := UpdateUserAccessToken(other.Id, first.Id, "stolen", []string{"profile:read"})
|
||||
require.ErrorIs(t, err, ErrAccessTokenNotFound)
|
||||
_, err = DeleteUserAccessToken(other.Id, first.Id)
|
||||
require.ErrorIs(t, err, ErrAccessTokenNotFound)
|
||||
stored, err := GetUserAccessToken(owner.Id, first.Id)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, first.Name, stored.Name)
|
||||
|
||||
assert.Equal(t, first.Scopes, stored.Scopes)
|
||||
|
||||
renamed, err := UpdateUserAccessToken(owner.Id, first.Id, "renamed", nil)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "renamed", renamed.Name)
|
||||
assert.Equal(t, first.Scopes, renamed.Scopes, "a rename keeps the grant")
|
||||
regranted, err := UpdateUserAccessToken(owner.Id, first.Id, "regranted", []string{"usage:read"})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "regranted", regranted.Name)
|
||||
assert.Equal(t, []string{"usage:read"}, regranted.GetScopes())
|
||||
deleted, err := DeleteUserAccessToken(owner.Id, first.Id)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, first.TokenHash, deleted.TokenHash)
|
||||
found, err := FindUserAccessTokenByHash(first.TokenHash)
|
||||
require.NoError(t, err)
|
||||
assert.Nil(t, found)
|
||||
}
|
||||
|
||||
func TestDisablingUserKeepsAccessTokensAndDeletingRemovesThem(t *testing.T) {
|
||||
setupUserUpdateTestState(t)
|
||||
user := createAccessTokenTestUser(t, "token-lifecycle-user")
|
||||
createTestUserAccessToken(t, user.Id, "nap_lifecycle_one", 0)
|
||||
createTestUserAccessToken(t, user.Id, "nap_lifecycle_two", time.Now().Add(time.Hour).Unix())
|
||||
countTokens := func() int64 {
|
||||
var count int64
|
||||
require.NoError(t, DB.Model(&UserAccessToken{}).Where("user_id = ?", user.Id).Count(&count).Error)
|
||||
return count
|
||||
}
|
||||
|
||||
user.Status = common.UserStatusDisabled
|
||||
require.NoError(t, user.Update(false))
|
||||
assert.EqualValues(t, 2, countTokens())
|
||||
|
||||
revoked, err := user.Delete()
|
||||
require.NoError(t, err)
|
||||
assert.EqualValues(t, 2, revoked)
|
||||
assert.Zero(t, countTokens())
|
||||
|
||||
hardDeleted := createAccessTokenTestUser(t, "token-hard-delete-user")
|
||||
createTestUserAccessToken(t, hardDeleted.Id, "nap_hard_delete", 0)
|
||||
revoked, err = hardDeleted.HardDelete()
|
||||
require.NoError(t, err)
|
||||
assert.EqualValues(t, 1, revoked)
|
||||
var remaining int64
|
||||
require.NoError(t, DB.Model(&UserAccessToken{}).Where("user_id = ?", hardDeleted.Id).Count(&remaining).Error)
|
||||
assert.Zero(t, remaining)
|
||||
}
|
||||
|
||||
func TestValidateAuthSessionWithTxChecksAccessTokenIdentity(t *testing.T) {
|
||||
setupUserUpdateTestState(t)
|
||||
now := time.Now()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
prepare func(user *User, token *UserAccessToken, identity *AuthSessionIdentity)
|
||||
wantErr bool
|
||||
}{
|
||||
{name: "active token", prepare: func(*User, *UserAccessToken, *AuthSessionIdentity) {}},
|
||||
{name: "expired token", wantErr: true, prepare: func(_ *User, token *UserAccessToken, _ *AuthSessionIdentity) {
|
||||
require.NoError(t, DB.Model(token).Update("expires_at", now.Add(-time.Second).Unix()).Error)
|
||||
}},
|
||||
{name: "deleted token", wantErr: true, prepare: func(user *User, token *UserAccessToken, _ *AuthSessionIdentity) {
|
||||
_, err := DeleteUserAccessToken(user.Id, token.Id)
|
||||
require.NoError(t, err)
|
||||
}},
|
||||
{name: "another user's token", wantErr: true, prepare: func(_ *User, _ *UserAccessToken, identity *AuthSessionIdentity) {
|
||||
other := createAccessTokenTestUser(t, "token-identity-other")
|
||||
identity.UserID = other.Id
|
||||
}},
|
||||
{name: "stale auth version", wantErr: true, prepare: func(user *User, _ *UserAccessToken, _ *AuthSessionIdentity) {
|
||||
require.NoError(t, DB.Model(user).Update("auth_version", 2).Error)
|
||||
}},
|
||||
{name: "disabled user", wantErr: true, prepare: func(user *User, _ *UserAccessToken, _ *AuthSessionIdentity) {
|
||||
require.NoError(t, DB.Model(user).Update("status", common.UserStatusDisabled).Error)
|
||||
}},
|
||||
{name: "unexpected session version", wantErr: true, prepare: func(_ *User, _ *UserAccessToken, identity *AuthSessionIdentity) {
|
||||
identity.SessionVersion = 2
|
||||
}},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
require.NoError(t, DB.Exec("DELETE FROM user_access_tokens").Error)
|
||||
require.NoError(t, DB.Exec("DELETE FROM users").Error)
|
||||
user := createAccessTokenTestUser(t, "token-identity-user")
|
||||
token := createTestUserAccessToken(t, user.Id, "nap_identity_token", now.Add(time.Hour).Unix())
|
||||
identity := AuthSessionIdentity{
|
||||
UserID: user.Id, SessionID: AccessTokenSessionID(token.Id),
|
||||
UserAuthVersion: 1, SessionVersion: AccessTokenSessionVersion,
|
||||
}
|
||||
tt.prepare(&user, &token, &identity)
|
||||
|
||||
txErr := DB.Transaction(func(tx *gorm.DB) error { return ValidateAuthSessionWithTx(tx, identity) })
|
||||
plainErr := ValidateAccessTokenIdentity(identity)
|
||||
if !tt.wantErr {
|
||||
assert.NoError(t, txErr)
|
||||
assert.NoError(t, plainErr)
|
||||
return
|
||||
}
|
||||
assert.ErrorIs(t, txErr, ErrUserSessionInactive)
|
||||
assert.ErrorIs(t, plainErr, ErrUserSessionInactive)
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user