go-kit/gin/jwtx/test/Tool_test.go

package test

import (
	"crypto/ed25519"
	"encoding/pem"
	"os"
	"testing"
	"time"

	"github.com/5-say/go-kit/gin/jwtx/tool"
	"github.com/golang-jwt/jwt/v5"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// TestGeneratePrivateKey 测试 GeneratePrivateKey 函数
func TestGeneratePrivateKey(t *testing.T) {
	t.Run("生成有效的私钥", func(t *testing.T) {
		privateKey, err := tool.GeneratePrivateKey()
		require.NoError(t, err)
		assert.NotNil(t, privateKey)
		assert.Equal(t, ed25519.PrivateKeySize, len(*privateKey))
	})

	t.Run("生成的私钥可以用于签名", func(t *testing.T) {
		privateKey, err := tool.GeneratePrivateKey()
		require.NoError(t, err)

		// 测试私钥可以用于签名
		message := []byte("test message")
		signature := ed25519.Sign(*privateKey, message)

		// 验证签名
		publicKey := privateKey.Public().(ed25519.PublicKey)
		isValid := ed25519.Verify(publicKey, message, signature)
		assert.True(t, isValid)
	})
}

// TestGenerateToken 测试 GenerateToken 函数
func TestGenerateToken(t *testing.T) {
	// 生成测试用的私钥
	privateKey, err := tool.GeneratePrivateKey()
	require.NoError(t, err)

	t.Run("生成有效的 JWT token", func(t *testing.T) {
		claims := jwt.MapClaims{
			"user_id": "12345",
			"role":    "admin",
			"iss":     "test-issuer",
			"exp":     time.Now().Add(time.Hour).Unix(),
			"iat":     time.Now().Unix(),
		}

		token, err := tool.GenerateToken(privateKey, claims)
		require.NoError(t, err)
		assert.NotEmpty(t, token)
		assert.Contains(t, token, ".") // JWT token 应该包含点分隔符
	})

	t.Run("生成的 token 可以被正确解析", func(t *testing.T) {
		claims := jwt.MapClaims{
			"user_id": "67890",
			"email":   "test@example.com",
			"exp":     time.Now().Add(time.Hour).Unix(),
		}

		token, err := tool.GenerateToken(privateKey, claims)
		require.NoError(t, err)

		// 使用公钥解析 token
		publicKey := privateKey.Public().(ed25519.PublicKey)
		parsedClaims, err := tool.ParseToken(token, &publicKey)
		require.NoError(t, err)

		assert.Equal(t, "67890", parsedClaims["user_id"])
		assert.Equal(t, "test@example.com", parsedClaims["email"])
	})

	t.Run("空声明生成 token", func(t *testing.T) {
		claims := jwt.MapClaims{}

		token, err := tool.GenerateToken(privateKey, claims)
		require.NoError(t, err)
		assert.NotEmpty(t, token)
	})
}

// TestParseToken 测试 ParseToken 函数
func TestParseToken(t *testing.T) {
	// 生成测试用的密钥对
	privateKey, err := tool.GeneratePrivateKey()
	require.NoError(t, err)
	publicKey := privateKey.Public().(ed25519.PublicKey)

	t.Run("解析有效的 token", func(t *testing.T) {
		originalClaims := jwt.MapClaims{
			"user_id": "test-user",
			"scope":   "read write",
			"exp":     time.Now().Add(time.Hour).Unix(),
		}

		token, err := tool.GenerateToken(privateKey, originalClaims)
		require.NoError(t, err)

		parsedClaims, err := tool.ParseToken(token, &publicKey)
		require.NoError(t, err)

		assert.Equal(t, originalClaims["user_id"], parsedClaims["user_id"])
		assert.Equal(t, originalClaims["scope"], parsedClaims["scope"])
	})

	t.Run("解析无效的 token 应该失败", func(t *testing.T) {
		invalidToken := "invalid.jwt.token"

		_, err := tool.ParseToken(invalidToken, &publicKey)
		assert.Error(t, err)
	})

	t.Run("使用错误的公钥解析应该失败", func(t *testing.T) {
		// 生成另一个密钥对
		wrongPrivateKey, err := tool.GeneratePrivateKey()
		require.NoError(t, err)
		wrongPublicKey := wrongPrivateKey.Public().(ed25519.PublicKey)

		claims := jwt.MapClaims{"test": "value"}
		token, err := tool.GenerateToken(privateKey, claims)
		require.NoError(t, err)

		// 使用错误的公钥解析
		_, err = tool.ParseToken(token, &wrongPublicKey)
		assert.Error(t, err)
	})

	t.Run("解析过期的 token 应该失败", func(t *testing.T) {
		// 创建过期的 claims
		expiredClaims := jwt.MapClaims{
			"user_id": "expired-user",
			"exp":     time.Now().Add(-time.Hour).Unix(), // 过去的时间
		}

		token, err := tool.GenerateToken(privateKey, expiredClaims)
		require.NoError(t, err)

		_, err = tool.ParseToken(token, &publicKey)
		assert.Error(t, err)
		assert.Contains(t, err.Error(), "expired")
	})
}

// TestGetPrivateKey 测试 GetPrivateKey 函数
func TestGetPrivateKey(t *testing.T) {
	// 创建临时文件用于测试
	tempFile, err := os.CreateTemp("", "test_private_key_*.pem")
	require.NoError(t, err)
	defer os.Remove(tempFile.Name())

	t.Run("从有效的 PEM 文件加载私钥", func(t *testing.T) {
		// 生成私钥并保存到文件
		privateKey, err := tool.GeneratePrivateKey()
		require.NoError(t, err)

		// 将私钥保存为 PEM 格式
		pemData := pemEncodePrivateKey(*privateKey)
		err = os.WriteFile(tempFile.Name(), pemData, 0600)
		require.NoError(t, err)

		// 从文件加载私钥
		loadedKey, err := tool.GetPrivateKey(tempFile.Name())
		require.NoError(t, err)
		assert.NotNil(t, loadedKey)
		assert.Equal(t, *privateKey, *loadedKey)
	})

	t.Run("加载不存在的文件应该失败", func(t *testing.T) {
		_, err := tool.GetPrivateKey("/nonexistent/file.pem")
		assert.Error(t, err)
	})

	t.Run("加载无效的 PEM 文件应该失败", func(t *testing.T) {
		// 写入无效的 PEM 数据
		err := os.WriteFile(tempFile.Name(), []byte("invalid pem data"), 0600)
		require.NoError(t, err)

		_, err = tool.GetPrivateKey(tempFile.Name())
		assert.Error(t, err)
	})

	t.Run("加载非私钥的 PEM 文件应该失败", func(t *testing.T) {
		// 写入公钥 PEM 数据（应该失败）
		privateKey, err := tool.GeneratePrivateKey()
		require.NoError(t, err)
		publicKey := privateKey.Public().(ed25519.PublicKey)

		pemData := pemEncodePublicKey(publicKey)
		err = os.WriteFile(tempFile.Name(), pemData, 0600)
		require.NoError(t, err)

		_, err = tool.GetPrivateKey(tempFile.Name())
		assert.Error(t, err)
	})
}

// TestIntegration 集成测试：完整的 JWT 流程
func TestIntegration(t *testing.T) {
	t.Run("完整的 JWT 生成和验证流程", func(t *testing.T) {
		// 1. 生成私钥
		privateKey, err := tool.GeneratePrivateKey()
		require.NoError(t, err)

		// 2. 创建 claims
		claims := jwt.MapClaims{
			"user_id": "integration-test",
			"name":    "Test User",
			"email":   "test@integration.com",
			"roles":   []string{"user", "admin"},
			"exp":     time.Now().Add(2 * time.Hour).Unix(),
			"iat":     time.Now().Unix(),
			"iss":     "integration-test-server",
		}

		// 3. 生成 JWT token
		token, err := tool.GenerateToken(privateKey, claims)
		require.NoError(t, err)
		assert.NotEmpty(t, token)

		// 4. 获取公钥
		publicKey := privateKey.Public().(ed25519.PublicKey)

		// 5. 解析和验证 token
		parsedClaims, err := tool.ParseToken(token, &publicKey)
		require.NoError(t, err)

		// 6. 验证 claims 内容
		assert.Equal(t, claims["user_id"], parsedClaims["user_id"])
		assert.Equal(t, claims["name"], parsedClaims["name"])
		assert.Equal(t, claims["email"], parsedClaims["email"])
		assert.Equal(t, claims["iss"], parsedClaims["iss"])

		// 7. 验证 token 可以用于实际场景
		// 例如：在中间件中验证 token 的有效性
		_, err = tool.ParseToken(token, &publicKey)
		assert.NoError(t, err, "token 应该可以重复验证")
	})
}

// pemEncodePrivateKey 将 Ed25519 私钥编码为 PEM 格式
func pemEncodePrivateKey(privateKey ed25519.PrivateKey) []byte {
	// 注意：ed25519.PrivateKey 实际上是 []byte，可以直接作为 PEM 数据
	block := &pem.Block{
		Type:  "PRIVATE KEY",
		Bytes: privateKey,
	}
	return pem.EncodeToMemory(block)
}

// pemEncodePublicKey 将 Ed25519 公钥编码为 PEM 格式
func pemEncodePublicKey(publicKey ed25519.PublicKey) []byte {
	block := &pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: publicKey,
	}
	return pem.EncodeToMemory(block)
}