From 5e6055a901108f40f4cb37d9876e9af96dc082f0 Mon Sep 17 00:00:00 2001
From: Carlo Alberto Ferraris <cafxx@strayorange.com>
Date: Tue, 25 Jan 2022 08:24:38 +0000
Subject: [PATCH] Handle Accept-Ranges/Range headers correctly

Fixes https://github.com/nytimes/gziphandler/issues/83
Fixes https://github.com/CAFxX/httpcompression/issues/6
---
 adapter.go         | 12 ++++++++++
 adapter_test.go    | 59 ++++++++++++++++++++++++++++++++++++++++++++++
 response_writer.go |  7 ++++++
 3 files changed, 78 insertions(+)

diff --git a/adapter.go b/adapter.go
index de29773..5bcab31 100644
--- a/adapter.go
+++ b/adapter.go
@@ -15,6 +15,7 @@ import (
 const (
 	vary            = "Vary"
 	acceptEncoding  = "Accept-Encoding"
+	acceptRanges    = "Accept-Ranges"
 	contentEncoding = "Content-Encoding"
 	contentType     = "Content-Type"
 	contentLength   = "Content-Length"
@@ -76,6 +77,17 @@ func Adapter(opts ...Option) (func(http.Handler) http.Handler, error) {
 				return
 			}
 
+			// We do not handle range requests when compression is used, as the
+			// range specified applies to the compressed data, not to the uncompressed one.
+			// So we would need to (1) ensure that compressors are deterministic and (2)
+			// generate the whole uncompressed response anyway, compress it, and then discard
+			// the bits outside of the range.
+			// Let's keep it simple, and simply ignore completely the range header.
+			// We also need to remove the Accept: Range header from any response that is
+			// compressed; this is done in the ResponseWriter.
+			// See https://github.com/nytimes/gziphandler/issues/83.
+			r.Header.Del("Range")
+
 			gw, _ := writerPool.Get().(*compressWriter)
 			if gw == nil {
 				gw = &compressWriter{}
diff --git a/adapter_test.go b/adapter_test.go
index 9700755..827b084 100644
--- a/adapter_test.go
+++ b/adapter_test.go
@@ -1045,6 +1045,65 @@ func TestWriteStringEquivalence(t *testing.T) {
 	}
 }
 
+func TestAcceptRanges(t *testing.T) {
+	// Tests for https://github.com/nytimes/gziphandler/issues/83 https://github.com/CAFxX/httpcompression/issues/6
+	t.Parallel()
+
+	cases := map[string]struct {
+		contentType    string
+		writeHeader    bool
+		_range         string
+		acceptEncoding string
+		body           string
+
+		expectRange           string
+		expectAcceptRanges    string
+		expectContentEncoding string
+	}{
+		// if the response is compressed, we do not support accept-ranges/range
+		"supported-encoding range":                      {"text/plain", false, "100-110bytes", "gzip", testBody, "", "", "gzip"},
+		"supported-encoding range explicit-writeheader": {"text/plain", true, "100-110bytes", "gzip", testBody, "", "", "gzip"},
+		// if the client does not accept one of the enabled encodings, we support accept-ranges/range
+		"unsupported-encoding range":                      {"text/plain", false, "100-110bytes", "unknown", testBody, "100-110bytes", "bytes", ""},
+		"unsupported-encoding range explicit-writeheader": {"text/plain", true, "100-110bytes", "unknown", testBody, "100-110bytes", "bytes", ""},
+		// if the content-type is not allowed to be compressed, we still strip the accept-ranges/range headers
+		// because we can't know this until the handler starts writing the response. See also the comments in adapter.go.
+		"not-whitelisted-type range":                      {"unknown/type", false, "100-110bytes", "gzip", testBody, "", "", ""},
+		"not-whitelisted-type range explicit-writeheader": {"unknown/type", true, "100-110bytes", "gzip", testBody, "", "", ""},
+	}
+
+	for n, c := range cases {
+		c := c
+		t.Run(n, func(t *testing.T) {
+			t.Parallel()
+
+			handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				assert.Equal(t, c.expectRange, r.Header.Get("Range"))
+				w.Header().Set(contentType, c.contentType)
+				w.Header().Set(acceptRanges, "bytes")
+				if c.writeHeader {
+					w.WriteHeader(http.StatusOK)
+				}
+				w.Write([]byte(c.body))
+			})
+
+			wrapper, err := DefaultAdapter(ContentTypes([]string{"text/plain"}, false))
+			assert.Nil(t, err, "DefaultAdapter returned error")
+
+			req, _ := http.NewRequest("GET", "/", nil)
+			req.Header.Set("Accept-Encoding", c.acceptEncoding)
+			req.Header.Set("Range", c._range)
+			resp := httptest.NewRecorder()
+			wrapper(handler).ServeHTTP(resp, req)
+			res := resp.Result()
+
+			assert.Equal(t, 200, res.StatusCode)
+			assert.Equal(t, c.expectAcceptRanges, res.Header.Get(acceptRanges))
+			assert.Equal(t, c.expectContentEncoding, res.Header.Get(contentEncoding))
+		})
+	}
+}
+
 // --------------------------------------------------------------------
 
 const (
diff --git a/response_writer.go b/response_writer.go
index 6b6e27c..fdd8365 100644
--- a/response_writer.go
+++ b/response_writer.go
@@ -146,6 +146,9 @@ func (w *compressWriter) startCompress(enc string) error {
 	// See: https://github.com/golang/go/issues/14975.
 	w.Header().Del(contentLength)
 
+	// See the comment about ranges in adapter.go
+	w.Header().Del(acceptRanges)
+
 	// Write the header to gzip response.
 	if w.code != 0 {
 		w.ResponseWriter.WriteHeader(w.code)
@@ -176,6 +179,10 @@ func (w *compressWriter) startCompress(enc string) error {
 
 // startPlain writes to sent bytes and buffer the underlying ResponseWriter without gzip.
 func (w *compressWriter) startPlain() error {
+	// See the comment about ranges in adapter.go; we need to do it even in this case
+	// because adapter will strip the range header anyway.
+	w.Header().Del(acceptRanges)
+
 	if w.code != 0 {
 		w.ResponseWriter.WriteHeader(w.code)
 		// Ensure that no other WriteHeader's happen