You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
121 lines
4.0 KiB
121 lines
4.0 KiB
package mcp
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/golang-jwt/jwt/v5"
|
|
)
|
|
|
|
const testKey = "test-token-key-please-ignore"
|
|
|
|
func mkToken(t *testing.T, uid string, key string, expired bool) string {
|
|
t.Helper()
|
|
exp := time.Now().Add(time.Hour)
|
|
if expired {
|
|
exp = time.Now().Add(-time.Hour)
|
|
}
|
|
tok := jwt.NewWithClaims(jwt.SigningMethodHS256, &jwt.RegisteredClaims{
|
|
ID: uid,
|
|
ExpiresAt: jwt.NewNumericDate(exp),
|
|
})
|
|
s, err := tok.SignedString([]byte(key))
|
|
if err != nil {
|
|
t.Fatalf("签 token 失败: %v", err)
|
|
}
|
|
return s
|
|
}
|
|
|
|
func mcpWith(key string) *Mcp {
|
|
return &Mcp{options: &Options{TokenKey: key}}
|
|
}
|
|
|
|
// 正常路径：能从 token 解出 uid
|
|
func TestResolveUID_FromToken(t *testing.T) {
|
|
m := mcpWith(testKey)
|
|
ctx := withAuthKey(context.Background(), mkToken(t, "u123", testKey, false))
|
|
uid, err := m.ResolveUID(ctx, "")
|
|
if err != nil || uid != "u123" {
|
|
t.Fatalf("期望 u123，得到 %q err=%v", uid, err)
|
|
}
|
|
}
|
|
|
|
// Bearer 前缀要能吃下：gateway 收的是裸 token，但标准 MCP 客户端会加前缀
|
|
func TestResolveUID_BearerPrefix(t *testing.T) {
|
|
m := mcpWith(testKey)
|
|
ctx := withAuthKey(context.Background(), "Bearer "+mkToken(t, "u123", testKey, false))
|
|
uid, err := m.ResolveUID(ctx, "")
|
|
if err != nil || uid != "u123" {
|
|
t.Fatalf("Bearer 前缀应被剥掉，得到 %q err=%v", uid, err)
|
|
}
|
|
}
|
|
|
|
// **核心安全断言**：参数里的 uid 与会话不符必须拒绝，
|
|
// 绝不能静默改用会话 uid（那会掩盖调用方的错误认知），更不能采信参数。
|
|
func TestResolveUID_RejectsMismatchedArgUID(t *testing.T) {
|
|
m := mcpWith(testKey)
|
|
ctx := withAuthKey(context.Background(), mkToken(t, "u123", testKey, false))
|
|
if _, err := m.ResolveUID(ctx, "someone_else"); err == nil {
|
|
t.Fatal("参数 uid 与会话不符时必须报错")
|
|
}
|
|
}
|
|
|
|
// 参数 uid 与会话一致时放行（百炼后台已配好的工具定义仍会传 uid）
|
|
func TestResolveUID_AllowsMatchingArgUID(t *testing.T) {
|
|
m := mcpWith(testKey)
|
|
ctx := withAuthKey(context.Background(), mkToken(t, "u123", testKey, false))
|
|
uid, err := m.ResolveUID(ctx, "u123")
|
|
if err != nil || uid != "u123" {
|
|
t.Fatalf("一致的 uid 应放行，得到 %q err=%v", uid, err)
|
|
}
|
|
}
|
|
|
|
// **核心安全断言**：没有 token 时必须拒绝，
|
|
// 绝不能回退成「用参数里的 uid」——那等于这层鉴权没做。
|
|
func TestResolveUID_NoTokenIsRejectedNotFallback(t *testing.T) {
|
|
m := mcpWith(testKey)
|
|
ctx := context.Background()
|
|
if uid, err := m.ResolveUID(ctx, "u123"); err == nil {
|
|
t.Fatalf("没有 token 时必须拒绝，却返回了 uid=%q", uid)
|
|
}
|
|
}
|
|
|
|
// 签名不对（伪造 token）必须拒绝
|
|
func TestResolveUID_RejectsWrongSignature(t *testing.T) {
|
|
m := mcpWith(testKey)
|
|
ctx := withAuthKey(context.Background(), mkToken(t, "u123", "another-key", false))
|
|
if _, err := m.ResolveUID(ctx, ""); err == nil {
|
|
t.Fatal("签名不匹配的 token 必须拒绝")
|
|
}
|
|
}
|
|
|
|
// 过期 token 必须拒绝
|
|
func TestResolveUID_RejectsExpired(t *testing.T) {
|
|
m := mcpWith(testKey)
|
|
ctx := withAuthKey(context.Background(), mkToken(t, "u123", testKey, true))
|
|
if _, err := m.ResolveUID(ctx, ""); err == nil {
|
|
t.Fatal("过期 token 必须拒绝")
|
|
}
|
|
}
|
|
|
|
// **核心安全断言**：服务端没配 TokenKey 时一律拒绝，不是放行。
|
|
// 配置漏了导致鉴权静默失效，比工具不可用严重得多。
|
|
func TestResolveUID_MissingTokenKeyDeniesAll(t *testing.T) {
|
|
m := mcpWith("")
|
|
ctx := withAuthKey(context.Background(), mkToken(t, "u123", testKey, false))
|
|
if _, err := m.ResolveUID(ctx, ""); err == nil {
|
|
t.Fatal("未配置 TokenKey 时必须拒绝所有用户数据类工具")
|
|
}
|
|
}
|
|
|
|
// token 里没有 uid（ID 为空）时不能返回空串当成合法身份 ——
|
|
// 空 uid 会让下游的 ApplyMemoryQuery 走 "1 = 0" 分支返回空，
|
|
// 但那是最后一道防线，不该指望它。
|
|
func TestResolveUID_RejectsEmptyIDClaim(t *testing.T) {
|
|
m := mcpWith(testKey)
|
|
ctx := withAuthKey(context.Background(), mkToken(t, "", testKey, false))
|
|
if uid, err := m.ResolveUID(ctx, ""); err == nil {
|
|
t.Fatalf("token 里没有 uid 时必须拒绝，却返回 %q", uid)
|
|
}
|
|
}
|
|
|