...
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18 package proxy
19
20
21
22 import (
23 "crypto/tls"
24 "crypto/x509"
25 "encoding/json"
26 "fmt"
27 "net"
28 "net/http"
29 "net/url"
30 "os"
31 "strings"
32 "sync"
33 "time"
34
35 "github.com/google/martian/v3"
36 "github.com/google/martian/v3/fifo"
37 "github.com/google/martian/v3/httpspec"
38 "github.com/google/martian/v3/martianlog"
39 "github.com/google/martian/v3/mitm"
40 )
41
42
43 type Proxy struct {
44
45 CACert *x509.Certificate
46
47
48 URL *url.URL
49
50
51 Initial []byte
52
53 mproxy *martian.Proxy
54 filename string
55 logger *Logger
56 ignoreHeaders map[string]bool
57 }
58
59
60 func ForRecording(filename string, port int) (*Proxy, error) {
61 p, err := newProxy(filename)
62 if err != nil {
63 return nil, err
64 }
65
66
67
68 stack, _ := httpspec.NewStack("httpr")
69 p.mproxy.SetRequestModifier(stack)
70 p.mproxy.SetResponseModifier(stack)
71
72
73 logGroup := fifo.NewGroup()
74 skipAuth := skipLoggingByHost("accounts.google.com")
75 logGroup.AddRequestModifier(skipAuth)
76 logGroup.AddResponseModifier(skipAuth)
77 p.logger = newLogger()
78 logGroup.AddRequestModifier(p.logger)
79 logGroup.AddResponseModifier(p.logger)
80
81 stack.AddRequestModifier(logGroup)
82 stack.AddResponseModifier(logGroup)
83
84
85 logger := martianlog.NewLogger()
86 logger.SetDecode(true)
87 stack.AddRequestModifier(logger)
88 stack.AddResponseModifier(logger)
89
90 if err := p.start(port); err != nil {
91 return nil, err
92 }
93 return p, nil
94 }
95
96 var (
97 configOnce sync.Once
98 cert *x509.Certificate
99 config *mitm.Config
100 configErr error
101 )
102
103 func newProxy(filename string) (*Proxy, error) {
104 configOnce.Do(func() {
105
106
107 x509c, priv, err := mitm.NewAuthority("cloud.google.com/go/httpreplay", "HTTPReplay Authority", 100*time.Hour)
108 if err != nil {
109 configErr = err
110 return
111 }
112 cert = x509c
113 config, configErr = mitm.NewConfig(x509c, priv)
114 if config != nil {
115 config.SetValidity(100 * time.Hour)
116 config.SetOrganization("cloud.google.com/go/httpreplay")
117 config.SkipTLSVerify(false)
118 }
119 })
120 if configErr != nil {
121 return nil, configErr
122 }
123 mproxy := martian.NewProxy()
124 mproxy.SetMITM(config)
125 return &Proxy{
126 mproxy: mproxy,
127 CACert: cert,
128 filename: filename,
129 ignoreHeaders: map[string]bool{},
130 }, nil
131 }
132
133 func (p *Proxy) start(port int) error {
134 l, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", port))
135 if err != nil {
136 return err
137 }
138 p.URL = &url.URL{Scheme: "http", Host: l.Addr().String()}
139 go p.mproxy.Serve(l)
140 return nil
141 }
142
143
144 func (p *Proxy) Transport() *http.Transport {
145 caCertPool := x509.NewCertPool()
146 caCertPool.AddCert(p.CACert)
147 return &http.Transport{
148 TLSClientConfig: &tls.Config{RootCAs: caCertPool},
149 Proxy: func(*http.Request) (*url.URL, error) { return p.URL, nil },
150 }
151 }
152
153
154
155
156
157
158
159 func (p *Proxy) RemoveRequestHeaders(patterns []string) {
160 for _, pat := range patterns {
161 p.logger.log.Converter.registerRemoveRequestHeaders(pat)
162 }
163 }
164
165
166
167
168
169 func (p *Proxy) ClearHeaders(patterns []string) {
170 for _, pat := range patterns {
171 p.logger.log.Converter.registerClearHeaders(pat)
172 }
173 }
174
175
176
177
178
179
180
181 func (p *Proxy) RemoveQueryParams(patterns []string) {
182 for _, pat := range patterns {
183 p.logger.log.Converter.registerRemoveParams(pat)
184 }
185 }
186
187
188
189
190
191 func (p *Proxy) ClearQueryParams(patterns []string) {
192 for _, pat := range patterns {
193 p.logger.log.Converter.registerClearParams(pat)
194 }
195 }
196
197
198
199 func (p *Proxy) IgnoreHeader(h string) {
200 p.ignoreHeaders[http.CanonicalHeaderKey(h)] = true
201 }
202
203
204 func (p *Proxy) Close() error {
205 p.mproxy.Close()
206 if p.logger != nil {
207 return p.writeLog()
208 }
209 return nil
210 }
211
212 func (p *Proxy) writeLog() error {
213 lg := p.logger.Extract()
214 lg.Initial = p.Initial
215 bytes, err := json.MarshalIndent(lg, "", " ")
216 if err != nil {
217 return err
218 }
219 return os.WriteFile(p.filename, bytes, 0600)
220 }
221
222
223 type skipLoggingByHost string
224
225 func (s skipLoggingByHost) ModifyRequest(req *http.Request) error {
226 if strings.HasPrefix(req.Host, string(s)) {
227 martian.NewContext(req).SkipLogging()
228 }
229 return nil
230 }
231
232 func (s skipLoggingByHost) ModifyResponse(res *http.Response) error {
233 return s.ModifyRequest(res.Request)
234 }
235
View as plain text