diff --git a/client/proxy.go b/client/proxy.go index 005be3a..a5da19e 100644 --- a/client/proxy.go +++ b/client/proxy.go @@ -426,7 +426,7 @@ func HandleTcpWorkConnection(localInfo *config.LocalSvrConf, proxyPlugin plugin. if proxyPlugin != nil { // if plugin is set, let plugin handle connections first workConn.Debug("handle by plugin: %s", proxyPlugin.Name()) - proxyPlugin.Handle(remote) + proxyPlugin.Handle(remote, workConn) workConn.Debug("handle by plugin finished") return } else { diff --git a/models/plugin/http_proxy.go b/models/plugin/http_proxy.go index f5fed6c..e5aefa8 100644 --- a/models/plugin/http_proxy.go +++ b/models/plugin/http_proxy.go @@ -106,13 +106,8 @@ func (hp *HttpProxy) Name() string { return PluginHttpProxy } -func (hp *HttpProxy) Handle(conn io.ReadWriteCloser) { - var wrapConn frpNet.Conn - if realConn, ok := conn.(frpNet.Conn); ok { - wrapConn = realConn - } else { - wrapConn = frpNet.WrapReadWriteCloserToConn(conn, realConn) - } +func (hp *HttpProxy) Handle(conn io.ReadWriteCloser, realConn frpNet.Conn) { + wrapConn := frpNet.WrapReadWriteCloserToConn(conn, realConn) sc, rd := frpNet.NewShareConn(wrapConn) request, err := http.ReadRequest(bufio.NewReader(rd)) diff --git a/models/plugin/plugin.go b/models/plugin/plugin.go index 0c60461..a0ae8e1 100644 --- a/models/plugin/plugin.go +++ b/models/plugin/plugin.go @@ -17,6 +17,8 @@ package plugin import ( "fmt" "io" + + frpNet "github.com/fatedier/frp/utils/net" ) // Creators is used for create plugins to handle connections. @@ -40,6 +42,6 @@ func Create(name string, params map[string]string) (p Plugin, err error) { type Plugin interface { Name() string - Handle(conn io.ReadWriteCloser) + Handle(conn io.ReadWriteCloser, realConn frpNet.Conn) Close() error } diff --git a/models/plugin/socks5.go b/models/plugin/socks5.go index b0f1bb2..d780488 100644 --- a/models/plugin/socks5.go +++ b/models/plugin/socks5.go @@ -43,16 +43,9 @@ func NewSocks5Plugin(params map[string]string) (p Plugin, err error) { return } -func (sp *Socks5Plugin) Handle(conn io.ReadWriteCloser) { +func (sp *Socks5Plugin) Handle(conn io.ReadWriteCloser, realConn frpNet.Conn) { defer conn.Close() - - var wrapConn frpNet.Conn - if realConn, ok := conn.(frpNet.Conn); ok { - wrapConn = realConn - } else { - wrapConn = frpNet.WrapReadWriteCloserToConn(conn, realConn) - } - + wrapConn := frpNet.WrapReadWriteCloserToConn(conn, realConn) sp.Server.ServeConn(wrapConn) } diff --git a/models/plugin/unix_domain_socket.go b/models/plugin/unix_domain_socket.go index 19e06f8..c3d50e9 100644 --- a/models/plugin/unix_domain_socket.go +++ b/models/plugin/unix_domain_socket.go @@ -20,6 +20,7 @@ import ( "net" frpIo "github.com/fatedier/frp/utils/io" + frpNet "github.com/fatedier/frp/utils/net" ) const PluginUnixDomainSocket = "unix_domain_socket" @@ -51,7 +52,7 @@ func NewUnixDomainSocketPlugin(params map[string]string) (p Plugin, err error) { return } -func (uds *UnixDomainSocketPlugin) Handle(conn io.ReadWriteCloser) { +func (uds *UnixDomainSocketPlugin) Handle(conn io.ReadWriteCloser, realConn frpNet.Conn) { localConn, err := net.DialUnix("unix", nil, uds.UnixAddr) if err != nil { return