internal/middleware/oauth_middleware_test.go
Rama main · réplica de lectura del repositorio original.
package middleware
import (
"context"
"errors"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
"xprem/internal/crypto"
"xprem/internal/oauth"
"xprem/internal/repository"
"xprem/internal/services"
"github.com/golang-jwt/jwt/v5"
"github.com/modelcontextprotocol/go-sdk/auth"
)
type fakeOAuthUserRepo struct {
user repository.User
err error
}
func (f *fakeOAuthUserRepo) GetUserByID(_ context.Context, id string) (repository.User, error) {
if f.err != nil {
return repository.User{}, f.err
}
if f.user.Id != id {
return repository.User{}, &repository.ErrResourceNotFound{Resource: "user", Identifier: id}
}
return f.user, nil
}
func (f *fakeOAuthUserRepo) BumpUserSessionVersion(_ context.Context, _ string) error {
return nil
}
// principalCapture records what the middleware put in the context of the
// request that reached the wrapped handler.
type principalCapture struct {
principal *services.DashboardPrincipal
}
func oauthTestSetup(t *testing.T, repo *fakeOAuthUserRepo) (*oauth.OAuthService, http.Handler, *principalCapture) {
t.Helper()
t.Setenv("BASE_URL", "https://ota.example.com")
t.Setenv("JWT_SECRET", "test-secret")
service := oauth.NewOAuthService(nil, nil, nil, repo)
capture := &principalCapture{}
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
capture.principal = services.PrincipalFromContext(r.Context())
w.WriteHeader(http.StatusOK)
})
return service, NewOAuthMiddleware(service)(next), capture
}
func doMCPRequest(handler http.Handler, token string) *httptest.ResponseRecorder {
req := httptest.NewRequest(http.MethodPost, "/mcp", nil)
if token != "" {
req.Header.Set("Authorization", "Bearer "+token)
}
res := httptest.NewRecorder()
handler.ServeHTTP(res, req)
return res
}
func enabledUser() repository.User {
return repository.User{Id: "user-1", Email: "a@b.c", IsAdmin: false, Enabled: true, SessionVersion: 3}
}
func TestOAuthMiddlewareMissingToken(t *testing.T) {
_, handler, _ := oauthTestSetup(t, &fakeOAuthUserRepo{user: enabledUser()})
res := doMCPRequest(handler, "")
if res.Code != http.StatusUnauthorized {
t.Fatalf("expected 401, got %d", res.Code)
}
challenge := res.Header().Get("WWW-Authenticate")
if !strings.Contains(challenge, `resource_metadata="https://ota.example.com/.well-known/oauth-protected-resource/mcp"`) {
t.Errorf("WWW-Authenticate must point at the resource metadata, got %q", challenge)
}
}
func TestOAuthMiddlewareValidToken(t *testing.T) {
repo := &fakeOAuthUserRepo{user: enabledUser()}
service, handler, capture := oauthTestSetup(t, repo)
token, err := service.IssueAccessToken(services.DashboardPrincipal{UserId: "user-1", SessionVersion: 3})
if err != nil {
t.Fatal(err)
}
res := doMCPRequest(handler, token)
if res.Code != http.StatusOK {
t.Fatalf("expected 200, got %d: %s", res.Code, res.Body.String())
}
if capture.principal == nil {
t.Fatal("no principal in the request context")
}
// The principal must come from the user row, not from the claims.
if capture.principal.Email != "a@b.c" || capture.principal.UserId != "user-1" {
t.Errorf("unexpected principal: %+v", capture.principal)
}
}
func TestOAuthMiddlewareRejections(t *testing.T) {
repo := &fakeOAuthUserRepo{user: enabledUser()}
service, handler, _ := oauthTestSetup(t, repo)
dashboardToken, _ := crypto.GenerateJWTToken("test-secret", jwt.MapClaims{
"sub": "admin-dashboard", "type": "token", "userId": "user-1", "sv": 3,
"exp": time.Now().Add(time.Hour).Unix(),
})
expiredToken, _ := crypto.GenerateJWTToken("test-secret", jwt.MapClaims{
"sub": "mcp", "aud": "https://ota.example.com/mcp", "userId": "user-1", "sv": 3,
"exp": time.Now().Add(-time.Minute).Unix(),
})
wrongAudience, _ := crypto.GenerateJWTToken("test-secret", jwt.MapClaims{
"sub": "mcp", "aud": "https://other.example.com/mcp", "userId": "user-1", "sv": 3,
"exp": time.Now().Add(time.Hour).Unix(),
})
staleSessionVersion, err := service.IssueAccessToken(services.DashboardPrincipal{UserId: "user-1", SessionVersion: 2})
if err != nil {
t.Fatal(err)
}
unknownUser, err := service.IssueAccessToken(services.DashboardPrincipal{UserId: "user-2", SessionVersion: 3})
if err != nil {
t.Fatal(err)
}
cases := map[string]string{
"dashboard token": dashboardToken,
"expired token": expiredToken,
"wrong audience": wrongAudience,
"stale session version": staleSessionVersion,
"unknown user": unknownUser,
}
for name, token := range cases {
t.Run(name, func(t *testing.T) {
res := doMCPRequest(handler, token)
if res.Code != http.StatusUnauthorized {
t.Fatalf("expected 401, got %d", res.Code)
}
if challenge := res.Header().Get("WWW-Authenticate"); !strings.Contains(challenge, "resource_metadata") {
t.Errorf("the 401 must carry the discovery challenge, got %q", challenge)
}
})
}
}
func TestOAuthMiddlewareDisabledUser(t *testing.T) {
user := enabledUser()
user.Enabled = false
repo := &fakeOAuthUserRepo{user: user}
service, handler, _ := oauthTestSetup(t, repo)
token, err := service.IssueAccessToken(services.DashboardPrincipal{UserId: "user-1", SessionVersion: 3})
if err != nil {
t.Fatal(err)
}
if res := doMCPRequest(handler, token); res.Code != http.StatusUnauthorized {
t.Fatalf("expected 401, got %d", res.Code)
}
}
func TestOAuthMiddlewareDatabaseOutage(t *testing.T) {
repo := &fakeOAuthUserRepo{err: errors.New("connection refused")}
service, handler, _ := oauthTestSetup(t, repo)
token, err := service.IssueAccessToken(services.DashboardPrincipal{UserId: "user-1", SessionVersion: 3})
if err != nil {
t.Fatal(err)
}
res := doMCPRequest(handler, token)
if res.Code != http.StatusInternalServerError {
t.Fatalf("an outage must be a 500, not a dead token: got %d", res.Code)
}
if res.Header().Get("WWW-Authenticate") != "" {
t.Error("a 500 must not carry WWW-Authenticate, the client would restart the flow")
}
if strings.Contains(res.Body.String(), "connection refused") {
t.Error("the response body must not leak the underlying database error")
}
}
func TestOAuthMiddlewareBindsSessionUser(t *testing.T) {
repo := &fakeOAuthUserRepo{user: enabledUser()}
service, _, _ := oauthTestSetup(t, repo)
token, err := service.IssueAccessToken(services.DashboardPrincipal{UserId: "user-1", SessionVersion: 3})
if err != nil {
t.Fatal(err)
}
// The streamable transport pins sessions to TokenInfo.UserID; assert the
// middleware actually plants it where the SDK reads it.
var seenInfo *auth.TokenInfo
handler := NewOAuthMiddleware(service)(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
seenInfo = auth.TokenInfoFromContext(r.Context())
}))
doMCPRequest(handler, token)
if seenInfo == nil || seenInfo.UserID != "user-1" {
t.Fatalf("TokenInfo.UserID must reach the transport, got %+v", seenInfo)
}
}