backend/pkg/library/centrifugo/blacklist_test.go

246 lines
5.8 KiB
Go

package centrifugo
import (
"context"
"testing"
"time"
"github.com/alicebob/miniredis/v2"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/zeromicro/go-zero/core/stores/redis"
)
func setupTestRedis(t *testing.T) (*redis.Redis, func()) {
mr, err := miniredis.Run()
require.NoError(t, err)
rds, err := redis.NewRedis(redis.RedisConf{
Host: mr.Addr(),
Type: "node",
})
require.NoError(t, err)
return rds, func() {
mr.Close()
}
}
func TestNewTokenBlacklist(t *testing.T) {
rds, cleanup := setupTestRedis(t)
defer cleanup()
blacklist := NewTokenBlacklist(rds)
assert.NotNil(t, blacklist)
assert.Equal(t, "centrifugo:blacklist:", blacklist.prefix)
}
func TestNewTokenBlacklistWithPrefix(t *testing.T) {
rds, cleanup := setupTestRedis(t)
defer cleanup()
blacklist := NewTokenBlacklistWithPrefix(rds, "custom:prefix:")
assert.NotNil(t, blacklist)
assert.Equal(t, "custom:prefix:", blacklist.prefix)
}
func TestRevokeToken(t *testing.T) {
rds, cleanup := setupTestRedis(t)
defer cleanup()
blacklist := NewTokenBlacklist(rds)
ctx := context.Background()
jti := "test-jti-123"
ttl := 1 * time.Hour
// 撤銷 Token
err := blacklist.RevokeToken(ctx, jti, ttl)
require.NoError(t, err)
// 檢查是否被撤銷
revoked, err := blacklist.IsTokenRevoked(ctx, jti)
require.NoError(t, err)
assert.True(t, revoked)
}
func TestRevokeToken_EmptyJTI(t *testing.T) {
rds, cleanup := setupTestRedis(t)
defer cleanup()
blacklist := NewTokenBlacklist(rds)
ctx := context.Background()
err := blacklist.RevokeToken(ctx, "", time.Hour)
assert.Error(t, err)
assert.Contains(t, err.Error(), "jti cannot be empty")
}
func TestIsTokenRevoked_NotRevoked(t *testing.T) {
rds, cleanup := setupTestRedis(t)
defer cleanup()
blacklist := NewTokenBlacklist(rds)
ctx := context.Background()
// 檢查未撤銷的 Token
revoked, err := blacklist.IsTokenRevoked(ctx, "non-existent-jti")
require.NoError(t, err)
assert.False(t, revoked)
}
func TestIsTokenRevoked_EmptyJTI(t *testing.T) {
rds, cleanup := setupTestRedis(t)
defer cleanup()
blacklist := NewTokenBlacklist(rds)
ctx := context.Background()
// 空 JTI 應該返回 false
revoked, err := blacklist.IsTokenRevoked(ctx, "")
require.NoError(t, err)
assert.False(t, revoked)
}
func TestRevokeUserTokens(t *testing.T) {
rds, cleanup := setupTestRedis(t)
defer cleanup()
blacklist := NewTokenBlacklist(rds)
ctx := context.Background()
userID := "user-123"
// 撤銷用戶所有 Token
err := blacklist.RevokeUserTokens(ctx, userID)
require.NoError(t, err)
// 獲取版本
version, err := blacklist.GetUserTokenVersion(ctx, userID)
require.NoError(t, err)
assert.Greater(t, version, int64(0))
}
func TestRevokeUserTokens_EmptyUserID(t *testing.T) {
rds, cleanup := setupTestRedis(t)
defer cleanup()
blacklist := NewTokenBlacklist(rds)
ctx := context.Background()
err := blacklist.RevokeUserTokens(ctx, "")
assert.Error(t, err)
assert.Contains(t, err.Error(), "userID cannot be empty")
}
func TestGetUserTokenVersion_NoVersion(t *testing.T) {
rds, cleanup := setupTestRedis(t)
defer cleanup()
blacklist := NewTokenBlacklist(rds)
ctx := context.Background()
// 未設置版本的用戶應該返回 0
version, err := blacklist.GetUserTokenVersion(ctx, "new-user")
require.NoError(t, err)
assert.Equal(t, int64(0), version)
}
func TestGetUserTokenVersion_EmptyUserID(t *testing.T) {
rds, cleanup := setupTestRedis(t)
defer cleanup()
blacklist := NewTokenBlacklist(rds)
ctx := context.Background()
version, err := blacklist.GetUserTokenVersion(ctx, "")
require.NoError(t, err)
assert.Equal(t, int64(0), version)
}
func TestIsTokenVersionValid(t *testing.T) {
rds, cleanup := setupTestRedis(t)
defer cleanup()
blacklist := NewTokenBlacklist(rds)
ctx := context.Background()
userID := "user-123"
// 未設置版本時,任何版本都應該有效
valid, err := blacklist.IsTokenVersionValid(ctx, userID, 0)
require.NoError(t, err)
assert.True(t, valid)
// 撤銷用戶 Token
err = blacklist.RevokeUserTokens(ctx, userID)
require.NoError(t, err)
// 獲取當前版本
currentVersion, err := blacklist.GetUserTokenVersion(ctx, userID)
require.NoError(t, err)
// 舊版本應該無效
valid, err = blacklist.IsTokenVersionValid(ctx, userID, currentVersion-1)
require.NoError(t, err)
assert.False(t, valid)
// 當前版本應該有效
valid, err = blacklist.IsTokenVersionValid(ctx, userID, currentVersion)
require.NoError(t, err)
assert.True(t, valid)
// 更新版本應該有效
valid, err = blacklist.IsTokenVersionValid(ctx, userID, currentVersion+1)
require.NoError(t, err)
assert.True(t, valid)
}
func TestRevokeUserTokens_MultipleRevokes(t *testing.T) {
rds, cleanup := setupTestRedis(t)
defer cleanup()
blacklist := NewTokenBlacklist(rds)
ctx := context.Background()
userID := "user-123"
// 第一次撤銷
err := blacklist.RevokeUserTokens(ctx, userID)
require.NoError(t, err)
version1, err := blacklist.GetUserTokenVersion(ctx, userID)
require.NoError(t, err)
// 等待一點時間確保時間戳不同
time.Sleep(10 * time.Millisecond)
// 第二次撤銷
err = blacklist.RevokeUserTokens(ctx, userID)
require.NoError(t, err)
version2, err := blacklist.GetUserTokenVersion(ctx, userID)
require.NoError(t, err)
// 第二次的版本應該更大
assert.Greater(t, version2, version1)
// 第一次的版本應該已經無效
valid, err := blacklist.IsTokenVersionValid(ctx, userID, version1)
require.NoError(t, err)
assert.False(t, valid)
}
func TestKeyGeneration(t *testing.T) {
rds, cleanup := setupTestRedis(t)
defer cleanup()
blacklist := NewTokenBlacklistWithPrefix(rds, "test:")
// 測試 key 生成
assert.Equal(t, "test:token:jti-123", blacklist.tokenKey("jti-123"))
assert.Equal(t, "test:user_version:user-456", blacklist.userVersionKey("user-456"))
}