add flow.ConnContext

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())
 	}()
 }

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{})

flow/flow.go

@@ -113,6 +113,7 @@ type Flow struct {
 	done   chan struct{}
 
 	Id          uuid.UUID
+	ConnContext *ConnContext
 }
 
 func (f *Flow) MarshalJSON() ([]byte, error) {

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
@@ -42,6 +44,7 @@ type connBuf struct {
 	r           *bufio.Reader
 	host        string
 	remoteAddr  string
+	connContext *flow.ConnContext
 }
 
 func newConnBuf(c net.Conn, req *http.Request) *connBuf {
@@ -50,6 +53,7 @@ func newConnBuf(c net.Conn, req *http.Request) *connBuf {
 		r:           bufio.NewReader(c),
 		host:        req.Host,
 		remoteAddr:  req.RemoteAddr,
+		connContext: req.Context().Value(flow.ConnContextKey).(*flow.ConnContext),
 	}
 }
 
@@ -87,6 +91,11 @@ func NewMiddle(proxy *Proxy, caPath string) (Interceptor, error) {
 	server := &http.Server{
 		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) {

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
+	activeConn map[net.Conn]*flow.ConnContext
-}
-
-type proxyContext struct {
-	client *connection.Client
 }
 
 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,
-		ConnState: func(c net.Conn, cs http.ConnState) {
+
-			if cs == http.StateNew {
+		ConnContext: func(ctx context.Context, c net.Conn) context.Context {
 			client := connection.NewClient(c)
-				proxy.activeConn[c] = &proxyContext{
+			connCtx := &flow.ConnContext{
-					client,
+				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 := 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)
 	}