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 {