starpivot/apps/services/modules/mcp/auth_mcptoken_test.go

package mcp

import (
	"testing"
	"time"

	"github.com/golang-jwt/jwt/v5"
)

// 签发在 user 模块、校验在 mcp 模块，两边隔着一个进程，
// 写错了不会编译报错，只会在真机上静默拒绝（或更糟：静默放行）。
// 这组用例把两边的口径钉在一起。

const testMcpKey = "test-mcp-key"

// signLikeUserModule 复刻 user/api_getmcptoken.go 的签发方式。
// ⚠️ 改那边的 claims 结构时，这里要一起改，否则这组用例就失去意义了。
func signLikeUserModule(t *testing.T, uid, audience string, ttl time.Duration, key string) string {
	t.Helper()
	claims := &jwt.RegisteredClaims{
		ID:        uid,
		ExpiresAt: jwt.NewNumericDate(time.Now().Add(ttl)),
		IssuedAt:  jwt.NewNumericDate(time.Now()),
	}
	if audience != "" {
		claims.Audience = jwt.ClaimStrings{audience}
	}
	s, err := jwt.NewWithClaims(jwt.SigningMethodHS256, claims).SignedString([]byte(key))
	if err != nil {
		t.Fatalf("签发失败: %v", err)
	}
	return s
}

func TestParseUIDWithAudience(t *testing.T) {
	const uid = "2095078768153460736"

	t.Run("正常令牌能解出uid", func(t *testing.T) {
		tok := signLikeUserModule(t, uid, mcpTokenAudience, time.Hour, testMcpKey)
		got, err := parseUIDWithAudience(tok, []byte(testMcpKey), mcpTokenAudience)
		if err != nil {
			t.Fatalf("不该失败: %v", err)
		}
		if got != uid {
			t.Fatalf("uid 不对: got=%s want=%s", got, uid)
		}
	})

	t.Run("换了密钥必须拒", func(t *testing.T) {
		tok := signLikeUserModule(t, uid, mcpTokenAudience, time.Hour, "another-key")
		if _, err := parseUIDWithAudience(tok, []byte(testMcpKey), mcpTokenAudience); err == nil {
			t.Fatal("别人签的令牌居然通过了")
		}
	})

	t.Run("用途不符必须拒", func(t *testing.T) {
		// 这条守的是「将来若有别处也用 McpTokenKey 签东西，别被 MCP 一并接受」
		tok := signLikeUserModule(t, uid, "someother", time.Hour, testMcpKey)
		if _, err := parseUIDWithAudience(tok, []byte(testMcpKey), mcpTokenAudience); err == nil {
			t.Fatal("用途不符的令牌居然通过了")
		}
	})

	t.Run("没有用途标记必须拒", func(t *testing.T) {
		tok := signLikeUserModule(t, uid, "", time.Hour, testMcpKey)
		if _, err := parseUIDWithAudience(tok, []byte(testMcpKey), mcpTokenAudience); err == nil {
			t.Fatal("没有 aud 的令牌居然通过了")
		}
	})

	t.Run("过期必须拒", func(t *testing.T) {
		tok := signLikeUserModule(t, uid, mcpTokenAudience, -time.Minute, testMcpKey)
		if _, err := parseUIDWithAudience(tok, []byte(testMcpKey), mcpTokenAudience); err == nil {
			t.Fatal("过期令牌居然通过了")
		}
	})

	t.Run("没有过期时间必须拒", func(t *testing.T) {
		// 没有 exp 的令牌等同长期有效，这套「短时效」的前提就没了
		s, err := jwt.NewWithClaims(jwt.SigningMethodHS256, &jwt.RegisteredClaims{
			ID:       uid,
			Audience: jwt.ClaimStrings{mcpTokenAudience},
		}).SignedString([]byte(testMcpKey))
		if err != nil {
			t.Fatalf("签发失败: %v", err)
		}
		if _, err := parseUIDWithAudience(s, []byte(testMcpKey), mcpTokenAudience); err == nil {
			t.Fatal("没有 exp 的令牌居然通过了")
		}
	})

	t.Run("Bearer前缀要能兼容", func(t *testing.T) {
		tok := signLikeUserModule(t, uid, mcpTokenAudience, time.Hour, testMcpKey)
		if _, err := parseUIDWithAudience("Bearer "+tok, []byte(testMcpKey), mcpTokenAudience); err != nil {
			t.Fatalf("带 Bearer 前缀不该失败: %v", err)
		}
	})
}

func TestResolveUIDWithToken(t *testing.T) {
	const uid = "2095078768153460736"
	m := &Mcp{}
	m.options = &Options{McpTokenKey: testMcpKey}

	t.Run("参数令牌可作为身份来源", func(t *testing.T) {
		tok := signLikeUserModule(t, uid, mcpTokenAudience, time.Hour, testMcpKey)
		got, err := m.ResolveUIDWithToken(t.Context(), "", tok)
		if err != nil {
			t.Fatalf("不该失败: %v", err)
		}
		if got != uid {
			t.Fatalf("uid 不对: %s", got)
		}
	})

	t.Run("参数uid与令牌不符必须拒", func(t *testing.T) {
		tok := signLikeUserModule(t, uid, mcpTokenAudience, time.Hour, testMcpKey)
		if _, err := m.ResolveUIDWithToken(t.Context(), "9999", tok); err == nil {
			t.Fatal("冒充别人居然通过了")
		}
	})

	t.Run("只给裸uid不给令牌必须拒", func(t *testing.T) {
		// 这条是整层鉴权的底线：裸 uid 谁都能填，绝不能作为身份来源
		if _, err := m.ResolveUIDWithToken(t.Context(), uid, ""); err == nil {
			t.Fatal("只凭一个 uid 参数就放行了")
		}
	})

	t.Run("没配密钥时带令牌也要拒", func(t *testing.T) {
		empty := &Mcp{}
		empty.options = &Options{}
		tok := signLikeUserModule(t, uid, mcpTokenAudience, time.Hour, testMcpKey)
		if _, err := empty.ResolveUIDWithToken(t.Context(), "", tok); err == nil {
			t.Fatal("密钥没配却放行了")
		}
	})
}