package compress import ( "context" "io/ioutil" "net/http" "net/http/httptest" "testing" "github.com/NYTimes/gziphandler" "github.com/containous/traefik/v2/pkg/config/dynamic" "github.com/containous/traefik/v2/pkg/testhelpers" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) const ( acceptEncodingHeader = "Accept-Encoding" contentEncodingHeader = "Content-Encoding" contentTypeHeader = "Content-Type" varyHeader = "Vary" gzipValue = "gzip" ) func TestShouldCompressWhenNoContentEncodingHeader(t *testing.T) { req := testhelpers.MustNewRequest(http.MethodGet, "http://localhost", nil) req.Header.Add(acceptEncodingHeader, gzipValue) baseBody := generateBytes(gziphandler.DefaultMinSize) next := http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) { _, err := rw.Write(baseBody) assert.NoError(t, err) }) handler := &compress{next: next} rw := httptest.NewRecorder() handler.ServeHTTP(rw, req) assert.Equal(t, gzipValue, rw.Header().Get(contentEncodingHeader)) assert.Equal(t, acceptEncodingHeader, rw.Header().Get(varyHeader)) if assert.ObjectsAreEqualValues(rw.Body.Bytes(), baseBody) { assert.Fail(t, "expected a compressed body", "got %v", rw.Body.Bytes()) } } func TestShouldNotCompressWhenContentEncodingHeader(t *testing.T) { req := testhelpers.MustNewRequest(http.MethodGet, "http://localhost", nil) req.Header.Add(acceptEncodingHeader, gzipValue) fakeCompressedBody := generateBytes(gziphandler.DefaultMinSize) next := http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) { rw.Header().Add(contentEncodingHeader, gzipValue) rw.Header().Add(varyHeader, acceptEncodingHeader) _, err := rw.Write(fakeCompressedBody) if err != nil { http.Error(rw, err.Error(), http.StatusInternalServerError) } }) handler := &compress{next: next} rw := httptest.NewRecorder() handler.ServeHTTP(rw, req) assert.Equal(t, gzipValue, rw.Header().Get(contentEncodingHeader)) assert.Equal(t, acceptEncodingHeader, rw.Header().Get(varyHeader)) assert.EqualValues(t, rw.Body.Bytes(), fakeCompressedBody) } func TestShouldNotCompressWhenNoAcceptEncodingHeader(t *testing.T) { req := testhelpers.MustNewRequest(http.MethodGet, "http://localhost", nil) fakeBody := generateBytes(gziphandler.DefaultMinSize) next := http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) { _, err := rw.Write(fakeBody) if err != nil { http.Error(rw, err.Error(), http.StatusInternalServerError) } }) handler := &compress{next: next} rw := httptest.NewRecorder() handler.ServeHTTP(rw, req) assert.Empty(t, rw.Header().Get(contentEncodingHeader)) assert.EqualValues(t, rw.Body.Bytes(), fakeBody) } func TestShouldNotCompressWhenSpecificContentType(t *testing.T) { baseBody := generateBytes(gziphandler.DefaultMinSize) next := http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) { _, err := rw.Write(baseBody) if err != nil { http.Error(rw, err.Error(), http.StatusInternalServerError) } }) testCases := []struct { desc string conf dynamic.Compress reqContentType string }{ { desc: "text/event-stream", conf: dynamic.Compress{ ExcludedContentTypes: []string{"text/event-stream"}, }, reqContentType: "text/event-stream", }, { desc: "application/grpc", conf: dynamic.Compress{}, reqContentType: "application/grpc", }, } for _, test := range testCases { test := test t.Run(test.desc, func(t *testing.T) { t.Parallel() req := testhelpers.MustNewRequest(http.MethodGet, "http://localhost", nil) req.Header.Add(acceptEncodingHeader, gzipValue) if test.reqContentType != "" { req.Header.Add(contentTypeHeader, test.reqContentType) } handler, err := New(context.Background(), next, test.conf, "test") require.NoError(t, err) rw := httptest.NewRecorder() handler.ServeHTTP(rw, req) assert.Empty(t, rw.Header().Get(acceptEncodingHeader)) assert.Empty(t, rw.Header().Get(contentEncodingHeader)) assert.EqualValues(t, rw.Body.Bytes(), baseBody) }) } } func TestIntegrationShouldNotCompress(t *testing.T) { fakeCompressedBody := generateBytes(100000) testCases := []struct { name string handler http.Handler expectedStatusCode int }{ { name: "when content already compressed", handler: http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) { rw.Header().Add(contentEncodingHeader, gzipValue) rw.Header().Add(varyHeader, acceptEncodingHeader) _, err := rw.Write(fakeCompressedBody) if err != nil { http.Error(rw, err.Error(), http.StatusInternalServerError) } }), expectedStatusCode: http.StatusOK, }, { name: "when content already compressed and status code Created", handler: http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) { rw.Header().Add(contentEncodingHeader, gzipValue) rw.Header().Add(varyHeader, acceptEncodingHeader) rw.WriteHeader(http.StatusCreated) _, err := rw.Write(fakeCompressedBody) if err != nil { http.Error(rw, err.Error(), http.StatusInternalServerError) } }), expectedStatusCode: http.StatusCreated, }, } for _, test := range testCases { t.Run(test.name, func(t *testing.T) { compress := &compress{next: test.handler} ts := httptest.NewServer(compress) defer ts.Close() req := testhelpers.MustNewRequest(http.MethodGet, ts.URL, nil) req.Header.Add(acceptEncodingHeader, gzipValue) resp, err := http.DefaultClient.Do(req) require.NoError(t, err) assert.Equal(t, test.expectedStatusCode, resp.StatusCode) assert.Equal(t, gzipValue, resp.Header.Get(contentEncodingHeader)) assert.Equal(t, acceptEncodingHeader, resp.Header.Get(varyHeader)) body, err := ioutil.ReadAll(resp.Body) require.NoError(t, err) assert.EqualValues(t, fakeCompressedBody, body) }) } } func TestShouldWriteHeaderWhenFlush(t *testing.T) { next := http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) { rw.Header().Add(contentEncodingHeader, gzipValue) rw.Header().Add(varyHeader, acceptEncodingHeader) rw.WriteHeader(http.StatusUnauthorized) rw.(http.Flusher).Flush() _, err := rw.Write([]byte("short")) if err != nil { http.Error(rw, err.Error(), http.StatusInternalServerError) } }) handler := &compress{next: next} ts := httptest.NewServer(handler) defer ts.Close() req := testhelpers.MustNewRequest(http.MethodGet, ts.URL, nil) req.Header.Add(acceptEncodingHeader, gzipValue) resp, err := http.DefaultClient.Do(req) require.NoError(t, err) assert.Equal(t, http.StatusUnauthorized, resp.StatusCode) assert.Equal(t, gzipValue, resp.Header.Get(contentEncodingHeader)) assert.Equal(t, acceptEncodingHeader, resp.Header.Get(varyHeader)) } func TestIntegrationShouldCompress(t *testing.T) { fakeBody := generateBytes(100000) testCases := []struct { name string handler http.Handler expectedStatusCode int }{ { name: "when AcceptEncoding header is present", handler: http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) { _, err := rw.Write(fakeBody) if err != nil { http.Error(rw, err.Error(), http.StatusInternalServerError) } }), expectedStatusCode: http.StatusOK, }, { name: "when AcceptEncoding header is present and status code Created", handler: http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) { rw.WriteHeader(http.StatusCreated) _, err := rw.Write(fakeBody) if err != nil { http.Error(rw, err.Error(), http.StatusInternalServerError) } }), expectedStatusCode: http.StatusCreated, }, } for _, test := range testCases { t.Run(test.name, func(t *testing.T) { compress := &compress{next: test.handler} ts := httptest.NewServer(compress) defer ts.Close() req := testhelpers.MustNewRequest(http.MethodGet, ts.URL, nil) req.Header.Add(acceptEncodingHeader, gzipValue) resp, err := http.DefaultClient.Do(req) require.NoError(t, err) assert.Equal(t, test.expectedStatusCode, resp.StatusCode) assert.Equal(t, gzipValue, resp.Header.Get(contentEncodingHeader)) assert.Equal(t, acceptEncodingHeader, resp.Header.Get(varyHeader)) body, err := ioutil.ReadAll(resp.Body) require.NoError(t, err) if assert.ObjectsAreEqualValues(body, fakeBody) { assert.Fail(t, "expected a compressed body", "got %v", body) } }) } } func generateBytes(len int) []byte { var value []byte for i := 0; i < len; i++ { value = append(value, 0x61+byte(i)) } return value }