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
|
package proxy
import (
"context"
"net"
"sync"
)
// maxCaptureBytes bounds how much of any single request or response
// mitmux buffers for history storage, independent of how much data
// actually flows through the proxy. Proxying itself always streams the
// full body regardless of this limit - only what gets stored is capped,
// so a multi-gigabyte download can't be turned into a memory exhaustion
// vector just because the history view wants to remember it.
const maxCaptureBytes = 10 << 20 // 10 MiB
// teeConn wraps a net.Conn, recording every byte read off the wire (up
// to maxCaptureBytes) so it can be attributed to a specific request or
// response later. Take returns everything recorded since the last call
// and resets the buffer, so callers must take exactly once per message
// they want attributed correctly - see forward() for why that's safe
// here (call sites synchronize on the request/response boundary itself).
type teeConn struct {
net.Conn
mu sync.Mutex
buf []byte
}
func newTeeConn(c net.Conn) *teeConn {
return &teeConn{Conn: c}
}
func (c *teeConn) Read(p []byte) (int, error) {
n, err := c.Conn.Read(p)
if n > 0 {
c.mu.Lock()
if room := maxCaptureBytes - len(c.buf); room > 0 {
end := n
if end > room {
end = room
}
c.buf = append(c.buf, p[:end]...)
}
c.mu.Unlock()
}
return n, err
}
// Take returns the bytes read since the last Take call (or since the
// connection was created) and resets the buffer.
func (c *teeConn) Take() []byte {
c.mu.Lock()
defer c.mu.Unlock()
out := c.buf
c.buf = nil
return out
}
// teeListener wraps a net.Listener so every accepted connection is
// tee-captured.
type teeListener struct {
net.Listener
}
func (l *teeListener) Accept() (net.Conn, error) {
c, err := l.Listener.Accept()
if err != nil {
return nil, err
}
return newTeeConn(c), nil
}
type contextKey int
const clientTeeKey contextKey = iota
// teeConnFromContext returns the teeConn wrapping the client connection
// the current request was read from, as attached via http.Server's
// ConnContext hook. Returns nil for HTTP/2 client connections, which
// aren't tee-captured (see capture.go).
func teeConnFromContext(ctx context.Context) *teeConn {
tc, _ := ctx.Value(clientTeeKey).(*teeConn)
return tc
}
func withClientTee(ctx context.Context, c net.Conn) context.Context {
tc, ok := c.(*teeConn)
if !ok {
return ctx
}
return context.WithValue(ctx, clientTeeKey, tc)
}
|