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