-
Notifications
You must be signed in to change notification settings - Fork 11
/
recorder.go
132 lines (111 loc) · 3.26 KB
/
recorder.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
package httpx
import (
"bufio"
"bytes"
"fmt"
"net/http"
"net/http/httputil"
"net/url"
"github.com/go-chi/chi/v5/middleware"
"github.com/nyaruka/gocommon/dates"
"github.com/pkg/errors"
)
// Recorder is a utility for creating traces of HTTP requests being handled
type Recorder struct {
Trace *Trace
ResponseWriter http.ResponseWriter
responseBody *bytes.Buffer
}
// NewRecorder creates a new recorder for an HTTP request. If `originalRequest` is true, it tries to reconstruct the
// original request object.
func NewRecorder(r *http.Request, w http.ResponseWriter, reconstruct bool) (*Recorder, error) {
or := r
if reconstruct {
or = reconstructOriginal(r)
}
requestTrace, err := httputil.DumpRequest(or, true)
if err != nil {
return nil, errors.Wrap(err, "error dumping request")
}
// if we cloned the request above, DumpRequest will have drained the body and saved a copy on the reconstructed
// request, so put that copy on the passed request as well
r.Body = or.Body
responseBody := &bytes.Buffer{}
wrapped := middleware.NewWrapResponseWriter(w, r.ProtoMajor)
wrapped.Tee(responseBody)
return &Recorder{
Trace: &Trace{
Request: or,
RequestTrace: requestTrace,
StartTime: dates.Now(),
},
ResponseWriter: wrapped,
responseBody: responseBody,
}, nil
}
// End is called when the response has been written and generates the trace
func (r *Recorder) End() error {
wrapped := r.ResponseWriter.(middleware.WrapResponseWriter)
// build an approximation of headers part
responseTrace := &bytes.Buffer{}
responseTrace.WriteString(fmt.Sprintf("HTTP/1.1 %d %s\r\n", wrapped.Status(), http.StatusText(wrapped.Status())))
r.ResponseWriter.Header().Write(responseTrace)
responseTrace.WriteString("\r\n")
// and parse as response object
response, err := http.ReadResponse(bufio.NewReader(bytes.NewReader(responseTrace.Bytes())), r.Trace.Request)
if err != nil {
return errors.Wrap(err, "error reading response trace")
}
r.Trace.Response = response
r.Trace.ResponseTrace = responseTrace.Bytes()
r.Trace.ResponseBody = r.responseBody.Bytes()
r.Trace.EndTime = dates.Now()
return nil
}
// tries to reconstruct the original client request from the received server request.
func reconstructOriginal(r *http.Request) *http.Request {
// create copy of request as we'll be modifying the headers and URL
o := r.Clone(r.Context())
header := r.Header.Clone()
host := r.URL.Host
if host == "" {
host = r.Host
}
if h := r.Header.Get("Host"); h != "" {
host = h
}
if h := r.Header.Get("X-Forwarded-Host"); h != "" {
host = h
}
scheme := r.URL.Scheme
if scheme == "" {
scheme = "http"
}
if h := r.Header.Get("X-Forwarded-Proto"); h != "" {
scheme = h
}
path := r.RequestURI
if h := r.Header.Get("X-Forwarded-Path"); h != "" {
path = h
}
for _, h := range stripHeaders {
header.Del(h)
}
// if all that gives us a valid URL, replace it on the request
u, _ := url.Parse(fmt.Sprintf("%s://%s%s", scheme, host, path))
if u != nil {
o.URL = u
o.RequestURI = path
o.Header = header
}
return o
}
// headers to strip from reconstructed requests (these are nginx and ELB additions)
var stripHeaders = []string{
"X-Forwarded-Proto",
"X-Forwarded-Host",
"X-Forwarded-Port",
"X-Forwarded-Path",
"X-Forwarded-For",
"X-Amzn-Trace-Id",
}