@@ -110,7 +110,8 @@ func runSseServer(configOpts *ServerOptions) error {
110110 TLSKeyFile : snConfig .ExternalPulsar .TLSKeyFile ,
111111 },
112112 }
113- pulsarSessionManager = session .NewPulsarSessionManager (managerConfig , mcpServer .PulsarSession , logger )
113+ // Pass nil as globalSession - in multi-session mode, every request must have a valid token
114+ pulsarSessionManager = session .NewPulsarSessionManager (managerConfig , nil , logger )
114115 logger .Info ("Multi-session Pulsar mode enabled" )
115116 }
116117
@@ -125,14 +126,16 @@ func runSseServer(configOpts *ServerOptions) error {
125126 // Handle per-user Pulsar sessions
126127 if pulsarSessionManager != nil {
127128 token := session .ExtractBearerToken (r )
129+ // Token is already validated in auth middleware, this should always succeed
128130 if pulsarSession , err := pulsarSessionManager .GetOrCreateSession (ctx , token ); err == nil {
129131 c = context2 .WithPulsarSession (c , pulsarSession )
130132 if token != "" {
131133 c = session .WithUserTokenHash (c , pulsarSessionManager .HashTokenForLog (token ))
132134 }
133135 } else {
134- logger .WithError (err ).Warn ("Failed to get per-user Pulsar session, using global" )
135- c = context2 .WithPulsarSession (c , mcpServer .PulsarSession )
136+ // Should not happen since middleware validates token first
137+ logger .WithError (err ).Error ("Unexpected auth error after middleware validation" )
138+ // Don't set PulsarSession - tool handlers will fail gracefully with "session not found"
136139 }
137140 } else {
138141 c = context2 .WithPulsarSession (c , mcpServer .PulsarSession )
@@ -144,16 +147,62 @@ func runSseServer(configOpts *ServerOptions) error {
144147
145148 // 4. Expose the full SSE URL to the user
146149 ssePath := sseServer .CompleteSsePath ()
150+ msgPath := sseServer .CompleteMessagePath ()
147151 fmt .Fprintf (os .Stderr , "StreamNative Cloud MCP Server listening on http://%s%s\n " ,
148152 configOpts .HTTPAddr , ssePath )
149153
150154 // 5. Run the HTTP listener in a goroutine
151155 errCh := make (chan error , 1 )
152- go func () {
153- if err := sseServer .Start (configOpts .HTTPAddr ); err != nil && ! errors .Is (err , http .ErrServerClosed ) {
154- errCh <- err // bubble up real crashes
156+ var httpServer * http.Server
157+
158+ if pulsarSessionManager != nil {
159+ // Multi-session mode: use custom handlers with auth middleware
160+ mux := http .NewServeMux ()
161+
162+ // Auth middleware wrapper that validates token before processing
163+ authMiddleware := func (next http.Handler ) http.Handler {
164+ return http .HandlerFunc (func (w http.ResponseWriter , r * http.Request ) {
165+ token := session .ExtractBearerToken (r )
166+ if token == "" {
167+ w .Header ().Set ("Content-Type" , "application/json" )
168+ http .Error (w , `{"error":"missing Authorization header"}` , http .StatusUnauthorized )
169+ return
170+ }
171+ // Pre-validate token by attempting to get/create session
172+ if _ , err := pulsarSessionManager .GetOrCreateSession (r .Context (), token ); err != nil {
173+ logger .WithError (err ).Warn ("Authentication failed" )
174+ w .Header ().Set ("Content-Type" , "application/json" )
175+ http .Error (w , `{"error":"authentication failed"}` , http .StatusUnauthorized )
176+ return
177+ }
178+ next .ServeHTTP (w , r )
179+ })
155180 }
156- }()
181+
182+ // Mount handlers with auth middleware
183+ mux .Handle (ssePath , authMiddleware (sseServer .SSEHandler ()))
184+ mux .Handle (msgPath , authMiddleware (sseServer .MessageHandler ()))
185+
186+ // Start custom HTTP server
187+ httpServer = & http.Server {
188+ Addr : configOpts .HTTPAddr ,
189+ Handler : mux ,
190+ ReadHeaderTimeout : 10 * time .Second , // Prevent Slowloris attacks
191+ }
192+ go func () {
193+ if err := httpServer .ListenAndServe (); err != nil && ! errors .Is (err , http .ErrServerClosed ) {
194+ errCh <- err
195+ }
196+ }()
197+ logger .Info ("SSE server started with authentication middleware" )
198+ } else {
199+ // Non-multi-session mode: use default Start()
200+ go func () {
201+ if err := sseServer .Start (configOpts .HTTPAddr ); err != nil && ! errors .Is (err , http .ErrServerClosed ) {
202+ errCh <- err // bubble up real crashes
203+ }
204+ }()
205+ }
157206
158207 // Give the server a moment to start
159208 time .Sleep (100 * time .Millisecond )
@@ -177,10 +226,20 @@ func runSseServer(configOpts *ServerOptions) error {
177226 pulsarSessionManager .Stop ()
178227 }
179228
180- // First try to shut down the SSE server
181- if err := sseServer .Shutdown (shCtx ); err != nil {
182- if ! errors .Is (err , http .ErrServerClosed ) {
183- logger .Errorf ("Error shutting down SSE server: %v" , err )
229+ // Shut down the HTTP server
230+ if httpServer != nil {
231+ // Multi-session mode: shut down custom HTTP server
232+ if err := httpServer .Shutdown (shCtx ); err != nil {
233+ if ! errors .Is (err , http .ErrServerClosed ) {
234+ logger .Errorf ("Error shutting down HTTP server: %v" , err )
235+ }
236+ }
237+ } else {
238+ // Non-multi-session mode: shut down SSE server
239+ if err := sseServer .Shutdown (shCtx ); err != nil {
240+ if ! errors .Is (err , http .ErrServerClosed ) {
241+ logger .Errorf ("Error shutting down SSE server: %v" , err )
242+ }
184243 }
185244 }
186245
0 commit comments