Ensure Written is set in GZIP ProxyResponseWriter (#9018) (#9026)

Fix #9001

The GZIP ProxyReponseWriter doesn't currently respond correctly
to requests about its Written status - leading to #9001.

This PR properly reimplements these methods.

diff --git a/modules/gzip/gzip.go b/modules/gzip/gzip.go
index 0d1007183027d81804f365da9a2446520f5b935e..a836f8a9166d85ab7f5938eb5b84787cff7c9253 100644 (file)
--- a/modules/gzip/gzip.go
+++ b/modules/gzip/gzip.go
@@ -123,7 +123,7 @@ func Middleware(options ...Options) macaron.Handler {
                // OK we should proxy the response writer
                // We are still not necessarily going to compress...
                proxyWriter := &ProxyResponseWriter{
-                       ResponseWriter: ctx.Resp,
+                       internal: ctx.Resp,
                }
                defer proxyWriter.Close()
 
@@ -137,19 +137,52 @@ func Middleware(options ...Options) macaron.Handler {
                }
 
                ctx.Next()
+               ctx.Resp = proxyWriter.internal
        }
 }
 
 // ProxyResponseWriter is a wrapped macaron ResponseWriter that may compress its contents
 type ProxyResponseWriter struct {
-       writer io.WriteCloser
-       macaron.ResponseWriter
-       stopped bool
+       writer   io.WriteCloser
+       internal macaron.ResponseWriter
+       stopped  bool
 
        code int
        buf  []byte
 }
 
+// Header returns the header map
+func (proxy *ProxyResponseWriter) Header() http.Header {
+       return proxy.internal.Header()
+}
+
+// Status returns the status code of the response or 0 if the response has not been written.
+func (proxy *ProxyResponseWriter) Status() int {
+       if proxy.code != 0 {
+               return proxy.code
+       }
+       return proxy.internal.Status()
+}
+
+// Written returns whether or not the ResponseWriter has been written.
+func (proxy *ProxyResponseWriter) Written() bool {
+       if proxy.code != 0 {
+               return true
+       }
+       return proxy.internal.Written()
+}
+
+// Size returns the size of the response body.
+func (proxy *ProxyResponseWriter) Size() int {
+       return proxy.internal.Size()
+}
+
+// Before allows for a function to be called before the ResponseWriter has been written to. This is
+// useful for setting headers or any other operations that must happen before a response has been written.
+func (proxy *ProxyResponseWriter) Before(before macaron.BeforeFunc) {
+       proxy.internal.Before(before)
+}
+
 // Write appends data to the proxied gzip writer.
 func (proxy *ProxyResponseWriter) Write(b []byte) (int, error) {
        // if writer is initialized, use the writer
@@ -210,7 +243,7 @@ func (proxy *ProxyResponseWriter) startGzip() error {
 
        // Write the header to gzip response.
        if proxy.code != 0 {
-               proxy.ResponseWriter.WriteHeader(proxy.code)
+               proxy.internal.WriteHeader(proxy.code)
                // Ensure that no other WriteHeader's happen
                proxy.code = 0
        }
@@ -220,7 +253,7 @@ func (proxy *ProxyResponseWriter) startGzip() error {
        // write the gzip header even if nothing was ever written.
        if len(proxy.buf) > 0 {
                // Initialize the GZIP response.
-               proxy.writer = writerPool.Get(proxy.ResponseWriter)
+               proxy.writer = writerPool.Get(proxy.internal)
 
                return proxy.writeBuf()
        }
@@ -229,11 +262,11 @@ func (proxy *ProxyResponseWriter) startGzip() error {
 
 func (proxy *ProxyResponseWriter) startPlain() error {
        if proxy.code != 0 {
-               proxy.ResponseWriter.WriteHeader(proxy.code)
+               proxy.internal.WriteHeader(proxy.code)
                proxy.code = 0
        }
        proxy.stopped = true
-       proxy.writer = noopCloser{proxy.ResponseWriter}
+       proxy.writer = noopCloser{proxy.internal}
        return proxy.writeBuf()
 }
 
@@ -295,13 +328,13 @@ func (proxy *ProxyResponseWriter) Flush() {
                gw.Flush()
        }
 
-       proxy.ResponseWriter.Flush()
+       proxy.internal.Flush()
 }
 
 // Hijack implements http.Hijacker. If the underlying ResponseWriter is a
 // Hijacker, its Hijack method is returned. Otherwise an error is returned.
 func (proxy *ProxyResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
-       hijacker, ok := proxy.ResponseWriter.(http.Hijacker)
+       hijacker, ok := proxy.internal.(http.Hijacker)
        if !ok {
                return nil, nil, fmt.Errorf("the ResponseWriter doesn't support the Hijacker interface")
        }