From 1871fe4b535d0ad60f617b2592c9b24fba5a90f2 Mon Sep 17 00:00:00 2001
From: lqqyt2423 <974923609@qq.com>
Date: Tue, 7 Jun 2022 16:51:57 +0800
Subject: [PATCH] add ClientDisconnected hook

---
 addon/addon.go | 10 +++++++++-
 proxy/proxy.go | 40 +++++++++++++++++++++++++++++++++-------
 2 files changed, 42 insertions(+), 8 deletions(-)

diff --git a/addon/addon.go b/addon/addon.go
index 1ff9b88..de560c3 100644
--- a/addon/addon.go
+++ b/addon/addon.go
@@ -14,6 +14,9 @@ type Addon interface {
 	// A client has connected to mitmproxy. Note that a connection can correspond to multiple HTTP requests.
 	ClientConnected(*connection.Client)
 
+	// A client connection has been closed (either by us or the client).
+	ClientDisconnected(*connection.Client)
+
 	// HTTP request headers were successfully read. At this point, the body is empty.
 	Requestheaders(*flow.Flow)
 
@@ -30,7 +33,8 @@ type Addon interface {
 // Base do nothing
 type Base struct{}
 
-func (addon *Base) ClientConnected(*connection.Client) {}
+func (addon *Base) ClientConnected(*connection.Client)    {}
+func (addon *Base) ClientDisconnected(*connection.Client) {}
 
 func (addon *Base) Requestheaders(*flow.Flow)  {}
 func (addon *Base) Request(*flow.Flow)         {}
@@ -46,6 +50,10 @@ func (addon *Log) ClientConnected(client *connection.Client) {
 	log.Infof("%v client connect\n", client.Conn.RemoteAddr())
 }
 
+func (addon *Log) ClientDisconnected(client *connection.Client) {
+	log.Infof("%v client disconnect\n", client.Conn.RemoteAddr())
+}
+
 func (addon *Log) Requestheaders(f *flow.Flow) {
 	log := log.WithField("in", "Log")
 	start := time.Now()
diff --git a/proxy/proxy.go b/proxy/proxy.go
index 140e856..2b59c16 100644
--- a/proxy/proxy.go
+++ b/proxy/proxy.go
@@ -2,7 +2,6 @@ package proxy
 
 import (
 	"bytes"
-	"context"
 	"crypto/tls"
 	"io"
 	"net"
@@ -31,6 +30,12 @@ type Proxy struct {
 	Interceptor       Interceptor
 	StreamLargeBodies int64 // 当请求或响应体大于此字节时，转为 stream 模式
 	Addons            []addon.Addon
+
+	activeConn map[net.Conn]*proxyContext
+}
+
+type proxyContext struct {
+	client *connection.Client
 }
 
 func NewProxy(opts *Options) (*Proxy, error) {
@@ -41,12 +46,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)
-			for _, addon := range proxy.Addons {
-				addon.ClientConnected(client)
+		ConnState: func(c net.Conn, cs http.ConnState) {
+			if cs == http.StateNew {
+				client := connection.NewClient(c)
+				proxy.activeConn[c] = &proxyContext{
+					client,
+				}
+
+				for _, addon := range proxy.Addons {
+					addon.ClientConnected(client)
+				}
+			} else if cs == http.StateClosed {
+				proxy.whenClientConnClose(c)
 			}
-			return ctx
 		},
 	}
 
@@ -89,6 +101,8 @@ func NewProxy(opts *Options) (*Proxy, error) {
 
 	proxy.Addons = make([]addon.Addon, 0)
 
+	proxy.activeConn = make(map[net.Conn]*proxyContext)
+
 	return proxy, nil
 }
 
@@ -294,7 +308,10 @@ func (proxy *Proxy) handleConnect(res http.ResponseWriter, req *http.Request) {
 
 	cconn.(*net.TCPConn).SetLinger(0) // send RST other than FIN when finished, to avoid TIME_WAIT state
 	cconn.(*net.TCPConn).SetKeepAlive(false)
-	defer cconn.Close()
+	defer func() {
+		cconn.Close()
+		proxy.whenClientConnClose(cconn)
+	}()
 
 	_, err = io.WriteString(cconn, "HTTP/1.1 200 Connection Established\r\n\r\n")
 	if err != nil {
@@ -304,3 +321,12 @@ func (proxy *Proxy) handleConnect(res http.ResponseWriter, req *http.Request) {
 
 	Transfer(log, conn, cconn)
 }
+
+func (proxy *Proxy) whenClientConnClose(c net.Conn) {
+	client := proxy.activeConn[c].client
+	for _, addon := range proxy.Addons {
+		addon.ClientDisconnected(client)
+	}
+
+	delete(proxy.activeConn, c)
+}