fix: prevent double content-encoding

This commit is contained in:
Daniel G. Taylor 2022-05-04 14:17:27 -07:00
parent 790cb63129
commit 848f168792
No known key found for this signature in database
GPG key ID: 74AE195C5112E534
2 changed files with 38 additions and 0 deletions

View file

@ -38,6 +38,11 @@ func (w *contentEncodingWriter) Write(data []byte) (int, error) {
return w.writer.Write(data)
}
if w.Header().Get("Content-Encoding") != "" {
// Content encoding was already set, so we should ignore this!
return w.ResponseWriter.Write(data)
}
// Buffer the data until we can decide whether to compress it or not.
w.buf.Write(data)

View file

@ -1,7 +1,9 @@
package middleware
import (
"bytes"
"compress/gzip"
"crypto/rand"
"io/ioutil"
"net/http"
"net/http/httptest"
@ -200,3 +202,34 @@ func TestContentEncodingPanicBuffered(t *testing.T) {
// when the content encoding writer is closed.
assert.Equal(t, http.StatusInternalServerError, w.Result().StatusCode)
}
func TestContentEncodingIgnore(t *testing.T) {
app, _ := newTestRouter(t)
app.Resource("/").Get("test", "test",
responses.OK(),
).Run(func(ctx huma.Context) {
ctx.Header().Set("Content-Encoding", "br")
buf := make([]byte, 2500)
// Randomize so it can't compress well, which keeps the response
// large enough to potentially double-compress.
rand.Read(buf)
buf[0] = 'H'
buf[1] = 'e'
buf[2] = 'l'
buf[3] = 'l'
buf[4] = 'o'
w := brotli.NewWriter(ctx)
w.Write(buf)
w.Close()
})
w := httptest.NewRecorder()
req, _ := http.NewRequest(http.MethodGet, "/", nil)
req.Header.Add("Accept-Encoding", "br")
app.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Result().StatusCode)
out := make([]byte, 2500)
_, err := brotli.NewReader(bytes.NewReader(w.Body.Bytes())).Read(out)
assert.NoError(t, err)
assert.Equal(t, []byte{'H', 'e', 'l', 'l', 'o'}, out[:5])
}