forked from sourcegraph/apiproxy
/
client.go
93 lines (77 loc) · 2.74 KB
/
client.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
package apiproxy
import (
"net/http"
"net/textproto"
"regexp"
"strings"
"sync"
)
// RequestModifyingTransport is an implementation of net/http.RoundTripper that
// allows headers to be overwritten on requests matching certain predicates.
//
// It gives more control over HTTP requests (e.g., caching) when using libraries
// whose only HTTP configuration point is a http.Client or http.RoundTripper.
type RequestModifyingTransport struct {
overrides map[*regexp.Regexp]requestOverride
overridesMu sync.Mutex
// Transport is the underlying transport. If nil, net/http.DefaultTransport
// is used.
Transport http.RoundTripper
}
// RoundTrip implements net/http.RoundTripper.
func (t *RequestModifyingTransport) RoundTrip(req *http.Request) (resp *http.Response, err error) {
req = t.applyOverrides(req)
transport := t.Transport
if transport == nil {
transport = http.DefaultTransport
}
return transport.RoundTrip(req)
}
// requestOverride represents how a request should be modified by
// RequestModifyingTransport.
type requestOverride struct {
setHeaders http.Header
runOnlyOnce bool
}
// Override instructs the transport to set the specified headers (overwriting
// existing headers of the same name) on a GET or HEAD request whose request URI
// matches the regexp. If runOnlyOnce is true, the override will be deleted
// after execution (and won't affect any future requests); otherwise, it will
// remain in effect for the lifetime of the transport.
func (t *RequestModifyingTransport) Override(requestURI *regexp.Regexp, setHeaders http.Header, runOnlyOnce bool) {
t.overridesMu.Lock()
defer t.overridesMu.Unlock()
if t.overrides == nil {
t.overrides = make(map[*regexp.Regexp]requestOverride)
}
t.overrides[requestURI] = requestOverride{setHeaders, runOnlyOnce}
}
var NoCache = http.Header{"Cache-Control": []string{"no-cache"}}
// applyOverrides applies the transport's request overrides to req. If any
// overrides apply, req is cloned and the overrides are applied to the clone.
func (t *RequestModifyingTransport) applyOverrides(req *http.Request) *http.Request {
// Only override GET and HEAD requests, just to be safe. We may want to
// revisit this constraint later.
if method := strings.ToUpper(req.Method); method != "GET" && method != "HEAD" {
return req
}
requestURI := req.URL.RequestURI()
t.overridesMu.Lock()
defer t.overridesMu.Unlock()
cloned := false
for requestURIRegexp, override := range t.overrides {
if requestURIRegexp.MatchString(requestURI) {
if !cloned {
req = cloneRequest(req)
cloned = true
}
for name, val := range override.setHeaders {
req.Header[textproto.CanonicalMIMEHeaderKey(name)] = val
}
if override.runOnlyOnce {
delete(t.overrides, requestURIRegexp)
}
}
}
return req
}