From 2c7d1c04089b7910a46674f031ffeb6306b16092 Mon Sep 17 00:00:00 2001 From: Robert Hurst Date: Sun, 2 Feb 2025 00:31:07 -0800 Subject: [PATCH] Track if the response writer has been written to when taken from the context. --- context-finalize.go | 8 +++++--- context.go | 44 +++++++++++++++++++++++++++++++++++--------- router.go | 5 +++++ 3 files changed, 45 insertions(+), 12 deletions(-) diff --git a/context-finalize.go b/context-finalize.go index 9bd713c..b81c0ed 100644 --- a/context-finalize.go +++ b/context-finalize.go @@ -111,15 +111,17 @@ func (c *Context) finalize() { if !c.inhibitResponse && hasBody { if is100Range || is204Or304 { - fmt.Printf("Response with status %d has body but no content is expected", c.Status) + fmt.Printf("response with status %d has body but no content is expected", c.Status) } else { _, err := io.Copy(c.bodyWriter, finalBodyReader) if finalBodyReaderCloser, ok := finalBodyReader.(io.Closer); ok { - finalBodyReaderCloser.Close() + if err := finalBodyReaderCloser.Close(); err != nil && PrintHandlerErrors { + fmt.Printf("Failed to close body read closer: %s", err) + } } if err != nil { c.Status = 500 - fmt.Printf("Error occurred when writing response body: %s", err) + fmt.Printf("error occurred when writing response body: %s", err) } } } diff --git a/context.go b/context.go index 42297c3..93d5491 100644 --- a/context.go +++ b/context.go @@ -64,7 +64,7 @@ var _ io.WriteCloser = &Context{} // NewContext creates a new Context from go's http.ResponseWriter and // http.Request. It also takes a variadic list of handlers. This is useful for -// creating a new Context outside of a router, and can be used by libraries +// creating a new Context outside a router, and can be used by libraries // which wish to extend or encapsulate the functionality of Navaros. func NewContext(res http.ResponseWriter, req *http.Request, handlers ...any) *Context { return NewContextWithNode(res, req, &HandlerNode{ @@ -75,7 +75,7 @@ func NewContext(res http.ResponseWriter, req *http.Request, handlers ...any) *Co // NewContextWithNode creates a new Context from go's http.ResponseWriter and // http.Request. It also takes a HandlerNode - a link in a chain of handlers. -// This is useful for creating a new Context outside of a router, and can be +// This is useful for creating a new Context outside a router, and can be // used by libraries which wish to extend or encapsulate the functionality of // Navaros. For example, implementing a custom router. func NewContextWithNode(res http.ResponseWriter, req *http.Request, firstHandlerNode *HandlerNode) *Context { @@ -309,7 +309,7 @@ func (c *Context) RequestTrailers() http.Header { return c.request.Trailer } -// RequestCookies returns the value of a request cookie by name. Returns nil +// RequestCookie returns the value of a request cookie by name. Returns nil // if the cookie does not exist. func (c *Context) RequestCookie(name string) (*http.Cookie, error) { return c.request.Cookie(name) @@ -330,10 +330,10 @@ func (c *Context) RequestBodyReader() io.ReadCloser { return http.MaxBytesReader(c.bodyWriter, c.request.Body, maxRequestBodySize) } -// Allows middleware to intercept the request body reader and replace it with -// their own. This is useful transformers that re-write the request body -// in a streaming fashion. It's also useful for transformers that re-encode -// the request body. +// SetRequestBodyReader Allows middleware to intercept the request body reader +// and replace it with their own. This is useful transformers that re-write the +// request body in a streaming fashion. It's also useful for transformers that +// re-encode the request body. func (c *Context) SetRequestBodyReader(reader io.Reader) { if readCloser, ok := reader.(io.ReadCloser); ok { c.request.Body = readCloser @@ -406,7 +406,11 @@ func (c *Context) Request() *http.Request { // ResponseWriter returns the underlying http.ResponseWriter object. func (c *Context) ResponseWriter() http.ResponseWriter { - return c.bodyWriter + return &ContextResponseWriter{ + hasWrittenHeaders: &c.hasWrittenHeaders, + hasWrittenBody: &c.hasWrittenBody, + bodyWriter: c.bodyWriter, + } } // Write writes bytes to the response body. This is useful for streaming the @@ -467,7 +471,7 @@ func (c *Context) Err() error { } // Value is a noop for compatibility with go's context.Context. -func (c *Context) Value(key any) any { +func (c *Context) Value(any) any { return nil } @@ -480,3 +484,25 @@ func (c *Context) marshallResponseBody() (io.Reader, error) { } return c.responseBodyMarshaller(c.Body) } + +type ContextResponseWriter struct { + hasWrittenHeaders *bool + hasWrittenBody *bool + bodyWriter http.ResponseWriter +} + +var _ http.ResponseWriter = &ContextResponseWriter{} + +func (c *ContextResponseWriter) Write(bytes []byte) (int, error) { + *c.hasWrittenBody = true + return c.bodyWriter.Write(bytes) +} + +func (c *ContextResponseWriter) WriteHeader(status int) { + *c.hasWrittenHeaders = true + c.bodyWriter.WriteHeader(status) +} + +func (c *ContextResponseWriter) Header() http.Header { + return c.bodyWriter.Header() +} diff --git a/router.go b/router.go index 81c7730..8451142 100644 --- a/router.go +++ b/router.go @@ -8,6 +8,11 @@ import ( var PrintHandlerErrors = false +// SetPrintHandlerErrors toggles the printing of handler errors. +func SetPrintHandlerErrors(enable bool) { + PrintHandlerErrors = enable +} + // Router is the main component of Navaros. It is an HTTP handler that can be // used to handle requests, and route them to the appropriate handlers. It // implements the http.Handler interface, and can be used as a handler in