Skip to content

Commit 6dd56dc

Browse files
rosylillyrosylilly
authored
Merge pull request #12 from isucon/support-multiple-content-encoding
Support multiple content encoding
2 parents b394848 + 66ced87 commit 6dd56dc

File tree

5 files changed

+177
-25
lines changed

5 files changed

+177
-25
lines changed

‎agent/agent.go

+2-1
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,8 @@ var (
2828
)
2929

3030
var (
31-
ErrTransportInvalid = errors.New("Specify transport option(WithCloneTransport or WithDefaultTransport)")
31+
ErrTransportInvalid = errors.New("Specify transport option(WithCloneTransport or WithDefaultTransport)")
32+
ErrUnknownContentEncoding = errors.New("Unknown content encoding")
3233
)
3334

3435
func init() {

‎agent/decompress.go

+31-23
Original file line numberDiff line numberDiff line change
@@ -3,30 +3,42 @@ package agent
33
import (
44
"compress/flate"
55
"compress/gzip"
6+
"fmt"
67
"io"
78
"net/http"
9+
"strings"
810

911
"github.com/dsnet/compress/brotli"
1012
)
1113

1214
func decompress(res *http.Response) (*http.Response, error) {
13-
ce := res.Header.Get("Content-Encoding")
14-
15-
var err error
16-
var body io.ReadCloser
17-
switch ce {
18-
case "br":
19-
body, err = brotli.NewReader(res.Body, &brotli.ReaderConfig{})
20-
case "gzip":
21-
body = &gzipReader{body: res.Body}
22-
case "deflate":
23-
body = flate.NewReader(res.Body)
24-
default:
15+
contentEncoding := res.Header.Get("Content-Encoding")
16+
if contentEncoding == "" {
2517
return res, nil
2618
}
2719

28-
if err != nil {
29-
return nil, err
20+
var err error
21+
var body io.ReadCloser = res.Body
22+
23+
encodings := strings.Split(contentEncoding, ",")
24+
for i := len(encodings) - 1; i >= 0; i-- {
25+
encoding := encodings[i]
26+
switch strings.TrimSpace(encoding) {
27+
case "br":
28+
body, err = brotli.NewReader(body, &brotli.ReaderConfig{})
29+
case "gzip":
30+
body = &gzipReader{body: body}
31+
case "deflate":
32+
body = flate.NewReader(body)
33+
case "identity", "":
34+
// nop
35+
default:
36+
err = fmt.Errorf("unknown content encoding: %s: %w", encoding, ErrUnknownContentEncoding)
37+
}
38+
39+
if err != nil {
40+
return nil, err
41+
}
3042
}
3143

3244
res.Header.Del("Content-Length")
@@ -43,19 +55,15 @@ type gzipReader struct {
4355
zerr error
4456
}
4557

46-
func (gz *gzipReader) Read(p []byte) (n int, err error) {
58+
func (gz *gzipReader) Read(p []byte) (int, error) {
4759
if gz.zr == nil {
48-
if gz.zerr == nil {
49-
gz.zr, gz.zerr = gzip.NewReader(gz.body)
50-
}
51-
if gz.zerr != nil {
52-
return 0, gz.zerr
60+
var err error
61+
gz.zr, err = gzip.NewReader(gz.body)
62+
if err != nil {
63+
return 0, err
5364
}
5465
}
5566

56-
if err != nil {
57-
return 0, err
58-
}
5967
return gz.zr.Read(p)
6068
}
6169

‎agent/decompress_test.go

+128
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import (
55
"compress/flate"
66
"compress/gzip"
77
"context"
8+
"errors"
89
"io"
910
"io/ioutil"
1011
"net/http"
@@ -82,6 +83,41 @@ func newCompressHTTPServer() *httptest.Server {
8283
w.WriteHeader(200)
8384
io.WriteString(w, "test it")
8485
})
86+
r.GET("/identity", func(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
87+
w.Header().Set("Content-Type", "text/plain")
88+
w.Header().Set("Content-Encoding", "identity")
89+
w.Header().Set("X-Content-Type-Options", "nosniff")
90+
w.Header().Set("Transfer-Encoding", "chunked")
91+
w.WriteHeader(200)
92+
io.WriteString(w, "test it")
93+
})
94+
r.GET("/multiple", func(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
95+
gw := gzip.NewWriter(w)
96+
defer gw.Close()
97+
98+
fw, err := flate.NewWriter(gw, 9)
99+
if err != nil {
100+
io.WriteString(w, err.Error())
101+
w.WriteHeader(http.StatusInternalServerError)
102+
return
103+
}
104+
defer fw.Close()
105+
106+
w.Header().Set("Content-Type", "text/plain")
107+
w.Header().Set("Content-Encoding", "deflate, gzip")
108+
w.Header().Set("X-Content-Type-Options", "nosniff")
109+
w.Header().Set("Transfer-Encoding", "chunked")
110+
w.WriteHeader(200)
111+
io.WriteString(fw, "test it")
112+
})
113+
r.GET("/unknown", func(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
114+
w.Header().Set("Content-Type", "text/plain")
115+
w.Header().Set("Content-Encoding", "unknown")
116+
w.Header().Set("X-Content-Type-Options", "nosniff")
117+
w.Header().Set("Transfer-Encoding", "chunked")
118+
w.WriteHeader(200)
119+
io.WriteString(w, "test it")
120+
})
85121

86122
return httptest.NewServer(r)
87123
}
@@ -217,6 +253,98 @@ func TestDeflateResponse(t *testing.T) {
217253
}
218254
}
219255

256+
func TestIdentityResponse(t *testing.T) {
257+
srv := newCompressHTTPServer()
258+
defer srv.Close()
259+
260+
agent, err := NewAgent(WithBaseURL(srv.URL), WithDefaultTransport())
261+
if err != nil {
262+
t.Fatalf("%+v", err)
263+
}
264+
265+
req, err := agent.GET("/identity")
266+
if err != nil {
267+
t.Fatalf("%+v", err)
268+
}
269+
270+
res, err := agent.Do(context.Background(), req)
271+
if err != nil {
272+
t.Fatalf("%+v", err)
273+
}
274+
275+
if res.StatusCode != 200 {
276+
t.Fatalf("%#v", res)
277+
}
278+
defer res.Body.Close()
279+
280+
body, err := ioutil.ReadAll(res.Body)
281+
if err != nil {
282+
t.Fatalf("%+v", err)
283+
}
284+
285+
if bytes.Compare(body, []byte("test it")) != 0 {
286+
t.Fatalf("%s missmatch %s", body, "test it")
287+
}
288+
}
289+
290+
func TestMultipleResponse(t *testing.T) {
291+
srv := newCompressHTTPServer()
292+
defer srv.Close()
293+
294+
agent, err := NewAgent(WithBaseURL(srv.URL), WithDefaultTransport())
295+
if err != nil {
296+
t.Fatalf("%+v", err)
297+
}
298+
299+
req, err := agent.GET("/multiple")
300+
if err != nil {
301+
t.Fatalf("%+v", err)
302+
}
303+
304+
res, err := agent.Do(context.Background(), req)
305+
if err != nil {
306+
t.Fatalf("%+v", err)
307+
}
308+
309+
if res.StatusCode != 200 {
310+
t.Fatalf("%#v", res)
311+
}
312+
defer res.Body.Close()
313+
314+
body, err := ioutil.ReadAll(res.Body)
315+
if err != nil {
316+
t.Fatalf("%+v", err)
317+
}
318+
319+
if bytes.Compare(body, []byte("test it")) != 0 {
320+
t.Fatalf("%s missmatch %s", body, "test it")
321+
}
322+
}
323+
324+
func TestUnknownResponse(t *testing.T) {
325+
srv := newCompressHTTPServer()
326+
defer srv.Close()
327+
328+
agent, err := NewAgent(WithBaseURL(srv.URL), WithDefaultTransport())
329+
if err != nil {
330+
t.Fatalf("%+v", err)
331+
}
332+
333+
req, err := agent.GET("/unknown")
334+
if err != nil {
335+
t.Fatalf("%+v", err)
336+
}
337+
338+
_, err = agent.Do(context.Background(), req)
339+
if err == nil {
340+
t.Fatalf("expected error but err is nil, %+v", err)
341+
} else {
342+
if !errors.Is(err, ErrUnknownContentEncoding) {
343+
t.Fatalf("%+v", err)
344+
}
345+
}
346+
}
347+
220348
func TestWithEcho(t *testing.T) {
221349
e := echo.New()
222350
e.GET("/", func(c echo.Context) error {

‎agent/option.go

+5-1
Original file line numberDiff line numberDiff line change
@@ -43,8 +43,12 @@ func WithTimeout(d time.Duration) AgentOption {
4343
}
4444

4545
func WithDefaultTransport() AgentOption {
46+
return WithTransport(DefaultTransport)
47+
}
48+
49+
func WithTransport(trs *http.Transport) AgentOption {
4650
return func(a *Agent) error {
47-
a.HttpClient.Transport = DefaultTransport
51+
a.HttpClient.Transport = trs
4852

4953
return nil
5054
}

‎agent/option_test.go

+11
Original file line numberDiff line numberDiff line change
@@ -102,6 +102,17 @@ func TestDefaultTransport(t *testing.T) {
102102
assert.Same(t, agent1.HttpClient.Transport, agent2.HttpClient.Transport)
103103
}
104104

105+
func TestTransport(t *testing.T) {
106+
trs := DefaultTransport.Clone()
107+
108+
agent1, err := NewAgent(WithTransport(trs))
109+
if err != nil {
110+
t.Fatal(err)
111+
}
112+
113+
assert.Same(t, agent1.HttpClient.Transport, trs)
114+
}
115+
105116
func TestCloneTransport(t *testing.T) {
106117
agent1, err := NewAgent(WithCloneTransport(DefaultTransport))
107118
if err != nil {

0 commit comments

Comments
 (0)