perms-system-server/internal/logic/auth/jwt_test.go

package auth

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"github.com/golang-jwt/jwt/v4"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"perms-system-server/internal/consts"
	"perms-system-server/internal/middleware"
	"strings"
	"testing"
	"time"
)

const testSecret = "test-jwt-secret-key"

// TC-0292: secret="s", expire=3600, userId=1, username="u", productCode="p", memberType=""
func TestGenerateAccessToken(t *testing.T) {
	tests := []struct {
		name         string
		secret       string
		expire       int64
		userId       int64
		username     string
		productCode  string
		memberType   string
		tokenVersion int64
	}{
		{
			name:        "normal generation",
			secret:      testSecret,
			expire:      3600,
			userId:      1,
			username:    "admin",
			productCode: "p1",
			memberType:  "ADMIN",
		},
		{
			name:        "empty productCode",
			secret:      testSecret,
			expire:      3600,
			userId:      3,
			username:    "user2",
			productCode: "",
			memberType:  "",
		},
		{
			name:         "super admin with tokenVersion",
			secret:       testSecret,
			expire:       7200,
			userId:       100,
			username:     "super",
			productCode:  "p1",
			memberType:   "SUPER_ADMIN",
			tokenVersion: 5,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			tokenStr, err := GenerateAccessToken(tt.secret, tt.expire, tt.userId, tt.username, tt.productCode, tt.memberType, tt.tokenVersion)
			require.NoError(t, err)
			assert.NotEmpty(t, tokenStr)

			token, err := jwt.ParseWithClaims(tokenStr, &middleware.Claims{}, func(token *jwt.Token) (interface{}, error) {
				return []byte(tt.secret), nil
			})
			require.NoError(t, err)
			assert.True(t, token.Valid)

			claims, ok := token.Claims.(*middleware.Claims)
			require.True(t, ok)
			assert.Equal(t, tt.userId, claims.UserId)
			assert.Equal(t, tt.username, claims.Username)
			assert.Equal(t, tt.productCode, claims.ProductCode)
			assert.Equal(t, tt.memberType, claims.MemberType)
			assert.Equal(t, tt.tokenVersion, claims.TokenVersion)

			// 项 ：`perms` 字段已从 Claims 结构体中移除。
			// 解析原始 JWT payload，确保 token JSON 中不存在 "perms" key。
			segments := strings.Split(tokenStr, ".")
			require.Len(t, segments, 3, "jwt must have 3 segments")
			payloadBytes, err := base64.RawURLEncoding.DecodeString(segments[1])
			require.NoError(t, err)
			var raw map[string]interface{}
			require.NoError(t, json.Unmarshal(payloadBytes, &raw))
			_, hasPerms := raw["perms"]
			assert.False(t, hasPerms, "access token payload must NOT contain perms field")
		})
	}
}

// TC-0296: expireSeconds=1, sleep 2s
func TestGenerateAccessToken_Expiry(t *testing.T) {
	tokenStr, err := GenerateAccessToken(testSecret, 1, 1, "u", "", "", 0)
	require.NoError(t, err)

	time.Sleep(2 * time.Second)

	_, err = jwt.ParseWithClaims(tokenStr, &middleware.Claims{}, func(token *jwt.Token) (interface{}, error) {
		return []byte(testSecret), nil
	})
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "token is expired")
}

// TC-0297: secret="s", expire=86400, userId=1, productCode="p"
func TestGenerateRefreshToken(t *testing.T) {
	tests := []struct {
		name        string
		secret      string
		expire      int64
		userId      int64
		productCode string
	}{
		{"normal", testSecret, 86400, 1, "p1"},
		{"empty productCode", testSecret, 86400, 2, ""},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			tokenStr, err := GenerateRefreshToken(tt.secret, tt.expire, tt.userId, tt.productCode, 0)
			require.NoError(t, err)
			assert.NotEmpty(t, tokenStr)

			claims, err := ParseRefreshToken(tokenStr, tt.secret)
			require.NoError(t, err)
			assert.Equal(t, tt.userId, claims.UserId)
			assert.Equal(t, tt.productCode, claims.ProductCode)
		})
	}
}

