Repository: avenga/couper
Pull request #245 (2021-05-25 15:01:56)

Refactor: Split gzip writer
Refactor the writer to prepare a better solution for #226 


Patch:
diff --git a/command/run.go b/command/run.go
index 2b7e16c7..c491a135 100644
--- a/command/run.go
+++ b/command/run.go
@@ -6,13 +6,15 @@ import (
 	"fmt"
 	"sync"
 
+	"github.com/sirupsen/logrus"
+
 	"github.com/avenga/couper/cache"
 	"github.com/avenga/couper/config"
 	"github.com/avenga/couper/config/env"
 	"github.com/avenga/couper/config/runtime"
 	"github.com/avenga/couper/errors"
 	"github.com/avenga/couper/server"
-	"github.com/sirupsen/logrus"
+	"github.com/avenga/couper/server/writer"
 )
 
 var _ Cmd = &Run{}
@@ -55,7 +57,7 @@ func (r *Run) Execute(args Args, config *config.Couper, logEntry *logrus.Entry)
 	}
 
 	if config.Settings.SecureCookies != "" &&
-		config.Settings.SecureCookies != server.SecureCookiesStrip {
+		config.Settings.SecureCookies != writer.SecureCookiesStrip {
 		return fmt.Errorf("invalid value for the -secure-cookies flag given: '%s' only 'strip' is supported", config.Settings.SecureCookies)
 	}
 
diff --git a/handler/transport/backend.go b/handler/transport/backend.go
index 9cb908ef..94a75f8a 100644
--- a/handler/transport/backend.go
+++ b/handler/transport/backend.go
@@ -6,7 +6,6 @@ import (
 	"encoding/base64"
 	"net/http"
 	"net/url"
-	"regexp"
 	"strings"
 	"time"
 
@@ -19,21 +18,12 @@ import (
 	"github.com/avenga/couper/eval"
 	"github.com/avenga/couper/handler/validation"
 	"github.com/avenga/couper/logging"
+	"github.com/avenga/couper/server/writer"
 	"github.com/avenga/couper/utils"
 )
 
-const (
-	GzipName              = "gzip"
-	AcceptEncodingHeader  = "Accept-Encoding"
-	ContentEncodingHeader = "Content-Encoding"
-	ContentLengthHeader   = "Content-Length"
-	VaryHeader            = "Vary"
-)
-
 var _ http.RoundTripper = &Backend{}
 
-var ReClientSupportsGZ = regexp.MustCompile(`(?i)\b` + GzipName + `\b`)
-
 // Backend represents the transport configuration.
 type Backend struct {
 	context          hcl.Body
@@ -96,11 +86,7 @@ func (b *Backend) RoundTrip(req *http.Request) (*http.Response, error) {
 		removeHopHeaders(req.Header)
 	}
 
-	if ReClientSupportsGZ.MatchString(req.Header.Get(AcceptEncodingHeader)) {
-		req.Header.Set(AcceptEncodingHeader, GzipName)
-	} else {
-		req.Header.Del(AcceptEncodingHeader)
-	}
+	writer.ModifyAcceptEncoding(req.Header)
 
 	if xff, ok := req.Context().Value(request.XFF).(string); ok {
 		if xff != "" {
@@ -129,13 +115,7 @@ func (b *Backend) RoundTrip(req *http.Request) (*http.Response, error) {
 		return nil, err
 	}
 
-	if strings.ToLower(beresp.Header.Get(ContentEncodingHeader)) == GzipName {
-		src, rerr := gzip.NewReader(beresp.Body)
-		if rerr == nil {
-			beresp.Header.Del(ContentEncodingHeader)
-			beresp.Body = eval.NewReadCloser(src, beresp.Body)
-		}
-	}
+	setGzipReader(beresp)
 
 	if !isProxyReq {
 		removeConnectionHeaders(req.Header)
@@ -327,6 +307,15 @@ func setUserAgent(outreq *http.Request) {
 	}
 }
 
+func setGzipReader(beresp *http.Response) {
+	if strings.ToLower(beresp.Header.Get(writer.ContentEncodingHeader)) == writer.GzipName {
+		if src, err := gzip.NewReader(beresp.Body); err == nil {
+			beresp.Header.Del(writer.ContentEncodingHeader)
+			beresp.Body = eval.NewReadCloser(src, beresp.Body)
+		}
+	}
+}
+
 // removeConnectionHeaders removes hop-by-hop headers listed in the "Connection" header of h.
 // See RFC 7230, section 6.1
 func removeConnectionHeaders(h http.Header) {
diff --git a/server/http.go b/server/http.go
index 8ebd9341..8bee5958 100644
--- a/server/http.go
+++ b/server/http.go
@@ -19,8 +19,8 @@ import (
 	"github.com/avenga/couper/errors"
 	"github.com/avenga/couper/eval"
 	"github.com/avenga/couper/handler"
-	"github.com/avenga/couper/handler/transport"
 	"github.com/avenga/couper/logging"
+	"github.com/avenga/couper/server/writer"
 )
 
 type muxers map[string]*Mux
@@ -201,27 +201,23 @@ func (s *HTTPServer) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 	}
 
 	h := mux.FindHandler(req)
-	w := NewRWWrapper(rw,
-		transport.ReClientSupportsGZ.MatchString(
-			req.Header.Get(transport.AcceptEncodingHeader),
-		),
-		s.settings.SecureCookies,
-	)
-	rw = w
 
 	clientReq := req.Clone(req.Context())
 
+	gw := writer.NewGzipWriter(rw, clientReq.Header)
+	w := writer.NewResponseWriter(gw, s.settings.SecureCookies)
+
 	if err = s.setGetBody(h, clientReq); err != nil {
-		mux.opts.ServerOptions.ServerErrTpl.ServeError(err).ServeHTTP(rw, req)
+		mux.opts.ServerOptions.ServerErrTpl.ServeError(err).ServeHTTP(w, clientReq)
 		return
 	}
 
 	ctx = s.evalCtx.WithClientRequest(clientReq)
 	*clientReq = *clientReq.WithContext(ctx)
 
-	s.accessLog.ServeHTTP(rw, clientReq, h, startTime)
+	s.accessLog.ServeHTTP(w, clientReq, h, startTime)
 
-	w.Close() // Closes the GZ writer.
+	_ = gw.Close() // Closes the GZ writer.
 }
 
 func (s *HTTPServer) setGetBody(h http.Handler, req *http.Request) error {
diff --git a/server/rw_wrapper.go b/server/rw_wrapper.go
deleted file mode 100644
index c9df43c9..00000000
--- a/server/rw_wrapper.go
+++ /dev/null
@@ -1,169 +0,0 @@
-package server
-
-import (
-	"bufio"
-	"bytes"
-	"compress/gzip"
-	"fmt"
-	"net"
-	"net/http"
-	"net/textproto"
-	"strconv"
-
-	"github.com/avenga/couper/handler/transport"
-	"github.com/avenga/couper/logging"
-)
-
-var (
-	_ http.Flusher         = &RWWrapper{}
-	_ http.Hijacker        = &RWWrapper{}
-	_ http.ResponseWriter  = &RWWrapper{}
-	_ logging.RecorderInfo = &RWWrapper{}
-)
-
-// RWWrapper wraps the <http.ResponseWriter>.
-type RWWrapper struct {
-	rw            http.ResponseWriter
-	gz            *gzip.Writer
-	headerBuffer  *bytes.Buffer
-	httpStatus    []byte
-	httpLineDelim []byte
-	secureCookies string
-	statusWritten bool
-	// logging info
-	statusCode      int
-	rawBytesWritten int
-	bytesWritten    int
-}
-
-// NewRWWrapper creates a new RWWrapper object.
-func NewRWWrapper(rw http.ResponseWriter, useGZ bool, secureCookies string) *RWWrapper {
-	w := &RWWrapper{
-		rw:            rw,
-		headerBuffer:  &bytes.Buffer{},
-		secureCookies: secureCookies,
-	}
-
-	if useGZ {
-		w.gz = gzip.NewWriter(rw)
-	}
-
-	return w
-}
-
-// Header wraps the Header method of the <http.ResponseWriter>.
-func (w *RWWrapper) Header() http.Header {
-	return w.rw.Header()
-}
-
-// Write wraps the Write method of the <http.ResponseWriter>.
-func (w *RWWrapper) Write(p []byte) (int, error) {
-	l := len(p)
-	w.rawBytesWritten += l
-	if !w.statusWritten {
-		if len(w.httpStatus) == 0 {
-			w.httpStatus = p[:]
-			// required for short writes without any additional header
-			// to detect EOH chunk later on
-			if l >= 2 {
-				w.httpLineDelim = p[l-2 : l]
-			}
-
-			return l, nil
-		}
-
-		// End-of-header
-		// http.Response.Write() EOH chunk is: '\r\n'
-		if bytes.Equal(w.httpLineDelim, p) {
-			reader := textproto.NewReader(bufio.NewReader(w.headerBuffer))
-			header, _ := reader.ReadMIMEHeader()
-			for k := range header {
-				w.rw.Header()[k] = header.Values(k)
-			}
-			w.WriteHeader(w.parseStatusCode(w.httpStatus))
-		}
-
-		if l >= 2 {
-			w.httpLineDelim = p[l-2 : l]
-		}
-		return w.headerBuffer.Write(p)
-	}
-
-	var n int
-	var writeErr error
-	if w.gz != nil {
-		n, writeErr = w.gz.Write(p)
-	} else {
-		n, writeErr = w.rw.Write(p)
-	}
-	w.bytesWritten += n
-	return n, writeErr
-}
-
-func (w *RWWrapper) Hijack() (net.Conn, *bufio.ReadWriter, error) {
-	hijack, ok := w.rw.(http.Hijacker)
-	if !ok {
-		return nil, nil, fmt.Errorf("can't switch protocols using non-Hijacker ResponseWriter type %T", w.rw)
-	}
-	return hijack.Hijack()
-}
-
-// Close closes the GZ writer.
-func (w *RWWrapper) Close() {
-	if w.gz != nil {
-		_ = w.gz.Close()
-	}
-}
-
-// Flush implements the <http.Flusher> interface.
-func (w *RWWrapper) Flush() {
-	if w.gz != nil {
-		_ = w.gz.Flush()
-	}
-
-	if rw, ok := w.rw.(http.Flusher); ok {
-		rw.Flush()
-	}
-}
-
-// WriteHeader wraps the WriteHeader method of the <http.ResponseWriter>.
-func (w *RWWrapper) WriteHeader(statusCode int) {
-	if w.statusWritten {
-		return
-	}
-
-	w.configureHeader()
-	w.rw.WriteHeader(statusCode)
-	w.statusWritten = true
-	w.statusCode = statusCode
-}
-
-func (w *RWWrapper) configureHeader() {
-	w.rw.Header().Set("Server", "couper.io")
-	w.rw.Header().Add(transport.VaryHeader, transport.AcceptEncodingHeader)
-
-	if w.gz != nil {
-		w.rw.Header().Del(transport.ContentLengthHeader)
-		w.rw.Header().Set(transport.ContentEncodingHeader, transport.GzipName)
-	}
-
-	if w.secureCookies == SecureCookiesStrip {
-		stripSecureCookies(w.rw.Header())
-	}
-}
-
-func (w *RWWrapper) parseStatusCode(p []byte) int {
-	if len(p) < 12 {
-		return 0
-	}
-	code, _ := strconv.Atoi(string(p[9:12]))
-	return code
-}
-
-func (w *RWWrapper) StatusCode() int {
-	return w.statusCode
-}
-
-func (w *RWWrapper) WrittenBytes() int {
-	return w.bytesWritten
-}
diff --git a/server/cookies.go b/server/writer/cookies.go
similarity index 98%
rename from server/cookies.go
rename to server/writer/cookies.go
index 9b9417fd..cfca86a7 100644
--- a/server/cookies.go
+++ b/server/writer/cookies.go
@@ -1,4 +1,4 @@
-package server
+package writer
 
 import (
 	"net/http"
diff --git a/server/writer/gzip.go b/server/writer/gzip.go
new file mode 100644
index 00000000..11b9a298
--- /dev/null
+++ b/server/writer/gzip.go
@@ -0,0 +1,95 @@
+package writer
+
+import (
+	"bufio"
+	"compress/gzip"
+	"fmt"
+	"net"
+	"net/http"
+	"regexp"
+)
+
+const (
+	AcceptEncodingHeader  = "Accept-Encoding"
+	ContentEncodingHeader = "Content-Encoding"
+	ContentLengthHeader   = "Content-Length"
+	GzipName              = "gzip"
+	VaryHeader            = "Vary"
+)
+
+var (
+	clientSupportsGZ = regexp.MustCompile(`(?i)\b` + GzipName + `\b`)
+
+	_ writer = &Gzip{}
+)
+
+type Gzip struct {
+	enabled bool
+	rw      http.ResponseWriter
+	w       *gzip.Writer
+	written int64
+}
+
+func NewGzipWriter(rw http.ResponseWriter, header http.Header) *Gzip {
+	gw := gzip.NewWriter(rw)
+	return &Gzip{
+		enabled: clientSupportsGZ.MatchString(header.Get(AcceptEncodingHeader)),
+		rw:      rw,
+		w:       gw,
+	}
+}
+
+func (g *Gzip) Write(p []byte) (n int, err error) {
+	if g.enabled {
+		return g.w.Write(p)
+	}
+	return g.rw.Write(p)
+}
+
+func (g *Gzip) Close() (err error) {
+	if g.enabled && g.w != nil {
+		err = g.w.Close()
+	}
+	return err
+}
+
+func (g *Gzip) Header() http.Header {
+	return g.rw.Header()
+}
+
+func (g *Gzip) WriteHeader(statusCode int) {
+	g.rw.Header().Add(VaryHeader, AcceptEncodingHeader)
+
+	if g.enabled {
+		g.rw.Header().Del(ContentLengthHeader)
+		g.rw.Header().Set(ContentEncodingHeader, GzipName)
+	}
+
+	g.rw.WriteHeader(statusCode)
+}
+
+func (g *Gzip) Flush() {
+	if g.enabled && g.w != nil {
+		_ = g.w.Flush()
+	}
+
+	if rw, ok := g.rw.(http.Flusher); ok {
+		rw.Flush()
+	}
+}
+
+func (g *Gzip) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	hijack, ok := g.rw.(http.Hijacker)
+	if !ok {
+		return nil, nil, fmt.Errorf("can't switch protocols using non-Hijacker gzip writer type %T", g.rw)
+	}
+	return hijack.Hijack()
+}
+
+func ModifyAcceptEncoding(header http.Header) {
+	if clientSupportsGZ.MatchString(header.Get(AcceptEncodingHeader)) {
+		header.Set(AcceptEncodingHeader, GzipName)
+	} else {
+		header.Del(AcceptEncodingHeader)
+	}
+}
diff --git a/server/writer/response.go b/server/writer/response.go
new file mode 100644
index 00000000..efff24a5
--- /dev/null
+++ b/server/writer/response.go
@@ -0,0 +1,141 @@
+package writer
+
+import (
+	"bufio"
+	"bytes"
+	"fmt"
+	"net"
+	"net/http"
+	"net/textproto"
+	"strconv"
+
+	"github.com/avenga/couper/logging"
+)
+
+type writer interface {
+	http.Flusher
+	http.Hijacker
+	http.ResponseWriter
+}
+
+var (
+	_ writer               = &Response{}
+	_ logging.RecorderInfo = &Response{}
+)
+
+// Response wraps the http.ResponseWriter.
+type Response struct {
+	rw            http.ResponseWriter
+	headerBuffer  *bytes.Buffer
+	httpStatus    []byte
+	httpLineDelim []byte
+	secureCookies string
+	statusWritten bool
+	// logging info
+	statusCode      int
+	rawBytesWritten int
+	bytesWritten    int
+}
+
+// NewResponseWriter creates a new Response object.
+func NewResponseWriter(rw http.ResponseWriter, secureCookies string) *Response {
+	return &Response{
+		rw:            rw,
+		headerBuffer:  &bytes.Buffer{},
+		secureCookies: secureCookies,
+	}
+}
+
+// Header wraps the Header method of the <http.ResponseWriter>.
+func (r *Response) Header() http.Header {
+	return r.rw.Header()
+}
+
+// Write wraps the Write method of the <http.ResponseWriter>.
+func (r *Response) Write(p []byte) (int, error) {
+	l := len(p)
+	r.rawBytesWritten += l
+	if !r.statusWritten {
+		if len(r.httpStatus) == 0 {
+			r.httpStatus = p[:]
+			// required for short writes without any additional header
+			// to detect EOH chunk later on
+			if l >= 2 {
+				r.httpLineDelim = p[l-2 : l]
+			}
+
+			return l, nil
+		}
+
+		// End-of-header
+		// http.Response.Write() EOH chunk is: '\r\n'
+		if bytes.Equal(r.httpLineDelim, p) {
+			reader := textproto.NewReader(bufio.NewReader(r.headerBuffer))
+			header, _ := reader.ReadMIMEHeader()
+			for k := range header {
+				r.rw.Header()[k] = header.Values(k)
+			}
+			r.WriteHeader(r.parseStatusCode(r.httpStatus))
+		}
+
+		if l >= 2 {
+			r.httpLineDelim = p[l-2 : l]
+		}
+		return r.headerBuffer.Write(p)
+	}
+
+	n, writeErr := r.rw.Write(p)
+	r.bytesWritten += n
+	return n, writeErr
+}
+
+func (r *Response) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	hijack, ok := r.rw.(http.Hijacker)
+	if !ok {
+		return nil, nil, fmt.Errorf("can't switch protocols using non-Hijacker ResponseWriter type %T", r.rw)
+	}
+	return hijack.Hijack()
+}
+
+// Flush implements the <http.Flusher> interface.
+func (r *Response) Flush() {
+	if rw, ok := r.rw.(http.Flusher); ok {
+		rw.Flush()
+	}
+}
+
+// WriteHeader wraps the WriteHeader method of the <http.ResponseWriter>.
+func (r *Response) WriteHeader(statusCode int) {
+	if r.statusWritten {
+		return
+	}
+
+	r.configureHeader()
+	r.rw.WriteHeader(statusCode)
+	r.statusWritten = true
+	r.statusCode = statusCode
+}
+
+func (r *Response) configureHeader() {
+	r.rw.Header().Set("Server", "couper.io")
+
+	if r.secureCookies == SecureCookiesStrip {
+		stripSecureCookies(r.rw.Header())
+	}
+}
+
+func (r *Response) parseStatusCode(p []byte) int {
+	if len(p) < 12 {
+		return 0
+	}
+	code, _ := strconv.Atoi(string(p[9:12]))
+	return code
+}
+
+func (r *Response) StatusCode() int {
+	return r.statusCode
+}
+
+func (r *Response) WrittenBytes() int {
+	return r.bytesWritten
+}


Test patch:
diff --git a/errors/template_test.go b/errors/template_test.go
index 15374e29..efbde46d 100644
--- a/errors/template_test.go
+++ b/errors/template_test.go
@@ -8,7 +8,7 @@ import (
 
 	"github.com/avenga/couper/errors"
 	"github.com/avenga/couper/internal/test"
-	"github.com/avenga/couper/server"
+	"github.com/avenga/couper/server/writer"
 )
 
 func TestTemplate_ServeError(t1 *testing.T) {
@@ -25,7 +25,7 @@ func TestTemplate_ServeError(t1 *testing.T) {
 	}
 	for _, tt := range tests {
 		t1.Run(tt.name, func(t2 *testing.T) {
-			rec := server.NewRWWrapper(httptest.NewRecorder(), false, "")
+			rec := writer.NewResponseWriter(httptest.NewRecorder(), "")
 			req := httptest.NewRequest(http.MethodGet, "/", nil)
 			errors.DefaultJSON.ServeError(tt.err).ServeHTTP(rec, req)
 
diff --git a/handler/endpoint_test.go b/handler/endpoint_test.go
index 9cfc302d..fd206705 100644
--- a/handler/endpoint_test.go
+++ b/handler/endpoint_test.go
@@ -20,7 +20,7 @@ import (
 	"github.com/avenga/couper/handler/producer"
 	"github.com/avenga/couper/handler/transport"
 	"github.com/avenga/couper/internal/test"
-	"github.com/avenga/couper/server"
+	"github.com/avenga/couper/server/writer"
 )
 
 func TestEndpoint_RoundTrip_Eval(t *testing.T) {
@@ -101,7 +101,7 @@ func TestEndpoint_RoundTrip_Eval(t *testing.T) {
 			*req = *req.WithContext(evalCtx.WithClientRequest(req))
 
 			rec := httptest.NewRecorder()
-			rw := server.NewRWWrapper(rec, false, "") // crucial for working ep due to res.Write()
+			rw := writer.NewResponseWriter(rec, "") // crucial for working ep due to res.Write()
 			ep.ServeHTTP(rw, req)
 			rec.Flush()
 			res := rec.Result()
@@ -228,7 +228,7 @@ func TestEndpoint_RoundTripContext_Variables_json_body(t *testing.T) {
 				*req = *req.WithContext(eval.NewContext(nil).WithClientRequest(req))
 
 				rec := httptest.NewRecorder()
-				rw := server.NewRWWrapper(rec, false, "") // crucial for working ep due to res.Write()
+				rw := writer.NewResponseWriter(rec, "") // crucial for working ep due to res.Write()
 				ep.ServeHTTP(rw, req)
 				rec.Flush()
 				res := rec.Result()
@@ -340,7 +340,7 @@ func TestEndpoint_RoundTripContext_Null_Eval(t *testing.T) {
 			*req = *req.WithContext(ctx)
 
 			rec := httptest.NewRecorder()
-			rw := server.NewRWWrapper(rec, false, "") // crucial for working ep due to res.Write()
+			rw := writer.NewResponseWriter(rec, "") // crucial for working ep due to res.Write()
 			ep.ServeHTTP(rw, req)
 			rec.Flush()
 			res := rec.Result()
diff --git a/handler/health_test.go b/handler/health_test.go
index ef4f5201..e9cb12ee 100644
--- a/handler/health_test.go
+++ b/handler/health_test.go
@@ -8,8 +8,9 @@ import (
 	"reflect"
 	"testing"
 
+	"github.com/avenga/couper/server/writer"
+
 	"github.com/avenga/couper/handler"
-	"github.com/avenga/couper/server"
 )
 
 func TestHealth_Match(t *testing.T) {
@@ -56,18 +57,20 @@ func TestHealth_ServeHTTP(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(subT *testing.T) {
 			h := handler.NewHealthCheck(tt.fields.path, tt.fields.shutdownCh)
-			rec := httptest.NewRecorder()
-			rw := server.NewRWWrapper(rec, tt.fields.gzip, "")
-			if tt.wantStatus >= 500 {
-				close(tt.fields.shutdownCh)
-			}
 
 			if tt.fields.gzip {
 				tt.req.Header.Set("Accept-Encoding", "gzip")
 			}
 
+			rec := httptest.NewRecorder()
+			gw := writer.NewGzipWriter(rec, tt.req.Header)
+			rw := writer.NewResponseWriter(gw, "")
+			if tt.wantStatus >= 500 {
+				close(tt.fields.shutdownCh)
+			}
+
 			h.ServeHTTP(rw, tt.req)
-			rw.Close()
+			_ = gw.Close()
 
 			rec.Flush()
 			res := rec.Result()
diff --git a/server/cookies_internal_test.go b/server/writer/cookies_internal_test.go
similarity index 98%
rename from server/cookies_internal_test.go
rename to server/writer/cookies_internal_test.go
index 74b54a08..a729e11d 100644
--- a/server/cookies_internal_test.go
+++ b/server/writer/cookies_internal_test.go
@@ -1,4 +1,4 @@
-package server
+package writer
 
 import (
 	"net/http"