diff --git a/internal/mcp/token_service_test.go b/internal/mcp/token_service_test.go new file mode 100644 index 0000000..8e21980 --- /dev/null +++ b/internal/mcp/token_service_test.go @@ -0,0 +1,341 @@ +package mcp_test + +import ( + "context" + "errors" + "sync" + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/yan1h/agent-coding-workflow/internal/audit" + "github.com/yan1h/agent-coding-workflow/internal/infra/errs" + "github.com/yan1h/agent-coding-workflow/internal/mcp" +) + +// ===== fake repo ===== + +type fakeRepo struct { + mu sync.Mutex + tokens map[uuid.UUID]*mcp.MCPToken + byHash map[string]uuid.UUID +} + +func newFakeRepo() *fakeRepo { + return &fakeRepo{ + tokens: map[uuid.UUID]*mcp.MCPToken{}, + byHash: map[string]uuid.UUID{}, + } +} + +func (r *fakeRepo) Create(ctx context.Context, t *mcp.MCPToken) (*mcp.MCPToken, error) { + r.mu.Lock() + defer r.mu.Unlock() + c := *t + c.CreatedAt = time.Now() + r.tokens[c.ID] = &c + r.byHash[string(c.TokenHash)] = c.ID + return &c, nil +} +func (r *fakeRepo) GetByHash(ctx context.Context, h []byte) (*mcp.MCPToken, error) { + r.mu.Lock() + defer r.mu.Unlock() + id, ok := r.byHash[string(h)] + if !ok { + return nil, errs.New(errs.CodeMcpTokenInvalid, "not found") + } + c := *r.tokens[id] + return &c, nil +} +func (r *fakeRepo) GetByID(ctx context.Context, id uuid.UUID) (*mcp.MCPToken, error) { + r.mu.Lock() + defer r.mu.Unlock() + t, ok := r.tokens[id] + if !ok { + return nil, errs.New(errs.CodeMcpTokenNotFound, "not found") + } + c := *t + return &c, nil +} +func (r *fakeRepo) List(ctx context.Context, caller *uuid.UUID, activeOnly bool) ([]*mcp.MCPToken, error) { + r.mu.Lock() + defer r.mu.Unlock() + var out []*mcp.MCPToken + for _, t := range r.tokens { + if caller != nil && (t.UserID != *caller || t.Issuer != mcp.IssuerAdmin) { + continue + } + if activeOnly && !t.IsActive(time.Now()) { + continue + } + c := *t + out = append(out, &c) + } + return out, nil +} +func (r *fakeRepo) UpdateLastUsed(ctx context.Context, id uuid.UUID) error { + r.mu.Lock() + defer r.mu.Unlock() + if t, ok := r.tokens[id]; ok { + now := time.Now() + t.LastUsedAt = &now + } + return nil +} +func (r *fakeRepo) Revoke(ctx context.Context, id uuid.UUID) (uuid.UUID, error) { + r.mu.Lock() + defer r.mu.Unlock() + t, ok := r.tokens[id] + if !ok || t.RevokedAt != nil { + return uuid.Nil, errs.New(errs.CodeMcpTokenNotFound, "not found") + } + now := time.Now() + t.RevokedAt = &now + return id, nil +} +func (r *fakeRepo) RevokeBySession(ctx context.Context, sid uuid.UUID) ([]uuid.UUID, error) { + r.mu.Lock() + defer r.mu.Unlock() + var revoked []uuid.UUID + now := time.Now() + for _, t := range r.tokens { + if t.Issuer == mcp.IssuerSystem && t.ACPSessionID != nil && *t.ACPSessionID == sid && t.RevokedAt == nil { + t.RevokedAt = &now + revoked = append(revoked, t.ID) + } + } + return revoked, nil +} +func (r *fakeRepo) ReapStaleSystemTokens(ctx context.Context) ([]uuid.UUID, error) { + return nil, nil +} +func (r *fakeRepo) PurgeExpired(ctx context.Context, before time.Time) (int64, error) { + return 0, nil +} +func (r *fakeRepo) DeleteByID(ctx context.Context, id uuid.UUID) error { + r.mu.Lock() + defer r.mu.Unlock() + delete(r.tokens, id) + return nil +} + +// ===== fake audit ===== + +type fakeAudit struct { + mu sync.Mutex + entries []audit.Entry +} + +func (a *fakeAudit) Record(ctx context.Context, e audit.Entry) error { + a.mu.Lock() + defer a.mu.Unlock() + a.entries = append(a.entries, e) + return nil +} +func (a *fakeAudit) actions() []string { + a.mu.Lock() + defer a.mu.Unlock() + out := make([]string, 0, len(a.entries)) + for _, e := range a.entries { + out = append(out, e.Action) + } + return out +} + +// ===== tests ===== + +func TestIssueAdmin_HappyPath(t *testing.T) { + t.Parallel() + repo := newFakeRepo() + au := &fakeAudit{} + svc := mcp.NewTokenService(repo, au) + + uid := uuid.New() + in := mcp.IssueAdminInput{ + UserID: uid, + Name: "test", + Scope: mcp.Scope{Tools: []string{"list_issues"}}, + ExpiresAt: time.Now().Add(time.Hour), + CreatedBy: uid, + } + res, err := svc.IssueAdmin(context.Background(), in) + require.NoError(t, err) + assert.NotEmpty(t, res.Plaintext) + assert.Equal(t, []string{"mcp.token.create"}, au.actions()) +} + +func TestIssueAdmin_NameRequired(t *testing.T) { + t.Parallel() + svc := mcp.NewTokenService(newFakeRepo(), &fakeAudit{}) + + _, err := svc.IssueAdmin(context.Background(), mcp.IssueAdminInput{ + UserID: uuid.New(), + Name: "", + ExpiresAt: time.Now().Add(time.Hour), + CreatedBy: uuid.New(), + }) + ae, ok := errs.As(err) + require.True(t, ok) + assert.Equal(t, errs.CodeMcpTokenNameRequired, ae.Code) +} + +func TestIssueAdmin_ExpiresRequired(t *testing.T) { + t.Parallel() + svc := mcp.NewTokenService(newFakeRepo(), &fakeAudit{}) + + _, err := svc.IssueAdmin(context.Background(), mcp.IssueAdminInput{ + UserID: uuid.New(), + Name: "x", + CreatedBy: uuid.New(), + }) + ae, ok := errs.As(err) + require.True(t, ok) + assert.Equal(t, errs.CodeMcpTokenExpiresRequired, ae.Code) +} + +func TestIssueSystemToken_Defaults(t *testing.T) { + t.Parallel() + repo := newFakeRepo() + au := &fakeAudit{} + svc := mcp.NewTokenService(repo, au) + + uid := uuid.New() + sid := uuid.New() + pid := uuid.New() + res, err := svc.IssueSystemToken(context.Background(), mcp.IssueSystemTokenInput{ + UserID: uid, + ACPSessionID: sid, + ProjectIDs: []uuid.UUID{pid}, + }) + require.NoError(t, err) + + got, err := repo.GetByID(context.Background(), res.ID) + require.NoError(t, err) + assert.Equal(t, mcp.IssuerSystem, got.Issuer) + require.NotNil(t, got.ACPSessionID) + assert.Equal(t, sid, *got.ACPSessionID) + assert.NotNil(t, got.ExpiresAt) + assert.WithinDuration(t, time.Now().Add(24*time.Hour), *got.ExpiresAt, time.Minute) + // 默认工具白名单包含 PM 写入但不含 list_recent_messages + assert.Contains(t, got.Scope.Tools, "create_issue") + assert.NotContains(t, got.Scope.Tools, "list_recent_messages") + assert.Equal(t, []string{"mcp.token.create_system"}, au.actions()) +} + +func TestAuthenticate_HappyPath(t *testing.T) { + t.Parallel() + repo := newFakeRepo() + svc := mcp.NewTokenService(repo, &fakeAudit{}) + + res, err := svc.IssueAdmin(context.Background(), mcp.IssueAdminInput{ + UserID: uuid.New(), + Name: "x", + ExpiresAt: time.Now().Add(time.Hour), + CreatedBy: uuid.New(), + }) + require.NoError(t, err) + + tok, err := svc.Authenticate(context.Background(), res.Plaintext) + require.NoError(t, err) + assert.Equal(t, res.ID, tok.ID) +} + +func TestAuthenticate_Empty(t *testing.T) { + t.Parallel() + svc := mcp.NewTokenService(newFakeRepo(), &fakeAudit{}) + + _, err := svc.Authenticate(context.Background(), "") + ae, ok := errs.As(err) + require.True(t, ok) + assert.Equal(t, errs.CodeMcpTokenInvalid, ae.Code) +} + +func TestAuthenticate_Revoked(t *testing.T) { + t.Parallel() + repo := newFakeRepo() + svc := mcp.NewTokenService(repo, &fakeAudit{}) + + res, _ := svc.IssueAdmin(context.Background(), mcp.IssueAdminInput{ + UserID: uuid.New(), + Name: "x", + ExpiresAt: time.Now().Add(time.Hour), + CreatedBy: uuid.New(), + }) + require.NoError(t, svc.Revoke(context.Background(), res.ID, uuid.New())) + + _, err := svc.Authenticate(context.Background(), res.Plaintext) + ae, ok := errs.As(err) + require.True(t, ok) + assert.Equal(t, errs.CodeMcpTokenRevoked, ae.Code) +} + +func TestAuthenticate_Expired(t *testing.T) { + t.Parallel() + repo := newFakeRepo() + svc := mcp.NewTokenService(repo, &fakeAudit{}) + + res, _ := svc.IssueAdmin(context.Background(), mcp.IssueAdminInput{ + UserID: uuid.New(), + Name: "x", + ExpiresAt: time.Now().Add(time.Hour), + CreatedBy: uuid.New(), + }) + // 直接改 fake repo 中 expires_at 为已过去 + past := time.Now().Add(-time.Hour) + repo.tokens[res.ID].ExpiresAt = &past + + _, err := svc.Authenticate(context.Background(), res.Plaintext) + ae, ok := errs.As(err) + require.True(t, ok) + assert.Equal(t, errs.CodeMcpTokenExpired, ae.Code) +} + +func TestAuthenticate_TamperedToken(t *testing.T) { + t.Parallel() + svc := mcp.NewTokenService(newFakeRepo(), &fakeAudit{}) + + _, err := svc.Authenticate(context.Background(), "mcpt_tampered_value") + ae, ok := errs.As(err) + require.True(t, ok) + assert.Equal(t, errs.CodeMcpTokenInvalid, ae.Code) +} + +func TestRevokeBySession_AuditPerToken(t *testing.T) { + t.Parallel() + repo := newFakeRepo() + au := &fakeAudit{} + svc := mcp.NewTokenService(repo, au) + + uid := uuid.New() + sid := uuid.New() + pid := uuid.New() + for i := 0; i < 3; i++ { + _, err := svc.IssueSystemToken(context.Background(), mcp.IssueSystemTokenInput{ + UserID: uid, ACPSessionID: sid, ProjectIDs: []uuid.UUID{pid}, + }) + require.NoError(t, err) + } + + require.NoError(t, svc.RevokeBySession(context.Background(), sid)) + + // 3 个 system 创建 + 3 个 by_session 撤销 = 6 条 audit + actions := au.actions() + assert.Len(t, actions, 6) + createCnt, revokeCnt := 0, 0 + for _, a := range actions { + switch a { + case "mcp.token.create_system": + createCnt++ + case "mcp.token.revoke_by_session": + revokeCnt++ + } + } + assert.Equal(t, 3, createCnt) + assert.Equal(t, 3, revokeCnt) +} + +// 编译期保证:类型实现接口 +var _ = errors.New