// TC-0300: 有效token+正确secret
func TestParseRefreshToken(t *testing.T) {
	validToken, err := GenerateRefreshToken(testSecret, 3600, 42, "prod", 0)
	require.NoError(t, err)

	t.Run("valid token", func(t *testing.T) {
		claims, err := ParseRefreshToken(validToken, testSecret)
		require.NoError(t, err)
		assert.Equal(t, int64(42), claims.UserId)
		assert.Equal(t, "prod", claims.ProductCode)
	})

	t.Run("wrong secret", func(t *testing.T) {
		_, err := ParseRefreshToken(validToken, "wrong-secret")
		assert.Error(t, err)
	})

	t.Run("invalid token string", func(t *testing.T) {
		_, err := ParseRefreshToken("not-a-valid-token", testSecret)
		assert.Error(t, err)
	})

	t.Run("empty token", func(t *testing.T) {
		_, err := ParseRefreshToken("", testSecret)
		assert.Error(t, err)
	})

	t.Run("expired token", func(t *testing.T) {
		expiredToken, err := GenerateRefreshToken(testSecret, 1, 1, "p", 0)
		require.NoError(t, err)
		time.Sleep(2 * time.Second)
		_, err = ParseRefreshToken(expiredToken, testSecret)
		assert.Error(t, err)
	})

	// TC-0305: AccessToken误用 — TokenType校验拒绝
	t.Run("access token used as refresh - should be rejected", func(t *testing.T) {
		accessToken, err := GenerateAccessToken(testSecret, 3600, 1, "u", "p", "M", 0)
		require.NoError(t, err)
		_, err = ParseRefreshToken(accessToken, testSecret)
		assert.Error(t, err, "BUG-002: access token 不应被 ParseRefreshToken 接受，应通过 TokenType 字段区分")
	})
}

// TC-0294: secret=""
func TestGenerateAccessToken_EmptySecret(t *testing.T) {
	tokenStr, err := GenerateAccessToken("", 3600, 1, "u", "p", "M", 0)
	require.NoError(t, err)
	assert.NotEmpty(t, tokenStr)

	token, err := jwt.ParseWithClaims(tokenStr, &middleware.Claims{}, func(token *jwt.Token) (interface{}, error) {
		return []byte(""), nil
	})
	require.NoError(t, err)
	assert.True(t, token.Valid)

	claims, ok := token.Claims.(*middleware.Claims)
	require.True(t, ok)
	assert.Equal(t, int64(1), claims.UserId)
}

// ---------------------------------------------------------------------------
// 覆盖目标：ParseWithHMAC 必须显式断言 token.Method 为
// *jwt.SigningMethodHMAC，拒绝任何非 HMAC 的 alg 头，包括 "none" / "RS256" 等。
// 这里不等同于 jwt-go v4 对 "alg=none" 的默认拒绝，而是深度防御的显式白名单校验，
// 杜绝未来迁移到 RSA/ECDSA 时攻击者把公钥当共享密钥伪造 HS256 token
// （CVE-2016-10555 同类问题、OWASP JWT / RFC 8725 要求）。
// ---------------------------------------------------------------------------

const h4Secret = "h4-audit-secret-key"

// b64url returns the jwt-style base64url (no padding) encoding.
func b64url(b []byte) string { return base64.RawURLEncoding.EncodeToString(b) }

// forgeToken 手动拼接一个 JWT：自定义 header.alg + payload，再用任意密钥做 HMAC 签名。
// 这用于模拟"攻击者伪造头部 alg 但签名仍走 HS256"的场景。
func forgeToken(t *testing.T, alg string, claims any, signingKey string) string {
	t.Helper()
	header := map[string]string{"alg": alg, "typ": "JWT"}
	hBytes, err := json.Marshal(header)
	require.NoError(t, err)
	pBytes, err := json.Marshal(claims)
	require.NoError(t, err)

	signingInput := b64url(hBytes) + "." + b64url(pBytes)
	mac := hmac.New(sha256.New, []byte(signingKey))
	mac.Write([]byte(signingInput))
	sig := mac.Sum(nil)
	return signingInput + "." + b64url(sig)
}

