...
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15 package proxy
16
17 import (
18 "bytes"
19 "fmt"
20 "io"
21 "net/http"
22 "strconv"
23 "sync"
24
25 "github.com/google/martian/v3"
26 )
27
28
29
30
31
32
33
34
35 const LogVersion = "0.2"
36
37
38 type Log struct {
39 Initial []byte
40 Version string
41 Converter *Converter
42 Entries []*Entry
43 }
44
45
46 type Entry struct {
47 ID string
48 Request *Request
49 Response *Response
50 }
51
52
53 type Request struct {
54 Method string
55 URL string
56 Header http.Header
57
58
59 MediaType string
60 BodyParts [][]byte
61 Trailer http.Header `json:",omitempty"`
62 }
63
64
65 type Response struct {
66 StatusCode int
67 Proto string
68 ProtoMajor int
69 ProtoMinor int
70 Header http.Header
71 Body []byte
72 Trailer http.Header `json:",omitempty"`
73 }
74
75
76 type Logger struct {
77 mu sync.Mutex
78 entries map[string]*Entry
79 log *Log
80 }
81
82
83 func newLogger() *Logger {
84 return &Logger{
85 log: &Log{
86 Version: LogVersion,
87 Converter: defaultConverter(),
88 },
89 entries: map[string]*Entry{},
90 }
91 }
92
93
94 func (l *Logger) ModifyRequest(req *http.Request) error {
95 if req.Method == "CONNECT" {
96 return nil
97 }
98 ctx := martian.NewContext(req)
99 if ctx.SkippingLogging() {
100 return nil
101 }
102 lreq, err := l.log.Converter.convertRequest(req)
103 if err != nil {
104 return err
105 }
106 id := ctx.ID()
107 entry := &Entry{ID: id, Request: lreq}
108
109 l.mu.Lock()
110 defer l.mu.Unlock()
111
112 if _, ok := l.entries[id]; ok {
113 panic(fmt.Sprintf("proxy: duplicate request ID: %s", id))
114 }
115 l.entries[id] = entry
116 l.log.Entries = append(l.log.Entries, entry)
117 return nil
118 }
119
120
121 func (l *Logger) ModifyResponse(res *http.Response) error {
122 ctx := martian.NewContext(res.Request)
123 if ctx.SkippingLogging() {
124 return nil
125 }
126 id := ctx.ID()
127 lres, err := l.log.Converter.convertResponse(res)
128 if err != nil {
129 return err
130 }
131
132 l.mu.Lock()
133 defer l.mu.Unlock()
134
135 if e, ok := l.entries[id]; ok {
136 e.Response = lres
137 }
138
139 return nil
140 }
141
142
143
144 func (l *Logger) Extract() *Log {
145 l.mu.Lock()
146 defer l.mu.Unlock()
147 r := l.log
148 l.log = nil
149 l.entries = nil
150 return r
151 }
152
153 func toHTTPResponse(lr *Response, req *http.Request) *http.Response {
154 res := &http.Response{
155 StatusCode: lr.StatusCode,
156 Proto: lr.Proto,
157 ProtoMajor: lr.ProtoMajor,
158 ProtoMinor: lr.ProtoMinor,
159 Header: lr.Header,
160 Body: io.NopCloser(bytes.NewReader(lr.Body)),
161 ContentLength: int64(len(lr.Body)),
162 }
163 res.Request = req
164
165
166 if req.Method == "HEAD" {
167 res.ContentLength = -1
168 if c := res.Header["Content-Length"]; len(c) == 1 {
169 if c64, err := strconv.ParseInt(c[0], 10, 64); err == nil {
170 res.ContentLength = c64
171 }
172 }
173 }
174 return res
175 }
176
View as plain text