-
Notifications
You must be signed in to change notification settings - Fork 11
/
context.go
119 lines (93 loc) · 3.26 KB
/
context.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
package filters
import (
"context"
"net"
"time"
)
type contextKey string
const (
ctxKeyDownstream = contextKey("downstream")
ctxKeyRequestNumber = contextKey("requestNumber")
ctxKeyMITMing = contextKey("mitming")
)
// Context is a wrapper for Context that exposes some additional
// information specific to its use in proxies.
type Context interface {
context.Context
// DownstreamConn retrieves the downstream connection from the given Context.
DownstreamConn() net.Conn
// RequestNumber indicates how many requests have been received on the current
// connection. The RequestNumber for the first request is 1, for the second is 2
// and so forth.
RequestNumber() int
// IncrementRequestNumber increments the request number by 1 and returns a new
// context.
IncrementRequestNumber() Context
// WithCancel mimics the method on context.Context
WithCancel() (Context, context.CancelFunc)
// WithDeadline mimics the method on context.Context
WithDeadline(deadline time.Time) (Context, context.CancelFunc)
// WithTimeout mimics the method on context.Context
WithTimeout(timeout time.Duration) (Context, context.CancelFunc)
// WithMITMing marks this context as being part of an MITM'ed connection.
WithMITMing() Context
// IsMITMing indicates whether or the proxy is MITMing the current connection.
IsMITMing() bool
// WithValue mimics the method on context.Context
WithValue(key, val interface{}) Context
}
// WrapContext wraps the given context.Context into a Context containing the
// given downstream net.Conn.
func WrapContext(ctx context.Context, downstream net.Conn) Context {
return (&ctext{ctx}).
WithValue(ctxKeyRequestNumber, 1).
WithValue(ctxKeyDownstream, func() net.Conn { return downstream })
}
// AdaptContext adapts a context.Context to the Context interface.
func AdaptContext(ctx context.Context) Context {
return &ctext{ctx}
}
// BackgroundContext creates a background Context without an associated
// connection.
func BackgroundContext() Context {
return WrapContext(context.Background(), nil)
}
// ctext implements Context
type ctext struct {
context.Context
}
func (ctx *ctext) DownstreamConn() net.Conn {
downstreamConn := ctx.Value(ctxKeyDownstream).(func() net.Conn)
if downstreamConn == nil {
return nil
}
return downstreamConn()
}
func (ctx *ctext) RequestNumber() int {
return ctx.Value(ctxKeyRequestNumber).(int)
}
func (ctx *ctext) IncrementRequestNumber() Context {
return ctx.WithValue(ctxKeyRequestNumber, ctx.RequestNumber()+1)
}
func (ctx *ctext) WithCancel() (Context, context.CancelFunc) {
result, cancel := context.WithCancel(ctx)
return &ctext{result}, cancel
}
func (ctx *ctext) WithDeadline(deadline time.Time) (Context, context.CancelFunc) {
result, cancel := context.WithDeadline(ctx, deadline)
return &ctext{result}, cancel
}
func (ctx *ctext) WithTimeout(timeout time.Duration) (Context, context.CancelFunc) {
result, cancel := context.WithTimeout(ctx, timeout)
return &ctext{result}, cancel
}
func (ctx *ctext) WithMITMing() Context {
return ctx.WithValue(ctxKeyMITMing, true)
}
func (ctx *ctext) IsMITMing() bool {
mitming := ctx.Value(ctxKeyMITMing)
return mitming != nil && mitming.(bool)
}
func (ctx *ctext) WithValue(key, val interface{}) Context {
return &ctext{context.WithValue(ctx, key, val)}
}