use writer that can delay sending status code

Author: aldas

context.go

@@ -415,6 +415,15 @@ func (c *Context) Render(code int, name string, data any) (err error) {
 	if c.echo.Renderer == nil {
 		return ErrRendererNotRegistered
 	}
+	// as Renderer.Render can fail, and in that case we need to delay sending status code to the client until
+	// (global) error handler decides the correct status code for the error to be sent to the client, so we need to write
+	//  the rendered template to the buffer first.
+	//
+	// html.Template.ExecuteTemplate() documentations writes:
+	// > If an error occurs executing the template or writing its output,
+	// > execution stops, but partial results may already have been written to
+	// > the output writer.
+
 	buf := new(bytes.Buffer)
 	if err = c.echo.Renderer.Render(c, buf, name, data); err != nil {
 		return
@@ -455,14 +464,12 @@ func (c *Context) jsonPBlob(code int, callback string, i any) (err error) {
 func (c *Context) json(code int, i any, indent string) error {
 	c.writeContentType(MIMEApplicationJSON)
 
-	if r, err := UnwrapResponse(c.response); err == nil {
-		// *echo.Response can delay sending status code until the first Write is called. As serialization can fail, we should delay
-		// sending the status code to the client until serialization is complete (first Write would be an indication it succeeded)
-		// Unsuccessful serialization error needs to go through the error handler and get a proper status code there.
-		r.Status = code
-	} else {
-		return fmt.Errorf("json: response does not unwrap to *echo.Response")
-	}
+	// as JSONSerializer.Serialize can fail, and in that case we need to delay sending status code to the client until
+	// (global) error handler decides correct status code for the error to be sent to the client.
+	// For that we need to use writer that can store the proposed status code until the first Write is called.
+	resp := c.Response()
+	c.SetResponse(&delayedStatusWriter{ResponseWriter: resp, status: code})
+	defer c.SetResponse(resp)
 
 	return c.echo.JSONSerializer.Serialize(c, i, indent)
 }

context_test.go

@@ -12,6 +12,7 @@ import (
 	"io"
 	"io/fs"
 	"log/slog"
+	"math"
 	"mime/multipart"
 	"net/http"
 	"net/http/httptest"
@@ -251,8 +252,8 @@ func TestContextJSONWithNotEchoResponse(t *testing.T) {
 
 	c.SetResponse(rec)
 
-	err := c.JSON(http.StatusOK, map[string]interface{}{"foo": "bar"})
-	assert.EqualError(t, err, "json: response does not unwrap to *echo.Response")
+	err := c.JSON(http.StatusCreated, map[string]float64{"foo": math.NaN()})
+	assert.EqualError(t, err, "json: unsupported value: NaN")
 
 	assert.Equal(t, http.StatusOK, rec.Code) // status code must not be sent to the client
 	assert.Empty(t, rec.Body.String())       // body must not be sent to the client

response.go

@@ -130,3 +130,43 @@ func UnwrapResponse(rw http.ResponseWriter) (*Response, error) {
 		}
 	}
 }
+
+// delayedStatusWriter is a wrapper around http.ResponseWriter that delays writing the status code until first Write is called.
+// This allows (global) error handler to decide correct status code to be sent to the client.
+type delayedStatusWriter struct {
+	http.ResponseWriter
+	commited bool
+	status   int
+}
+
+func (w *delayedStatusWriter) WriteHeader(statusCode int) {
+	// in case something else writes status code explicitly before us we need mark response commited
+	w.commited = true
+	w.ResponseWriter.WriteHeader(statusCode)
+}
+
+func (w *delayedStatusWriter) Write(data []byte) (int, error) {
+	if !w.commited {
+		w.commited = true
+		if w.status == 0 {
+			w.status = http.StatusOK
+		}
+		w.ResponseWriter.WriteHeader(w.status)
+	}
+	return w.ResponseWriter.Write(data)
+}
+
+func (w *delayedStatusWriter) Flush() {
+	err := http.NewResponseController(w.ResponseWriter).Flush()
+	if err != nil && errors.Is(err, http.ErrNotSupported) {
+		panic(errors.New("response writer flushing is not supported"))
+	}
+}
+
+func (w *delayedStatusWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	return http.NewResponseController(w.ResponseWriter).Hijack()
+}
+
+func (w *delayedStatusWriter) Unwrap() http.ResponseWriter {
+	return w.ResponseWriter
+}