fix: proper support for chat session cors (#1109)

The previous approach was missing support for OPTIONS requests, causing
the preflight to break. It was impractical to add support for that due
to the location the middleware was attached (after the actual CORS
middleware and after routes were matched, thus excluding OPTIONS), so
instead I went back to the original approach of integrating with the
existing CORS middleware

Author: chase-crumbaugh

server/cmd/gram/start.go

@@ -539,7 +539,7 @@ func newStartCommand() *cli.Command {
 			ragService := rag.NewToolsetVectorStore(logger, tracerProvider, db, baseChatClient)
 
 			mux := goahttp.NewMuxer()
-			mux.Use(middleware.CORSMiddleware(c.String("environment"), c.String("server-url")))
+			mux.Use(middleware.CORSMiddleware(c.String("environment"), c.String("server-url"), chatSessionsManager))
 			mux.Use(middleware.NewHTTPLoggingMiddleware(logger))
 			mux.Use(customdomains.Middleware(logger, db, c.String("environment"), serverURL))
 			mux.Use(middleware.SessionMiddleware)

server/internal/chat/impl.go

@@ -90,14 +90,10 @@ func Attach(mux goahttp.Muxer, service *Service) {
 	endpoints.Use(middleware.MapErrors())
 	endpoints.Use(middleware.TraceMethods(service.tracer))
 
-	chatSessionMiddleware := middleware.ChatSessionMiddleware(service.chatSessions)
-
 	server := srv.New(endpoints, mux, goahttp.RequestDecoder, goahttp.ResponseEncoder, nil, nil)
-	server.Use(chatSessionMiddleware)
 	srv.Mount(mux, server)
 
-	wrappedHandler := chatSessionMiddleware(oops.ErrHandle(service.logger, service.HandleCompletion))
-	o11y.AttachHandler(mux, "POST", "/chat/completions", wrappedHandler.ServeHTTP)
+	o11y.AttachHandler(mux, "POST", "/chat/completions", oops.ErrHandle(service.logger, service.HandleCompletion).ServeHTTP)
 }
 
 func (s *Service) APIKeyAuth(ctx context.Context, key string, schema *security.APIKeyScheme) (context.Context, error) {
@@ -194,7 +190,7 @@ func (s *Service) ListChats(ctx context.Context, payload *gen.ListChatsPayload)
 	} else {
 		chats, err := s.repo.ListChatsForUser(ctx, repo.ListChatsForUserParams{
 			ProjectID: *authCtx.ProjectID,
-			UserID:    conv.ToPGText(authCtx.UserID),
+			UserID:    conv.ToPGText(authCtx.UserID), // TODO: make this work for external user ids (Elements)
 		})
 		if err != nil {
 			return nil, oops.E(oops.CodeUnexpected, err, "failed to list chats").Log(ctx, s.logger)

server/internal/instances/impl.go

@@ -124,7 +124,6 @@ func Attach(mux goahttp.Muxer, service *Service) {
 	endpoints.Use(middleware.TraceMethods(service.tracer))
 
 	server := srv.New(endpoints, mux, goahttp.RequestDecoder, goahttp.ResponseEncoder, nil, nil)
-	server.Use(middleware.ChatSessionMiddleware(service.chatSessions))
 	srv.Mount(mux, server)
 
 	o11y.AttachHandler(mux, "POST", "/rpc/instances.invoke/tool", func(w http.ResponseWriter, r *http.Request) {

server/internal/mcp/impl.go

@@ -20,7 +20,6 @@ import (
 	"github.com/go-chi/chi/v5"
 	"github.com/google/uuid"
 	"github.com/jackc/pgx/v5/pgxpool"
-	"github.com/speakeasy-api/gram/server/internal/middleware"
 	"github.com/speakeasy-api/gram/server/internal/rag"
 	tm "github.com/speakeasy-api/gram/server/internal/telemetry"
 	"go.opentelemetry.io/otel/metric"
@@ -175,21 +174,12 @@ func NewService(
 }
 
 func Attach(mux goahttp.Muxer, service *Service, metadataService *mcpmetadata.Service) {
-	chatSessionMiddleware := middleware.ChatSessionMiddleware(service.chatSessionsManager)
-
-	// Wraps handler functions with chat session middleware
-	withMiddleware := func(handler http.Handler) http.Handler {
-		return chatSessionMiddleware(handler)
-	}
-
-	handler := withMiddleware(oops.ErrHandle(service.logger, service.ServePublic))
-	o11y.AttachHandler(mux, "POST", "/mcp/{mcpSlug}", handler.ServeHTTP)
-
-	o11y.AttachHandler(mux, "GET", "/mcp/{mcpSlug}", withMiddleware(oops.ErrHandle(service.logger, func(w http.ResponseWriter, r *http.Request) error {
+	o11y.AttachHandler(mux, "POST", "/mcp/{mcpSlug}", oops.ErrHandle(service.logger, service.ServePublic).ServeHTTP)
+	o11y.AttachHandler(mux, "GET", "/mcp/{mcpSlug}", oops.ErrHandle(service.logger, func(w http.ResponseWriter, r *http.Request) error {
 		return service.HandleGetServer(w, r, metadataService)
-	})).ServeHTTP)
-	o11y.AttachHandler(mux, "GET", "/mcp/{mcpSlug}/install", withMiddleware(oops.ErrHandle(service.logger, metadataService.ServeInstallPage)).ServeHTTP)
-	o11y.AttachHandler(mux, "POST", "/mcp/{project}/{toolset}/{environment}", withMiddleware(oops.ErrHandle(service.logger, service.ServeAuthenticated)).ServeHTTP)
+	}).ServeHTTP)
+	o11y.AttachHandler(mux, "GET", "/mcp/{mcpSlug}/install", oops.ErrHandle(service.logger, metadataService.ServeInstallPage).ServeHTTP)
+	o11y.AttachHandler(mux, "POST", "/mcp/{project}/{toolset}/{environment}", oops.ErrHandle(service.logger, service.ServeAuthenticated).ServeHTTP)
 
 	// OAuth 2.1 Authorization Server Metadata
 	o11y.AttachHandler(mux, "GET", "/.well-known/oauth-authorization-server/mcp/{mcpSlug}", oops.ErrHandle(service.logger, service.HandleWellKnownOAuthServerMetadata).ServeHTTP)

server/internal/middleware/chat_session.go

@@ -0,0 +1,54 @@
+package middleware
+
+import (
+	"fmt"
+	"net/http"
+	"slices"
+	"strings"
+
+	"github.com/speakeasy-api/gram/server/internal/auth/chatsessions"
+	"github.com/speakeasy-api/gram/server/internal/constants"
+)
+
+var chatSessionsAllowedRoutes = []string{
+	"/chat/completions",
+	"/mcp",
+	// "/rpc/chat", // TODO: Support listing / creating chats for elements
+}
+
+// This isn't practical to do as a proper middleware because it needs to interoperate with the CORSMiddleware which does things like returning early for OPTIONS requests.
+// Instead, we combine it with the CORSMiddleware so that all CORS stuff is handled in one place.
+func chatSessionsCORS(chatSessionsManager *chatsessions.Manager) func(next http.Handler) http.Handler {
+	return func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if r.Method == http.MethodOptions {
+				// Slightly non-ideal, but later in the file we validate the origin of the request against the audience claim
+				w.Header().Set("Access-Control-Allow-Origin", r.Header.Get("Origin")) // Allow the origin of the request for OPTIONS requests because we don't know what origins to allow until we get the token on the actual request
+				w.WriteHeader(http.StatusNoContent)
+				return
+			}
+
+			chatSession := r.Header.Get(constants.ChatSessionsTokenHeader)
+			if chatSession == "" {
+				next.ServeHTTP(w, r)
+				return
+			}
+
+			claims, err := chatSessionsManager.ValidateToken(r.Context(), chatSession)
+			if err != nil {
+				http.Error(w, "unauthorized", http.StatusUnauthorized)
+				return
+			}
+
+			// If the request origin is in the allowed origins, set the allowed origin in the context to be used in the CORS middleware
+			if slices.Contains(claims.Audience, r.Header.Get("Origin")) {
+				w.Header().Set("Access-Control-Allow-Origin", r.Header.Get("Origin"))
+			} else {
+				http.Error(w, fmt.Sprintf("Origin %s does not match audience claim: %s", r.Header.Get("Origin"), strings.Join(claims.Audience, ", ")), http.StatusForbidden)
+				return
+			}
+
+			next.ServeHTTP(w, r)
+		})
+	}
+}

server/internal/middleware/chat_session_cors.go

@@ -5,14 +5,16 @@ import (
 	"net/url"
 	"slices"
 	"strings"
+
+	"github.com/speakeasy-api/gram/server/internal/auth/chatsessions"
 )
 
 var mcpOpenAccessControlRoutes = []string{
 	"/.well-known/oauth-authorization-server/mcp",
 	"/.well-known/oauth-protected-resource/mcp",
 }
 
-func CORSMiddleware(env string, serverURL string) func(next http.Handler) http.Handler {
+func CORSMiddleware(env string, serverURL string, chatSessionsManager *chatsessions.Manager) func(next http.Handler) http.Handler {
 	return func(next http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			switch env {
@@ -29,8 +31,6 @@ func CORSMiddleware(env string, serverURL string) func(next http.Handler) http.H
 				// No CORS headers set for unspecified environments
 			}
 
-			// NOTE: The chatSession middleware may also set the Access-Control-Allow-Origin header for chat sessions.
-
 			w.Header().Set("Access-Control-Allow-Methods", "POST, GET, OPTIONS, PUT, DELETE")
 			w.Header().Set("Access-Control-Allow-Headers", "Accept, Content-Type, Content-Length, Accept-Encoding, Authorization, User-Agent, Gram-Session, Gram-Project, Gram-Token, idempotency-key, Gram-Admin-Override, Gram-Chat-ID, Gram-Chat-Session, MCP-Protocol-Version")
 			w.Header().Set("Access-Control-Expose-Headers", "Accept, Content-Type, Content-Length, Accept-Encoding, x-trace-id, Gram-Session, Gram-Chat-ID, Gram-Chat-Session")
@@ -46,6 +46,14 @@ func CORSMiddleware(env string, serverURL string) func(next http.Handler) http.H
 				w.Header().Del("Access-Control-Allow-Credentials")
 			}
 
+			// Special CORS handling for chat sessions-enabled routes
+			if slices.ContainsFunc(chatSessionsAllowedRoutes, func(route string) bool {
+				return strings.HasPrefix(r.URL.Path, route)
+			}) {
+				chatSessionsCORS(chatSessionsManager)(next).ServeHTTP(w, r)
+				return
+			}
+
 			if r.Method == "OPTIONS" {
 				w.WriteHeader(http.StatusOK)
 				return