test(mcp): auth middleware + rate limiter tests

2026-05-08 08:11:02 +08:00
parent e1e6508c4c
commit c046d77a6e
2 changed files with 204 additions and 0 deletions
@@ -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)
+}
@@ -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())
+}