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 }