fix(router): preserve SSE response writer compatibility

Author: endigma

docs-website/router/custom-modules.mdx

@@ -338,6 +338,29 @@ func (m *CustomModule) RouterOnRequest(ctx core.RequestContext, next http.Handle
 }
 ```
 
+#### Streaming responses
+
+Custom response writers used for subscriptions must implement `http.Flusher`.
+
+To retain SSE write timeouts, they must also either implement `SetWriteDeadline(time.Time) error` **or** expose the wrapped response writer:
+
+```go
+// required for streaming to function
+func (w *headerCapturingWriter) Flush() {
+	w.ResponseWriter.(http.Flusher).Flush()
+}
+
+// Forward write deadlines directly
+func (w *headerCapturingWriter) SetWriteDeadline(deadline time.Time) error {
+	return http.NewResponseController(w.ResponseWriter).SetWriteDeadline(deadline)
+}
+
+// or expose the wrapped writer.
+func (w *headerCapturingWriter) Unwrap() http.ResponseWriter {
+	return w.ResponseWriter
+}
+```
+
 ### Request Handler lifecycle
 
 The current module handler allow to intercept and modify request / response subgraphs.

router-tests/subscriptions/http_subscriptions_test.go

@@ -614,7 +614,9 @@ func testSSENonFlusherWriter(t *testing.T) {
 			require.Equal(
 				t,
 				`event: next
-data: {"errors":[{"message":"subscription response writer does not support flushing"}]}`,
+data: {"errors":[{"message":"subscription response writer does not support flushing"}]}
+
+`,
 				string(body),
 			)
 		})

router/core/errors.go

@@ -299,6 +299,17 @@ func writeRequestErrors(params writeRequestErrorsParams) {
 			}
 			params.logger.Error("Error writing response", zap.Error(err))
 		}
+		return
+	}
+
+	if wgRequestParams.UseSse {
+		if _, err := params.writer.Write([]byte("\n\n")); err != nil && params.logger != nil {
+			if rErrors.IsBrokenPipe(err) {
+				params.logger.Warn("Broken pipe, error writing response", zap.Error(err))
+				return
+			}
+			params.logger.Error("Error writing response", zap.Error(err))
+		}
 	}
 }
 

router/core/graphql_handler.go

