-
Notifications
You must be signed in to change notification settings - Fork 1.4k
/
compression.go
175 lines (155 loc) · 5.38 KB
/
compression.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
// Copyright The OpenTelemetry Authors
// SPDX-License-Identifier: Apache-2.0
// This file contains helper functions regarding compression/decompression for confighttp.
package confighttp // import "go.opentelemetry.io/collector/config/confighttp"
import (
"bytes"
"compress/gzip"
"compress/zlib"
"fmt"
"io"
"net/http"
"github.com/golang/snappy"
"github.com/klauspost/compress/zstd"
"go.opentelemetry.io/collector/config/configcompression"
)
type compressRoundTripper struct {
rt http.RoundTripper
compressionType configcompression.Type
compressor *compressor
}
func newCompressRoundTripper(rt http.RoundTripper, compressionType configcompression.Type) (*compressRoundTripper, error) {
encoder, err := newCompressor(compressionType)
if err != nil {
return nil, err
}
return &compressRoundTripper{
rt: rt,
compressionType: compressionType,
compressor: encoder,
}, nil
}
func (r *compressRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
if req.Header.Get(headerContentEncoding) != "" {
// If the header already specifies a content encoding then skip compression
// since we don't want to compress it again. This is a safeguard that normally
// should not happen since CompressRoundTripper is not intended to be used
// with http clients which already do their own compression.
return r.rt.RoundTrip(req)
}
// Compress the body.
buf := bytes.NewBuffer([]byte{})
if err := r.compressor.compress(buf, req.Body); err != nil {
return nil, err
}
// Create a new request since the docs say that we cannot modify the "req"
// (see https://golang.org/pkg/net/http/#RoundTripper).
cReq, err := http.NewRequestWithContext(req.Context(), req.Method, req.URL.String(), buf)
if err != nil {
return nil, err
}
// Clone the headers and add the encoding header.
cReq.Header = req.Header.Clone()
cReq.Header.Add(headerContentEncoding, string(r.compressionType))
return r.rt.RoundTrip(cReq)
}
type decompressor struct {
errHandler func(w http.ResponseWriter, r *http.Request, errorMsg string, statusCode int)
base http.Handler
decoders map[string]func(body io.ReadCloser) (io.ReadCloser, error)
}
// httpContentDecompressor offloads the task of handling compressed HTTP requests
// by identifying the compression format in the "Content-Encoding" header and re-writing
// request body so that the handlers further in the chain can work on decompressed data.
// It supports gzip and deflate/zlib compression.
func httpContentDecompressor(h http.Handler, eh func(w http.ResponseWriter, r *http.Request, errorMsg string, statusCode int), decoders map[string]func(body io.ReadCloser) (io.ReadCloser, error)) http.Handler {
errHandler := defaultErrorHandler
if eh != nil {
errHandler = eh
}
d := &decompressor{
errHandler: errHandler,
base: h,
decoders: map[string]func(body io.ReadCloser) (io.ReadCloser, error){
"": func(io.ReadCloser) (io.ReadCloser, error) {
// Not a compressed payload. Nothing to do.
return nil, nil
},
"gzip": func(body io.ReadCloser) (io.ReadCloser, error) {
gr, err := gzip.NewReader(body)
if err != nil {
return nil, err
}
return gr, nil
},
"zstd": func(body io.ReadCloser) (io.ReadCloser, error) {
zr, err := zstd.NewReader(
body,
// Concurrency 1 disables async decoding. We don't need async decoding, it is pointless
// for our use-case (a server accepting decoding http requests).
// Disabling async improves performance (I benchmarked it previously when working
// on https://github.com/open-telemetry/opentelemetry-collector-contrib/pull/23257).
zstd.WithDecoderConcurrency(1),
)
if err != nil {
return nil, err
}
return zr.IOReadCloser(), nil
},
"zlib": func(body io.ReadCloser) (io.ReadCloser, error) {
zr, err := zlib.NewReader(body)
if err != nil {
return nil, err
}
return zr, nil
},
"snappy": func(body io.ReadCloser) (io.ReadCloser, error) {
sr := snappy.NewReader(body)
sb := new(bytes.Buffer)
_, err := io.Copy(sb, sr)
if err != nil {
return nil, err
}
if err = body.Close(); err != nil {
return nil, err
}
return io.NopCloser(sb), nil
},
},
}
d.decoders["deflate"] = d.decoders["zlib"]
for key, dec := range decoders {
d.decoders[key] = dec
}
return d
}
func (d *decompressor) ServeHTTP(w http.ResponseWriter, r *http.Request) {
newBody, err := d.newBodyReader(r)
if err != nil {
d.errHandler(w, r, err.Error(), http.StatusBadRequest)
return
}
if newBody != nil {
defer newBody.Close()
// "Content-Encoding" header is removed to avoid decompressing twice
// in case the next handler(s) have implemented a similar mechanism.
r.Header.Del("Content-Encoding")
// "Content-Length" is set to -1 as the size of the decompressed body is unknown.
r.Header.Del("Content-Length")
r.ContentLength = -1
r.Body = newBody
}
d.base.ServeHTTP(w, r)
}
func (d *decompressor) newBodyReader(r *http.Request) (io.ReadCloser, error) {
encoding := r.Header.Get(headerContentEncoding)
decoder, ok := d.decoders[encoding]
if !ok {
return nil, fmt.Errorf("unsupported %s: %s", headerContentEncoding, encoding)
}
return decoder(r.Body)
}
// defaultErrorHandler writes the error message in plain text.
func defaultErrorHandler(w http.ResponseWriter, _ *http.Request, errMsg string, statusCode int) {
http.Error(w, errMsg, statusCode)
}