From 338d5bae37fdd7ceba4aefcaeda20e681eb2da0a Mon Sep 17 00:00:00 2001
From: fatedier <fatedier@gmail.com>
Date: Thu, 18 Jan 2018 00:45:11 +0800
Subject: [PATCH] fix panic when using socks5 plugin with encryption and
 compression, fix #446

---
 client/proxy.go                     |  2 +-
 models/plugin/http_proxy.go         |  9 ++-------
 models/plugin/plugin.go             |  4 +++-
 models/plugin/socks5.go             | 11 ++---------
 models/plugin/unix_domain_socket.go |  3 ++-
 5 files changed, 10 insertions(+), 19 deletions(-)

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