@@ -330,6 +330,7 @@ func (h *GraphQLHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		resolveCtx, writer, writerErr = GetSubscriptionResponseWriter(resolveCtx, r, w, SubscriptionResponseWriterOptions{
 			ApolloSubscriptionMultipartPrintBoundary: h.apolloSubscriptionMultipartPrintBoundary,
 			SSEWriteTimeout:                          h.sseServerWriteTimeout,
+			Logger:                                   reqCtx.logger,
 		})
 		if writerErr != nil {
 			reqCtx.logger.Error("unable to get subscription response writer", zap.Error(writerErr))

router/core/subscription_response_writer.go

@@ -14,6 +14,7 @@ import (
 
 	"github.com/wundergraph/astjson"
 	"github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve"
+	"go.uber.org/zap"
 )
 
 const (
@@ -37,6 +38,7 @@ type withFlushWriter interface {
 type SubscriptionResponseWriterOptions struct {
 	ApolloSubscriptionMultipartPrintBoundary bool
 	SSEWriteTimeout                          time.Duration
+	Logger                                   *zap.Logger
 }
 
 type HttpFlushWriter struct {
@@ -51,6 +53,7 @@ type HttpFlushWriter struct {
 	buf             *bytes.Buffer
 	firstMessage    bool
 	sseWriteTimeout time.Duration
+	logger          *zap.Logger
 	// apolloSubscriptionMultipartPrintBoundary if set to true will send the multipart boundary at the end of the message to allow
 	// misbehaving client (like apollo client) to read the message just sent before the next one or the heartbeat
 	apolloSubscriptionMultipartPrintBoundary bool
@@ -189,16 +192,25 @@ func (f *HttpFlushWriter) Flush() (err error) {
 
 func (f *HttpFlushWriter) writeAndFlushSSE(write func() error) (err error) {
 	if f.sseWriteTimeout > 0 {
-		if err := f.responseControl.SetWriteDeadline(time.Now().Add(f.sseWriteTimeout)); err != nil {
-			// Failing closed prevents a response writer without deadline support from
-			// reintroducing an unbounded shared-trigger stall.
-			return fmt.Errorf("set SSE write deadline: %w", err)
-		}
-		defer func() {
-			if clearErr := f.responseControl.SetWriteDeadline(time.Time{}); clearErr != nil {
-				err = errors.Join(err, fmt.Errorf("clear SSE write deadline: %w", clearErr))
+		if deadlineErr := f.responseControl.SetWriteDeadline(time.Now().Add(f.sseWriteTimeout)); deadlineErr != nil {
+			if !errors.Is(deadlineErr, http.ErrNotSupported) {
+				return fmt.Errorf("set SSE write deadline: %w", deadlineErr)
+			}
+
+			f.sseWriteTimeout = 0
+			if f.logger != nil {
+				f.logger.Warn(
+					"SSE write timeout disabled because response writer does not support write deadlines",
+					zap.Error(deadlineErr),
+				)
 			}
-		}()
+		} else {
+			defer func() {
+				if clearErr := f.responseControl.SetWriteDeadline(time.Time{}); clearErr != nil {
+					err = errors.Join(err, fmt.Errorf("clear SSE write deadline: %w", clearErr))
+				}
+			}()
+		}
 	}
 
 	if err := write(); err != nil {
@@ -231,6 +243,7 @@ func GetSubscriptionResponseWriter(ctx *resolve.Context, r *http.Request, w http
 		buf:                                      &bytes.Buffer{},
 		firstMessage:                             true,
 		sseWriteTimeout:                          opts.SSEWriteTimeout,
+		logger:                                   opts.Logger,
 		apolloSubscriptionMultipartPrintBoundary: opts.ApolloSubscriptionMultipartPrintBoundary,
 	}
 

router/core/subscription_response_writer_test.go

@@ -12,6 +12,8 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve"
+	"go.uber.org/zap"
+	"go.uber.org/zap/zaptest/observer"
 )
 
 type deadlineRecorder struct {
@@ -244,16 +246,21 @@ func TestGetSubscriptionResponseWriter(t *testing.T) {
 		assert.Nil(t, writer)
 	})
 
-	t.Run("fails closed when an SSE deadline is configured but unsupported", func(t *testing.T) {
+	t.Run("disables the SSE timeout when write deadlines are unsupported", func(t *testing.T) {
 		recorder := httptest.NewRecorder()
 		req := httptest.NewRequest(http.MethodPost, "/graphql", nil)
 		req.Header.Set("Accept", sseMimeType)
+		logCore, logs := observer.New(zap.WarnLevel)
 
-		_, writer, err := GetSubscriptionResponseWriter(resolve.NewContext(context.Background()), req, recorder, SubscriptionResponseWriterOptions{SSEWriteTimeout: time.Second})
-		require.Error(t, err)
-		assert.ErrorIs(t, err, http.ErrNotSupported)
-		assert.ErrorContains(t, err, "set SSE write deadline")
-		assert.Nil(t, writer)
+		_, writer, err := GetSubscriptionResponseWriter(resolve.NewContext(context.Background()), req, recorder, SubscriptionResponseWriterOptions{
+			SSEWriteTimeout: time.Second,
+			Logger:          zap.New(logCore),
+		})
+		require.NoError(t, err)
+		require.NotNil(t, writer)
+		require.NoError(t, writer.Heartbeat())
+		assert.True(t, recorder.Flushed)
+		assert.Equal(t, 1, logs.FilterMessage("SSE write timeout disabled because response writer does not support write deadlines").Len())
 	})
 
 	t.Run("does not require deadline support when the timeout is disabled", func(t *testing.T) {