diff --git a/internal/mcp/auth_middleware_test.go b/internal/mcp/auth_middleware_test.go new file mode 100644 index 0000000..e93a8d7 --- /dev/null +++ b/internal/mcp/auth_middleware_test.go @@ -0,0 +1,102 @@ +package mcp_test + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/yan1h/agent-coding-workflow/internal/mcp" +) + +func TestAuthMiddleware_BearerHeader(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) + + mw := mcp.NewAuthMiddleware(svc) + h := mw(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + s, ok := mcp.FromContext(r.Context()) + require.True(t, ok) + require.NotNil(t, s) + _, _ = w.Write([]byte(s.TokenID)) + })) + + req := httptest.NewRequest("GET", "/mcp", nil) + req.Header.Set("Authorization", "Bearer "+res.Plaintext) + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + assert.Equal(t, http.StatusOK, rec.Code) + body, _ := io.ReadAll(rec.Body) + assert.Equal(t, res.ID.String(), string(body)) +} + +func TestAuthMiddleware_QueryFallback(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(), + }) + + mw := mcp.NewAuthMiddleware(svc) + h := mw(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + + req := httptest.NewRequest("GET", "/mcp?token="+res.Plaintext, nil) + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + assert.Equal(t, http.StatusOK, rec.Code) +} + +func TestAuthMiddleware_MissingToken(t *testing.T) { + t.Parallel() + svc := mcp.NewTokenService(newFakeRepo(), &fakeAudit{}) + + mw := mcp.NewAuthMiddleware(svc) + h := mw(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + t.Fatal("should not call inner handler") + })) + + req := httptest.NewRequest("GET", "/mcp", nil) + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + assert.Equal(t, http.StatusUnauthorized, rec.Code) +} + +func TestAuthMiddleware_InvalidToken(t *testing.T) { + t.Parallel() + svc := mcp.NewTokenService(newFakeRepo(), &fakeAudit{}) + + mw := mcp.NewAuthMiddleware(svc) + h := mw(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + t.Fatal("should not call") + })) + + req := httptest.NewRequest("GET", "/mcp", nil) + req.Header.Set("Authorization", "Bearer mcpt_invalid") + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + assert.Equal(t, http.StatusUnauthorized, rec.Code) +} diff --git a/internal/mcp/ratelimit_test.go b/internal/mcp/ratelimit_test.go new file mode 100644 index 0000000..99a0f3f --- /dev/null +++ b/internal/mcp/ratelimit_test.go @@ -0,0 +1,102 @@ +package mcp_test + +import ( + "context" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + + "github.com/yan1h/agent-coding-workflow/internal/mcp" +) + +func TestRateLimiter_AllowUntilCap(t *testing.T) { + t.Parallel() + now := time.Date(2026, 5, 7, 12, 0, 0, 0, time.UTC) + limiter := mcp.NewRateLimiter(mcp.RateLimitConfig{ + PerTokenPerMinute: 3, + GlobalPerMinute: 100, + }, func() time.Time { return now }) + + for i := 0; i < 3; i++ { + ok, _ := limiter.Allow("token-A") + assert.True(t, ok, "request %d should pass", i) + } + ok, reason := limiter.Allow("token-A") + assert.False(t, ok) + assert.Equal(t, "per_token", reason) +} + +func TestRateLimiter_GlobalCap(t *testing.T) { + t.Parallel() + now := time.Date(2026, 5, 7, 12, 0, 0, 0, time.UTC) + limiter := mcp.NewRateLimiter(mcp.RateLimitConfig{ + PerTokenPerMinute: 100, + GlobalPerMinute: 2, + }, func() time.Time { return now }) + + ok, _ := limiter.Allow("token-A") + assert.True(t, ok) + ok, _ = limiter.Allow("token-B") + assert.True(t, ok) + ok, reason := limiter.Allow("token-C") + assert.False(t, ok) + assert.Equal(t, "global", reason) +} + +func TestRateLimiter_WindowReset(t *testing.T) { + t.Parallel() + current := time.Date(2026, 5, 7, 12, 0, 0, 0, time.UTC) + limiter := mcp.NewRateLimiter(mcp.RateLimitConfig{ + PerTokenPerMinute: 1, + GlobalPerMinute: 100, + }, func() time.Time { return current }) + + ok, _ := limiter.Allow("A") + assert.True(t, ok) + ok, _ = limiter.Allow("A") + assert.False(t, ok) + + // 跨分钟窗口 + current = current.Add(time.Minute) + ok, _ = limiter.Allow("A") + assert.True(t, ok, "after window reset should allow again") +} + +func TestRateLimitMiddleware_429(t *testing.T) { + t.Parallel() + now := time.Date(2026, 5, 7, 12, 0, 0, 0, time.UTC) + limiter := mcp.NewRateLimiter(mcp.RateLimitConfig{ + PerTokenPerMinute: 1, + GlobalPerMinute: 100, + }, func() time.Time { return now }) + + // 内层 handler 计数被调次数 + var hits atomic.Int32 + innerH := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + hits.Add(1) + w.WriteHeader(http.StatusOK) + }) + + // 模拟 auth 中间件已注入 ctx:直接 context.WithValue + mcp.AuthContextKey + withSession := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + ctx := context.WithValue(r.Context(), mcp.AuthContextKey, &mcp.AuthSession{ + TokenID: "fake-token-id", + UserID: "fake-user", + }) + mcp.NewRateLimitMiddleware(limiter)(innerH).ServeHTTP(w, r.WithContext(ctx)) + }) + + req := httptest.NewRequest("GET", "/mcp", nil) + rec := httptest.NewRecorder() + withSession.ServeHTTP(rec, req) + assert.Equal(t, http.StatusOK, rec.Code) + + rec2 := httptest.NewRecorder() + withSession.ServeHTTP(rec2, req) + assert.Equal(t, http.StatusTooManyRequests, rec2.Code) + assert.Equal(t, int32(1), hits.Load()) +}