Skip to content

Commit

Permalink
Ensure Written is set in GZIP ProxyResponseWriter (#9018) (#9026)
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
zeripath authored Nov 15, 2019
1 parent 2f73fff commit 3c3823d
Showing 1 changed file with 43 additions and 10 deletions.
53 changes: 43 additions & 10 deletions modules/gzip/gzip.go
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand All @@ -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
Expand Down Expand Up @@ -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
}
Expand All @@ -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()
}
Expand All @@ -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()
}

Expand Down Expand Up @@ -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")
}
Expand Down

0 comments on commit 3c3823d

Please sign in to comment.