diff --git a/addon/addon.go b/addon/addon.go index de560c3..b4a9f76 100644 --- a/addon/addon.go +++ b/addon/addon.go @@ -67,6 +67,6 @@ func (addon *Log) Requestheaders(f *flow.Flow) { if f.Response != nil && f.Response.Body != nil { contentLen = len(f.Response.Body) } - log.Infof("%v %v %v %v - %v ms\n", f.Request.Method, f.Request.URL.String(), StatusCode, contentLen, time.Since(start).Milliseconds()) + log.Infof("%v %v %v %v %v - %v ms\n", f.ConnContext.Client.Conn.RemoteAddr(), f.Request.Method, f.Request.URL.String(), StatusCode, contentLen, time.Since(start).Milliseconds()) }() } diff --git a/flow/conncontext.go b/flow/conncontext.go new file mode 100644 index 0000000..175c8dc --- /dev/null +++ b/flow/conncontext.go @@ -0,0 +1,9 @@ +package flow + +import "github.com/lqqyt2423/go-mitmproxy/connection" + +type ConnContext struct { + Client *connection.Client +} + +var ConnContextKey = new(struct{}) diff --git a/flow/flow.go b/flow/flow.go index ffc4453..52abc42 100644 --- a/flow/flow.go +++ b/flow/flow.go @@ -112,7 +112,8 @@ type Flow struct { Stream bool done chan struct{} - Id uuid.UUID + Id uuid.UUID + ConnContext *ConnContext } func (f *Flow) MarshalJSON() ([]byte, error) { diff --git a/proxy/middle.go b/proxy/middle.go index a5d093c..08fc4bf 100644 --- a/proxy/middle.go +++ b/proxy/middle.go @@ -2,6 +2,7 @@ package proxy import ( "bufio" + "context" "crypto/tls" "net" "net/http" @@ -9,6 +10,7 @@ import ( "time" "github.com/lqqyt2423/go-mitmproxy/cert" + "github.com/lqqyt2423/go-mitmproxy/flow" ) // 模拟了标准库中 server 运行,目的是仅通过当前进程内存转发 socket 数据,不需要经过 tcp 或 unix socket @@ -39,17 +41,19 @@ func newPipes(req *http.Request) (net.Conn, *connBuf) { // add Peek method for conn type connBuf struct { net.Conn - r *bufio.Reader - host string - remoteAddr string + r *bufio.Reader + host string + remoteAddr string + connContext *flow.ConnContext } func newConnBuf(c net.Conn, req *http.Request) *connBuf { return &connBuf{ - Conn: c, - r: bufio.NewReader(c), - host: req.Host, - remoteAddr: req.RemoteAddr, + Conn: c, + r: bufio.NewReader(c), + host: req.Host, + remoteAddr: req.RemoteAddr, + connContext: req.Context().Value(flow.ConnContextKey).(*flow.ConnContext), } } @@ -85,8 +89,13 @@ func NewMiddle(proxy *Proxy, caPath string) (Interceptor, error) { } server := &http.Server{ - Handler: m, - IdleTimeout: 5 * time.Second, + Handler: m, + IdleTimeout: 5 * time.Second, + + ConnContext: func(ctx context.Context, c net.Conn) context.Context { + return context.WithValue(ctx, flow.ConnContextKey, c.(*tls.Conn).NetConn().(*connBuf).connContext) + }, + TLSNextProto: make(map[string]func(*http.Server, *tls.Conn, http.Handler)), // disable http2 TLSConfig: &tls.Config{ GetCertificate: func(chi *tls.ClientHelloInfo) (*tls.Certificate, error) { diff --git a/proxy/proxy.go b/proxy/proxy.go index 2b59c16..90dc5c4 100644 --- a/proxy/proxy.go +++ b/proxy/proxy.go @@ -2,6 +2,7 @@ package proxy import ( "bytes" + "context" "crypto/tls" "io" "net" @@ -31,11 +32,7 @@ type Proxy struct { StreamLargeBodies int64 // 当请求或响应体大于此字节时,转为 stream 模式 Addons []addon.Addon - activeConn map[net.Conn]*proxyContext -} - -type proxyContext struct { - client *connection.Client + activeConn map[net.Conn]*flow.ConnContext } func NewProxy(opts *Options) (*Proxy, error) { @@ -46,13 +43,19 @@ func NewProxy(opts *Options) (*Proxy, error) { Addr: opts.Addr, Handler: proxy, IdleTimeout: 5 * time.Second, + + ConnContext: func(ctx context.Context, c net.Conn) context.Context { + client := connection.NewClient(c) + connCtx := &flow.ConnContext{ + Client: client, + } + proxy.activeConn[c] = connCtx + return context.WithValue(ctx, flow.ConnContextKey, connCtx) + }, + ConnState: func(c net.Conn, cs http.ConnState) { if cs == http.StateNew { - client := connection.NewClient(c) - proxy.activeConn[c] = &proxyContext{ - client, - } - + client := proxy.activeConn[c].Client for _, addon := range proxy.Addons { addon.ClientConnected(client) } @@ -101,7 +104,7 @@ func NewProxy(opts *Options) (*Proxy, error) { proxy.Addons = make([]addon.Addon, 0) - proxy.activeConn = make(map[net.Conn]*proxyContext) + proxy.activeConn = make(map[net.Conn]*flow.ConnContext) return proxy, nil } @@ -183,6 +186,7 @@ func (proxy *Proxy) ServeHTTP(res http.ResponseWriter, req *http.Request) { f := flow.NewFlow() f.Request = flow.NewRequest(req) + f.ConnContext = req.Context().Value(flow.ConnContextKey).(*flow.ConnContext) defer f.Finish() // trigger addon event Requestheaders @@ -323,7 +327,7 @@ func (proxy *Proxy) handleConnect(res http.ResponseWriter, req *http.Request) { } func (proxy *Proxy) whenClientConnClose(c net.Conn) { - client := proxy.activeConn[c].client + client := proxy.activeConn[c].Client for _, addon := range proxy.Addons { addon.ClientDisconnected(client) }