add ClientConnected hook

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

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,
+	}
+}

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{