add some guardrails and notes to the example code

Author: robbyt

examples/custom_middleware/example/jsonenforcer.go

@@ -9,6 +9,15 @@ import (
 )
 
 // ResponseWrapper captures response data for transformation
+//
+// This is a simple example for demonstration purposes and is not intended for
+// production use. Limitations:
+// - Does not preserve optional HTTP interfaces (http.Hijacker, http.Flusher, http.Pusher)
+// - Not safe for concurrent writes from multiple goroutines within the same request
+// - No memory limits on buffered content
+//
+// Each request gets its own ResponseWrapper instance, so different requests won't
+// interfere with each other.
 type ResponseWrapper struct {
 	original httpserver.ResponseWriter
 	buffer   *bytes.Buffer
@@ -63,9 +72,8 @@ func (rw *ResponseWrapper) Size() int {
 
 // transformToJSON wraps non-JSON content in a JSON response
 func transformToJSON(data []byte) ([]byte, error) {
-	// First check if the data is valid JSON regardless of content type
-	var test any
-	if json.Unmarshal(data, &test) == nil {
+	// Use json.Valid for efficient validation without unmarshaling
+	if json.Valid(data) {
 		return data, nil // Valid JSON, return as-is
 	}
 
@@ -91,6 +99,25 @@ func New() httpserver.HandlerFunc {
 
 		// RESPONSE PHASE: Transform response to JSON
 		originalData := wrapper.buffer.Bytes()
+		statusCode := wrapper.Status()
+		if statusCode == 0 {
+			statusCode = http.StatusOK
+		}
+
+		// Copy headers to original writer
+		originalWriter := wrapper.original
+		for key, values := range wrapper.Header() {
+			for _, value := range values {
+				originalWriter.Header().Add(key, value)
+			}
+		}
+
+		// Check if this status code should have no body per HTTP spec
+		// 204 No Content and 304 Not Modified MUST NOT have a message body
+		if statusCode == http.StatusNoContent || statusCode == http.StatusNotModified {
+			originalWriter.WriteHeader(statusCode)
+			return
+		}
 
 		// Transform captured data to JSON
 		if len(originalData) == 0 && wrapper.status == 0 {
@@ -104,22 +131,10 @@ func New() httpserver.HandlerFunc {
 			jsonData = []byte(`{"error":"Unable to encode response"}`)
 		}
 
-		// Copy headers to original writer
-		originalWriter := wrapper.original
-		for key, values := range wrapper.Header() {
-			for _, value := range values {
-				originalWriter.Header().Add(key, value)
-			}
-		}
-
 		// Ensure JSON content type
 		originalWriter.Header().Set("Content-Type", "application/json")
 
 		// Write status and transformed data
-		statusCode := wrapper.Status()
-		if statusCode == 0 {
-			statusCode = http.StatusOK
-		}
 		originalWriter.WriteHeader(statusCode)
 		if _, err := originalWriter.Write(jsonData); err != nil {
 			// Response is already committed, cannot recover from write error

examples/custom_middleware/example/jsonenforcer_test.go

@@ -103,11 +103,11 @@ func TestJSONEnforcerMiddleware(t *testing.T) {
 		)
 	})
 
-	t.Run("handles empty response body", func(t *testing.T) {
+	t.Run("handles 204 No Content correctly", func(t *testing.T) {
 		req := httptest.NewRequest("GET", "/test", nil)
 		rec := httptest.NewRecorder()
 
-		// Handler that writes nothing
+		// Handler that returns 204 No Content
 		emptyHandler := func(w http.ResponseWriter, r *http.Request) {
 			w.WriteHeader(http.StatusNoContent)
 		}
@@ -123,13 +123,34 @@ func TestJSONEnforcerMiddleware(t *testing.T) {
 		route.ServeHTTP(rec, req)
 
 		assert.Equal(t, http.StatusNoContent, rec.Code, "status code should be preserved")
-		assert.Equal(t, "application/json", rec.Header().Get("Content-Type"))
-		assert.Equal(
-			t,
-			`{"response":""}`,
-			rec.Body.String(),
-			"empty response should be wrapped",
+		assert.Empty(t, rec.Body.String(), "204 responses must have no body per HTTP spec")
+		assert.Equal(t, 0, rec.Body.Len(), "content length should be 0")
+	})
+
+	t.Run("handles 304 Not Modified correctly", func(t *testing.T) {
+		req := httptest.NewRequest("GET", "/test", nil)
+		rec := httptest.NewRecorder()
+
+		// Handler that returns 304 Not Modified
+		notModifiedHandler := func(w http.ResponseWriter, r *http.Request) {
+			w.WriteHeader(http.StatusNotModified)
+			// Even if handler writes, middleware should not include body
+			_, err := w.Write([]byte("should be ignored"))
+			assert.NoError(t, err)
+		}
+
+		route, err := httpserver.NewRouteFromHandlerFunc(
+			"test",
+			"/test",
+			notModifiedHandler,
+			New(),
 		)
+		assert.NoError(t, err)
+
+		route.ServeHTTP(rec, req)
+
+		assert.Equal(t, http.StatusNotModified, rec.Code, "status code should be preserved")
+		assert.Empty(t, rec.Body.String(), "304 responses must have no body per HTTP spec")
 	})
 
 	t.Run("transforms HTML response to JSON", func(t *testing.T) {