From dee64ddddf7637cd9898b995da572ce348b27443 Mon Sep 17 00:00:00 2001
From: lqqyt2423 <974923609@qq.com>
Date: Mon, 6 Jun 2022 17:46:39 +0800
Subject: [PATCH] add ClientConnected hook

---
 addon/addon.go           | 10 ++++++++++
 connection/connection.go | 19 +++++++++++++++++++
 proxy/proxy.go           |  9 +++++++++
 3 files changed, 38 insertions(+)
 create mode 100644 connection/connection.go

diff --git a/addon/addon.go b/addon/addon.go
index d1dd587..1ff9b88 100644
--- a/addon/addon.go
+++ b/addon/addon.go
@@ -3,6 +3,7 @@ package addon
 import (
 	"time"
 
+	"github.com/lqqyt2423/go-mitmproxy/connection"
 	"github.com/lqqyt2423/go-mitmproxy/flow"
 	_log "github.com/sirupsen/logrus"
 )
@@ -10,6 +11,9 @@ import (
 var log = _log.WithField("at", "addon")
 
 type Addon interface {
+	// A client has connected to mitmproxy. Note that a connection can correspond to multiple HTTP requests.
+	ClientConnected(*connection.Client)
+
 	// HTTP request headers were successfully read. At this point, the body is empty.
 	Requestheaders(*flow.Flow)
 
@@ -26,6 +30,8 @@ type Addon interface {
 // Base do nothing
 type Base struct{}
 
+func (addon *Base) ClientConnected(*connection.Client) {}
+
 func (addon *Base) Requestheaders(*flow.Flow)  {}
 func (addon *Base) Request(*flow.Flow)         {}
 func (addon *Base) Responseheaders(*flow.Flow) {}
@@ -36,6 +42,10 @@ type Log struct {
 	Base
 }
 
+func (addon *Log) ClientConnected(client *connection.Client) {
+	log.Infof("%v client connect\n", client.Conn.RemoteAddr())
+}
+
 func (addon *Log) Requestheaders(f *flow.Flow) {
 	log := log.WithField("in", "Log")
 	start := time.Now()
diff --git a/connection/connection.go b/connection/connection.go
new file mode 100644
index 0000000..25a843b
--- /dev/null
+++ b/connection/connection.go
@@ -0,0 +1,19 @@
+package connection
+
+import (
+	"net"
+
+	uuid "github.com/satori/go.uuid"
+)
+
+type Client struct {
+	Id   uuid.UUID
+	Conn net.Conn
+}
+
+func NewClient(c net.Conn) *Client {
+	return &Client{
+		Id:   uuid.NewV4(),
+		Conn: c,
+	}
+}
diff --git a/proxy/proxy.go b/proxy/proxy.go
index f4b840b..140e856 100644
--- a/proxy/proxy.go
+++ b/proxy/proxy.go
@@ -2,6 +2,7 @@ package proxy
 
 import (
 	"bytes"
+	"context"
 	"crypto/tls"
 	"io"
 	"net"
@@ -9,6 +10,7 @@ import (
 	"time"
 
 	"github.com/lqqyt2423/go-mitmproxy/addon"
+	"github.com/lqqyt2423/go-mitmproxy/connection"
 	"github.com/lqqyt2423/go-mitmproxy/flow"
 	_log "github.com/sirupsen/logrus"
 )
@@ -39,6 +41,13 @@ 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)
+			}
+			return ctx
+		},
 	}
 
 	proxy.Client = &http.Client{