// forgeTokenNoSig 拼接一个没有签名的 token（alg=none 典型攻击，第三段签名留空）。
func forgeTokenNoSig(t *testing.T, alg string, claims any) string {
	t.Helper()
	header := map[string]string{"alg": alg, "typ": "JWT"}
	hBytes, err := json.Marshal(header)
	require.NoError(t, err)
	pBytes, err := json.Marshal(claims)
	require.NoError(t, err)
	return b64url(hBytes) + "." + b64url(pBytes) + "."
}

// validRefreshClaims 返回一组完整、未过期的 refresh claims，用于伪造攻击 token。
func validRefreshClaims() RefreshClaims {
	now := time.Now()
	return RefreshClaims{
		TokenType:    consts.TokenTypeRefresh,
		UserId:       7,
		ProductCode:  "h4_pc",
		TokenVersion: 0,
		RegisteredClaims: jwt.RegisteredClaims{
			ExpiresAt: jwt.NewNumericDate(now.Add(1 * time.Hour)),
			IssuedAt:  jwt.NewNumericDate(now),
		},
	}
}

// TC-0951: 正常 HS256 token 必须被 ParseWithHMAC 正确接受。
func TestParseWithHMAC_HS256_Valid(t *testing.T) {
	tok, err := GenerateRefreshToken(h4Secret, 3600, 7, "h4_pc", 0)
	require.NoError(t, err)

	token, err := ParseWithHMAC(tok, h4Secret, &RefreshClaims{})
	require.NoError(t, err)
	assert.True(t, token.Valid)
	claims, ok := token.Claims.(*RefreshClaims)
	require.True(t, ok)
	assert.Equal(t, int64(7), claims.UserId)
	assert.Equal(t, consts.TokenTypeRefresh, claims.TokenType)
}

// TC-0952: alg=none 的伪造 token 必须被拒绝。
// jwt-go v4 默认就会拦住 "none"，但显式 HMAC 断言保证即使 lib 行为变化我们仍 fail-close。
func TestParseWithHMAC_AlgNone_Rejected(t *testing.T) {
	forged := forgeTokenNoSig(t, "none", validRefreshClaims())

	_, err := ParseWithHMAC(forged, h4Secret, &RefreshClaims{})
	require.Error(t, err, "alg=none 必须被 ParseWithHMAC 拒绝")
}

// TC-0953: 攻击者把 header alg 改成 RS256 但仍用 secret 作 HS256 签名
// （RSA 公钥 → HMAC secret 混淆攻击）。必须被 ParseWithHMAC 显式拒绝：
// 命中 keyfunc 的 `token.Method.(*SigningMethodHMAC)` 断言失败分支。
func TestParseWithHMAC_RS256HeaderButHMACSigned_Rejected(t *testing.T) {
	forged := forgeToken(t, "RS256", validRefreshClaims(), h4Secret)

	_, err := ParseWithHMAC(forged, h4Secret, &RefreshClaims{})
	require.Error(t, err, "alg=RS256 必须被 ParseWithHMAC 拒绝")
	assert.Contains(t, err.Error(), "unexpected signing method",
		"错误信息必须明确指出 alg 与预期不符（便于运维快速定位攻击尝试）")
}

// TC-0954: alg=ES256 同样应被拒绝（非 HMAC 算法一律拒绝）。
func TestParseWithHMAC_ES256HeaderButHMACSigned_Rejected(t *testing.T) {
	forged := forgeToken(t, "ES256", validRefreshClaims(), h4Secret)

	_, err := ParseWithHMAC(forged, h4Secret, &RefreshClaims{})
	require.Error(t, err)
	assert.Contains(t, err.Error(), "unexpected signing method")
}

