From 46f165105dc123d32b808cf3b9f1bf0c891c0a01 Mon Sep 17 00:00:00 2001
From: lqqyt2423 <974923609@qq.com>
Date: Mon, 13 Jun 2022 12:40:42 +0800
Subject: [PATCH] http add server hook

---
 addon/addon.go      | 16 +++++++++++
 flow/conncontext.go | 20 ++++---------
 proxy/proxy.go      | 70 ++++++++++++++++++++++++++++++++++++++-------
 3 files changed, 81 insertions(+), 25 deletions(-)

diff --git a/addon/addon.go b/addon/addon.go
index b4a9f76..5b12c85 100644
--- a/addon/addon.go
+++ b/addon/addon.go
@@ -17,6 +17,12 @@ type Addon interface {
 	// A client connection has been closed (either by us or the client).
 	ClientDisconnected(*connection.Client)
 
+	// Mitmproxy has connected to a server.
+	ServerConnected(*flow.ConnContext)
+
+	// A server connection has been closed (either by us or the server).
+	ServerDisconnected(*flow.ConnContext)
+
 	// HTTP request headers were successfully read. At this point, the body is empty.
 	Requestheaders(*flow.Flow)
 
@@ -35,6 +41,8 @@ type Base struct{}
 
 func (addon *Base) ClientConnected(*connection.Client)    {}
 func (addon *Base) ClientDisconnected(*connection.Client) {}
+func (addon *Base) ServerConnected(*flow.ConnContext)     {}
+func (addon *Base) ServerDisconnected(*flow.ConnContext)  {}
 
 func (addon *Base) Requestheaders(*flow.Flow)  {}
 func (addon *Base) Request(*flow.Flow)         {}
@@ -54,6 +62,14 @@ func (addon *Log) ClientDisconnected(client *connection.Client) {
 	log.Infof("%v client disconnect\n", client.Conn.RemoteAddr())
 }
 
+func (addon *Log) ServerConnected(connCtx *flow.ConnContext) {
+	log.Infof("%v server connect %v\n", connCtx.Client.Conn.RemoteAddr(), connCtx.Server.Conn.RemoteAddr())
+}
+
+func (addon *Log) ServerDisconnected(connCtx *flow.ConnContext) {
+	log.Infof("%v server disconnect %v\n", connCtx.Client.Conn.RemoteAddr(), connCtx.Server.Conn.RemoteAddr())
+}
+
 func (addon *Log) Requestheaders(f *flow.Flow) {
 	log := log.WithField("in", "Log")
 	start := time.Now()
diff --git a/flow/conncontext.go b/flow/conncontext.go
index 1335864..cbdfeb1 100644
--- a/flow/conncontext.go
+++ b/flow/conncontext.go
@@ -23,16 +23,7 @@ func NewConnContext(c net.Conn) *ConnContext {
 	}
 }
 
-type serverConn struct {
-	net.Conn
-}
-
-func (c *serverConn) Close() error {
-	log.Debugln("in http serverConn close")
-	return c.Conn.Close()
-}
-
-func (connCtx *ConnContext) InitHttpServer(SslInsecure bool) {
+func (connCtx *ConnContext) InitHttpServer(sslInsecure bool, connWrap func(net.Conn) net.Conn, whenServerConnected func()) {
 	if connCtx.Server != nil {
 		return
 	}
@@ -55,15 +46,16 @@ func (connCtx *ConnContext) InitHttpServer(SslInsecure bool) {
 					return nil, err
 				}
 
-				cw := &serverConn{c}
+				cw := connWrap(c)
 				server.Conn = cw
+				defer whenServerConnected()
 				return cw, nil
 			},
 			ForceAttemptHTTP2: false, // disable http2
 
 			DisableCompression: true, // To get the original response from the server, set Transport.DisableCompression to true.
 			TLSClientConfig: &tls.Config{
-				InsecureSkipVerify: SslInsecure,
+				InsecureSkipVerify: sslInsecure,
 				KeyLogWriter:       GetTlsKeyLogWriter(),
 			},
 		},
@@ -75,7 +67,7 @@ func (connCtx *ConnContext) InitHttpServer(SslInsecure bool) {
 	connCtx.Server = server
 }
 
-func (connCtx *ConnContext) InitHttpsServer(SslInsecure bool) {
+func (connCtx *ConnContext) InitHttpsServer(sslInsecure bool) {
 	if connCtx.Server != nil {
 		return
 	}
@@ -97,7 +89,7 @@ func (connCtx *ConnContext) InitHttpsServer(SslInsecure bool) {
 
 			DisableCompression: true, // To get the original response from the server, set Transport.DisableCompression to true.
 			TLSClientConfig: &tls.Config{
-				InsecureSkipVerify: SslInsecure,
+				InsecureSkipVerify: sslInsecure,
 				KeyLogWriter:       GetTlsKeyLogWriter(),
 			},
 		},
diff --git a/proxy/proxy.go b/proxy/proxy.go
index 8cad148..8901288 100644
--- a/proxy/proxy.go
+++ b/proxy/proxy.go
@@ -6,7 +6,6 @@ import (
 	"io"
 	"net"
 	"net/http"
-	"sync"
 
 	"github.com/lqqyt2423/go-mitmproxy/addon"
 	"github.com/lqqyt2423/go-mitmproxy/flow"
@@ -50,21 +49,56 @@ func (l *proxyListener) Accept() (net.Conn, error) {
 
 type proxyConn struct {
 	net.Conn
-	proxy     *Proxy
-	connCtx   *flow.ConnContext
-	closeOnce sync.Once
+	proxy    *Proxy
+	connCtx  *flow.ConnContext
+	closed   bool
+	closeErr error
 }
 
 func (c *proxyConn) Close() error {
 	log.Debugln("in proxyConn close")
+	if c.closed {
+		return c.closeErr
+	}
 
-	c.closeOnce.Do(func() {
-		for _, addon := range c.proxy.Addons {
-			addon.ClientDisconnected(c.connCtx.Client)
-		}
-	})
+	c.closed = true
+	c.closeErr = c.Conn.Close()
+
+	for _, addon := range c.proxy.Addons {
+		addon.ClientDisconnected(c.connCtx.Client)
+	}
+
+	if c.connCtx.Server != nil && c.connCtx.Server.Conn != nil {
+		c.connCtx.Server.Conn.Close()
+	}
+
+	return c.closeErr
+}
 
-	return c.Conn.Close()
+type serverConn struct {
+	net.Conn
+	proxy    *Proxy
+	connCtx  *flow.ConnContext
+	closed   bool
+	closeErr error
+}
+
+func (c *serverConn) Close() error {
+	log.Debugln("in http serverConn close")
+	if c.closed {
+		return c.closeErr
+	}
+
+	c.closed = true
+	c.closeErr = c.Conn.Close()
+
+	for _, addon := range c.proxy.Addons {
+		addon.ServerDisconnected(c.connCtx)
+	}
+
+	c.connCtx.Client.Conn.Close()
+
+	return c.closeErr
 }
 
 func NewProxy(opts *Options) (*Proxy, error) {
@@ -246,7 +280,21 @@ func (proxy *Proxy) ServeHTTP(res http.ResponseWriter, req *http.Request) {
 		}
 	}
 
-	f.ConnContext.InitHttpServer(proxy.Opts.SslInsecure)
+	f.ConnContext.InitHttpServer(
+		proxy.Opts.SslInsecure,
+		func(c net.Conn) net.Conn {
+			return &serverConn{
+				Conn:    c,
+				proxy:   proxy,
+				connCtx: f.ConnContext,
+			}
+		},
+		func() {
+			for _, addon := range proxy.Addons {
+				addon.ServerConnected(f.ConnContext)
+			}
+		},
+	)
 
 	proxyRes, err := f.ConnContext.Server.Client.Do(proxyReq)
 	if err != nil {