264 lines
6.7 KiB
Go
264 lines
6.7 KiB
Go
package centrifugo
|
|
|
|
import (
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/golang-jwt/jwt/v5"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestNewTokenGenerator(t *testing.T) {
|
|
secret := "test-secret"
|
|
gen := NewTokenGenerator(secret)
|
|
|
|
assert.NotNil(t, gen)
|
|
assert.Equal(t, secret, gen.config.Secret)
|
|
assert.Equal(t, time.Hour, gen.config.ExpireIn)
|
|
}
|
|
|
|
func TestNewTokenGeneratorWithConfig(t *testing.T) {
|
|
config := TokenConfig{
|
|
Secret: "custom-secret",
|
|
ExpireIn: 24 * time.Hour,
|
|
}
|
|
gen := NewTokenGeneratorWithConfig(config)
|
|
|
|
assert.NotNil(t, gen)
|
|
assert.Equal(t, config.Secret, gen.config.Secret)
|
|
assert.Equal(t, config.ExpireIn, gen.config.ExpireIn)
|
|
}
|
|
|
|
func TestNewTokenGeneratorWithConfig_DefaultExpire(t *testing.T) {
|
|
config := TokenConfig{
|
|
Secret: "test-secret",
|
|
ExpireIn: 0, // 應該使用預設值
|
|
}
|
|
gen := NewTokenGeneratorWithConfig(config)
|
|
|
|
assert.Equal(t, time.Hour, gen.config.ExpireIn)
|
|
}
|
|
|
|
func TestQuickConnectionToken(t *testing.T) {
|
|
gen := NewTokenGenerator("test-secret")
|
|
userID := "user-123"
|
|
|
|
token, err := gen.QuickConnectionToken(userID)
|
|
|
|
require.NoError(t, err)
|
|
assert.NotEmpty(t, token)
|
|
|
|
// 驗證 Token 內容
|
|
claims := &ConnectionClaims{}
|
|
parsedToken, err := jwt.ParseWithClaims(token, claims, func(token *jwt.Token) (interface{}, error) {
|
|
return []byte("test-secret"), nil
|
|
})
|
|
|
|
require.NoError(t, err)
|
|
assert.True(t, parsedToken.Valid)
|
|
assert.Equal(t, userID, claims.Sub)
|
|
assert.NotNil(t, claims.ExpiresAt)
|
|
assert.NotNil(t, claims.IssuedAt)
|
|
}
|
|
|
|
func TestGenerateConnectionToken(t *testing.T) {
|
|
gen := NewTokenGenerator("test-secret")
|
|
|
|
tests := []struct {
|
|
name string
|
|
opts ConnectionTokenOptions
|
|
checkFn func(t *testing.T, claims *ConnectionClaims)
|
|
}{
|
|
{
|
|
name: "basic token",
|
|
opts: ConnectionTokenOptions{
|
|
UserID: "user-123",
|
|
},
|
|
checkFn: func(t *testing.T, claims *ConnectionClaims) {
|
|
assert.Equal(t, "user-123", claims.Sub)
|
|
assert.Nil(t, claims.Info)
|
|
assert.Nil(t, claims.Channels)
|
|
},
|
|
},
|
|
{
|
|
name: "token with info",
|
|
opts: ConnectionTokenOptions{
|
|
UserID: "user-456",
|
|
Info: map[string]interface{}{
|
|
"name": "Daniel",
|
|
"role": "admin",
|
|
},
|
|
},
|
|
checkFn: func(t *testing.T, claims *ConnectionClaims) {
|
|
assert.Equal(t, "user-456", claims.Sub)
|
|
assert.NotNil(t, claims.Info)
|
|
assert.Equal(t, "Daniel", claims.Info["name"])
|
|
assert.Equal(t, "admin", claims.Info["role"])
|
|
},
|
|
},
|
|
{
|
|
name: "token with channels",
|
|
opts: ConnectionTokenOptions{
|
|
UserID: "user-789",
|
|
Channels: []string{"chat:room-1", "chat:room-2"},
|
|
},
|
|
checkFn: func(t *testing.T, claims *ConnectionClaims) {
|
|
assert.Equal(t, "user-789", claims.Sub)
|
|
assert.Equal(t, []string{"chat:room-1", "chat:room-2"}, claims.Channels)
|
|
},
|
|
},
|
|
{
|
|
name: "token with custom expire",
|
|
opts: ConnectionTokenOptions{
|
|
UserID: "user-abc",
|
|
ExpireAt: ptrTime(time.Now().Add(48 * time.Hour)),
|
|
},
|
|
checkFn: func(t *testing.T, claims *ConnectionClaims) {
|
|
assert.Equal(t, "user-abc", claims.Sub)
|
|
// 檢查過期時間大約在 48 小時後
|
|
expireTime := claims.ExpiresAt.Time
|
|
assert.True(t, expireTime.After(time.Now().Add(47*time.Hour)))
|
|
},
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
token, err := gen.GenerateConnectionToken(tt.opts)
|
|
require.NoError(t, err)
|
|
assert.NotEmpty(t, token)
|
|
|
|
// 解析 Token
|
|
claims := &ConnectionClaims{}
|
|
parsedToken, err := jwt.ParseWithClaims(token, claims, func(token *jwt.Token) (interface{}, error) {
|
|
return []byte("test-secret"), nil
|
|
})
|
|
|
|
require.NoError(t, err)
|
|
assert.True(t, parsedToken.Valid)
|
|
tt.checkFn(t, claims)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestQuickSubscriptionToken(t *testing.T) {
|
|
gen := NewTokenGenerator("test-secret")
|
|
userID := "user-123"
|
|
channel := "private:room-456"
|
|
|
|
token, err := gen.QuickSubscriptionToken(userID, channel)
|
|
|
|
require.NoError(t, err)
|
|
assert.NotEmpty(t, token)
|
|
|
|
// 驗證 Token 內容
|
|
claims := &SubscriptionClaims{}
|
|
parsedToken, err := jwt.ParseWithClaims(token, claims, func(token *jwt.Token) (interface{}, error) {
|
|
return []byte("test-secret"), nil
|
|
})
|
|
|
|
require.NoError(t, err)
|
|
assert.True(t, parsedToken.Valid)
|
|
assert.Equal(t, userID, claims.Sub)
|
|
assert.Equal(t, channel, claims.Channel)
|
|
}
|
|
|
|
func TestGenerateSubscriptionToken(t *testing.T) {
|
|
gen := NewTokenGenerator("test-secret")
|
|
|
|
opts := SubscriptionTokenOptions{
|
|
UserID: "user-123",
|
|
Channel: "private:room-456",
|
|
Info: map[string]interface{}{
|
|
"role": "moderator",
|
|
},
|
|
}
|
|
|
|
token, err := gen.GenerateSubscriptionToken(opts)
|
|
|
|
require.NoError(t, err)
|
|
assert.NotEmpty(t, token)
|
|
|
|
// 驗證 Token 內容
|
|
claims := &SubscriptionClaims{}
|
|
parsedToken, err := jwt.ParseWithClaims(token, claims, func(token *jwt.Token) (interface{}, error) {
|
|
return []byte("test-secret"), nil
|
|
})
|
|
|
|
require.NoError(t, err)
|
|
assert.True(t, parsedToken.Valid)
|
|
assert.Equal(t, opts.UserID, claims.Sub)
|
|
assert.Equal(t, opts.Channel, claims.Channel)
|
|
assert.Equal(t, "moderator", claims.Info["role"])
|
|
}
|
|
|
|
func TestGenerateAnonymousToken(t *testing.T) {
|
|
gen := NewTokenGenerator("test-secret")
|
|
|
|
token, err := gen.GenerateAnonymousToken()
|
|
|
|
require.NoError(t, err)
|
|
assert.NotEmpty(t, token)
|
|
|
|
// 驗證 Token 內容
|
|
claims := &ConnectionClaims{}
|
|
parsedToken, err := jwt.ParseWithClaims(token, claims, func(token *jwt.Token) (interface{}, error) {
|
|
return []byte("test-secret"), nil
|
|
})
|
|
|
|
require.NoError(t, err)
|
|
assert.True(t, parsedToken.Valid)
|
|
assert.Equal(t, "", claims.Sub) // 匿名用戶的 UserID 為空
|
|
}
|
|
|
|
func TestTokenExpiration(t *testing.T) {
|
|
// 創建一個很短過期時間的生成器
|
|
gen := NewTokenGeneratorWithConfig(TokenConfig{
|
|
Secret: "test-secret",
|
|
ExpireIn: 1 * time.Second,
|
|
})
|
|
|
|
token, err := gen.QuickConnectionToken("user-123")
|
|
require.NoError(t, err)
|
|
|
|
// 立即驗證應該成功
|
|
claims := &ConnectionClaims{}
|
|
parsedToken, err := jwt.ParseWithClaims(token, claims, func(token *jwt.Token) (interface{}, error) {
|
|
return []byte("test-secret"), nil
|
|
})
|
|
require.NoError(t, err)
|
|
assert.True(t, parsedToken.Valid)
|
|
|
|
// 等待過期
|
|
time.Sleep(2 * time.Second)
|
|
|
|
// 過期後驗證應該失敗
|
|
claims2 := &ConnectionClaims{}
|
|
_, err = jwt.ParseWithClaims(token, claims2, func(token *jwt.Token) (interface{}, error) {
|
|
return []byte("test-secret"), nil
|
|
})
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "token is expired")
|
|
}
|
|
|
|
func TestTokenWithWrongSecret(t *testing.T) {
|
|
gen := NewTokenGenerator("correct-secret")
|
|
|
|
token, err := gen.QuickConnectionToken("user-123")
|
|
require.NoError(t, err)
|
|
|
|
// 使用錯誤的密鑰驗證
|
|
claims := &ConnectionClaims{}
|
|
_, err = jwt.ParseWithClaims(token, claims, func(token *jwt.Token) (interface{}, error) {
|
|
return []byte("wrong-secret"), nil
|
|
})
|
|
|
|
assert.Error(t, err)
|
|
}
|
|
|
|
// 輔助函數
|
|
func ptrTime(t time.Time) *time.Time {
|
|
return &t
|
|
}
|