// TC-0955: alg=HS256 但用错误的 secret 签名应被拒绝（签名校验失败路径）。
func TestParseWithHMAC_HS256WrongSecret_Rejected(t *testing.T) {
	tok, err := GenerateRefreshToken("attacker-guessed-secret", 3600, 7, "h4_pc", 0)
	require.NoError(t, err)

	_, err = ParseWithHMAC(tok, h4Secret, &RefreshClaims{})
	require.Error(t, err, "签名校验失败必须回错，不得放行")
}

// TC-0956: ParseRefreshToken（对外真实入口）也走 HMAC 断言，alg=RS256 必须被拒。
// 保证 ParseWithHMAC 不是孤立函数，而是已被真实调用链使用。
func TestParseRefreshToken_RS256Header_Rejected(t *testing.T) {
	forged := forgeToken(t, "RS256", validRefreshClaims(), h4Secret)
	_, err := ParseRefreshToken(forged, h4Secret)
	require.Error(t, err, "ParseRefreshToken 必须转交 ParseWithHMAC 拒绝 RS256 伪造 token")
}

// TC-0957: ParseRefreshToken 对 alg=none 的 token 也必须拒绝。
func TestParseRefreshToken_AlgNone_Rejected(t *testing.T) {
	forged := forgeTokenNoSig(t, "none", validRefreshClaims())
	_, err := ParseRefreshToken(forged, h4Secret)
	require.Error(t, err)
}

// TC-0958:  回归 —— 格式错误的 token（非三段式）必须 error 而不是 panic。
func TestParseWithHMAC_Malformed_Rejected(t *testing.T) {
	cases := []string{
		"",
		"not-a-token",
		"only.two",
		"a.b.c.d", // 四段
	}
	for _, s := range cases {
		t.Run("malformed:"+s, func(t *testing.T) {
			_, err := ParseWithHMAC(s, h4Secret, &RefreshClaims{})
			require.Error(t, err)
		})
	}
}

// TC-0959: payload 中 TokenType 非 refresh 的 HS256 token 应被 ParseRefreshToken
// 以 ErrTokenTypeMismatch 拒绝。确认  修复不会误吞该业务校验。
func TestParseRefreshToken_AccessTokenRejectedWithTypeMismatch(t *testing.T) {
	accessTok, err := GenerateAccessToken(h4Secret, 3600, 7, "u", "p", "M", 0)
	require.NoError(t, err)
	_, err = ParseRefreshToken(accessTok, h4Secret)
	require.Error(t, err)
	assert.Equal(t, ErrTokenTypeMismatch, err,
		"的 ParseWithHMAC 不能吞掉业务层 TokenType 校验错误")
}

// TC-0960: 伪造 alg=HS256 但 header.typ 异常（如 "JWT"→"xxx"）也不能绕过
// HMAC 校验。此用例用来证明只要底层签名正确，header 其余字段不影响放行/拒绝的核心语义。
// 反之，任何 alg 头不是 HS* 的一律拒，和 typ 无关。
func TestParseWithHMAC_HS256UnusualTyp_Accepted(t *testing.T) {
	// header.alg = HS256, header.typ = "JWT+weird"，签名正确 → 应放行（typ 不参与断言）
	header := map[string]string{"alg": "HS256", "typ": "JWT+weird"}
	hBytes, _ := json.Marshal(header)
	claims := validRefreshClaims()
	pBytes, _ := json.Marshal(claims)
	signingInput := b64url(hBytes) + "." + b64url(pBytes)
	mac := hmac.New(sha256.New, []byte(h4Secret))
	mac.Write([]byte(signingInput))
	tok := signingInput + "." + b64url(mac.Sum(nil))

	_, err := ParseWithHMAC(tok, h4Secret, &RefreshClaims{})
	require.NoError(t, err,
		"HMAC 断言只看 alg，typ 不属于签名算法白名单范畴，正常 HS256 应放行")
}

// 辅助：保持 strings 导入被使用，避免 go vet 警告。
var _ = strings.Split

// 确保 middleware.Claims 在包内可被用于 TypeRefresh / TypeAccess 等正反测试（未来扩展）。
var _ = middleware.Claims{}