From d2659e454c78d6e18694acf5e62d997cbc58af2d Mon Sep 17 00:00:00 2001
From: lqqyt2423 <974923609@qq.com>
Date: Tue, 7 Jun 2022 18:34:43 +0800
Subject: [PATCH] add flow.ConnContext

---
 addon/addon.go      |  2 +-
 flow/conncontext.go |  9 +++++++++
 flow/flow.go        |  3 ++-
 proxy/middle.go     | 27 ++++++++++++++++++---------
 proxy/proxy.go      | 28 ++++++++++++++++------------
 5 files changed, 46 insertions(+), 23 deletions(-)
 create mode 100644 flow/conncontext.go

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