From f90028cf961b3419e87b62115ba517a67f7eb53e Mon Sep 17 00:00:00 2001
From: fatedier <fatedier@gmail.com>
Date: Fri, 10 Mar 2017 01:42:06 +0800
Subject: [PATCH] Use encryption in frp protocol.

---
 client/control.go             | 20 ++++++++++++++++----
 client/proxy.go               |  3 +++
 server/control.go             | 22 +++++++++++++++++-----
 server/proxy.go               |  1 +
 utils/version/version_test.go |  2 +-
 5 files changed, 38 insertions(+), 10 deletions(-)

diff --git a/client/control.go b/client/control.go
index 9cc139a6..1f5e588a 100644
--- a/client/control.go
+++ b/client/control.go
@@ -23,6 +23,7 @@ import (
 
 	"github.com/fatedier/frp/models/config"
 	"github.com/fatedier/frp/models/msg"
+	"github.com/fatedier/frp/utils/crypto"
 	"github.com/fatedier/frp/utils/log"
 	"github.com/fatedier/frp/utils/net"
 	"github.com/fatedier/frp/utils/util"
@@ -135,7 +136,7 @@ func (ctl *Control) NewWorkConn() {
 
 	var startMsg msg.StartWorkConn
 	if err = msg.ReadMsgInto(workConn, &startMsg); err != nil {
-		ctl.Error("work connection closed and no response from server, %v", err)
+		ctl.Error("work connection closed, %v", err)
 		workConn.Close()
 		return
 	}
@@ -205,12 +206,17 @@ func (ctl *Control) reader() {
 			ctl.Error("panic error: %v", err)
 		}
 	}()
+	defer close(ctl.closedCh)
 
+	encReader, err := crypto.NewReader(ctl.conn, []byte(config.ClientCommonCfg.PrivilegeToken))
+	if err != nil {
+		ctl.conn.Error("crypto new reader error: %v", err)
+		return
+	}
 	for {
-		if m, err := msg.ReadMsg(ctl.conn); err != nil {
+		if m, err := msg.ReadMsg(encReader); err != nil {
 			if err == io.EOF {
 				ctl.Debug("read from control connection EOF")
-				close(ctl.closedCh)
 				return
 			} else {
 				ctl.Warn("read error: %v", err)
@@ -223,12 +229,18 @@ func (ctl *Control) reader() {
 }
 
 func (ctl *Control) writer() {
+	encWriter, err := crypto.NewWriter(ctl.conn, []byte(config.ClientCommonCfg.PrivilegeToken))
+	if err != nil {
+		ctl.conn.Error("crypto new writer error: %v", err)
+		ctl.conn.Close()
+		return
+	}
 	for {
 		if m, ok := <-ctl.sendCh; !ok {
 			ctl.Info("control writer is closing")
 			return
 		} else {
-			if err := msg.WriteMsg(ctl.conn, m); err != nil {
+			if err := msg.WriteMsg(encWriter, m); err != nil {
 				ctl.Warn("write message to control connection error: %v", err)
 				return
 			}
diff --git a/client/proxy.go b/client/proxy.go
index 5467bc39..a5a400ba 100644
--- a/client/proxy.go
+++ b/client/proxy.go
@@ -23,8 +23,11 @@ import (
 	"github.com/fatedier/frp/utils/net"
 )
 
+// Proxy defines how to work for different proxy type.
 type Proxy interface {
 	Run()
+
+	// InWorkConn accept work connections registered to server.
 	InWorkConn(conn net.Conn)
 	Close()
 }
diff --git a/server/control.go b/server/control.go
index e3027ba5..0c2d45ca 100644
--- a/server/control.go
+++ b/server/control.go
@@ -9,6 +9,7 @@ import (
 	"github.com/fatedier/frp/models/config"
 	"github.com/fatedier/frp/models/consts"
 	"github.com/fatedier/frp/models/msg"
+	"github.com/fatedier/frp/utils/crypto"
 	"github.com/fatedier/frp/utils/errors"
 	"github.com/fatedier/frp/utils/net"
 	"github.com/fatedier/frp/utils/shutdown"
@@ -81,14 +82,14 @@ func NewControl(svr *Service, ctlConn net.Conn, loginMsg *msg.Login) *Control {
 
 // Start send a login success message to client and start working.
 func (ctl *Control) Start() {
-	go ctl.writer()
-
-	ctl.sendCh <- &msg.LoginResp{
+	loginRespMsg := &msg.LoginResp{
 		Version: version.Full(),
 		RunId:   ctl.runId,
 		Error:   "",
 	}
+	msg.WriteMsg(ctl.conn, loginRespMsg)
 
+	go ctl.writer()
 	for i := 0; i < ctl.poolCount; i++ {
 		ctl.sendCh <- &msg.ReqWorkConn{}
 	}
@@ -182,12 +183,18 @@ func (ctl *Control) writer() {
 	defer ctl.allShutdown.Start()
 	defer ctl.writerShutdown.Done()
 
+	encWriter, err := crypto.NewWriter(ctl.conn, []byte(config.ServerCommonCfg.PrivilegeToken))
+	if err != nil {
+		ctl.conn.Error("crypto new writer error: %v", err)
+		ctl.allShutdown.Start()
+		return
+	}
 	for {
 		if m, ok := <-ctl.sendCh; !ok {
 			ctl.conn.Info("control writer is closing")
 			return
 		} else {
-			if err := msg.WriteMsg(ctl.conn, m); err != nil {
+			if err := msg.WriteMsg(encWriter, m); err != nil {
 				ctl.conn.Warn("write message to control connection error: %v", err)
 				return
 			}
@@ -205,8 +212,13 @@ func (ctl *Control) reader() {
 	defer ctl.allShutdown.Start()
 	defer ctl.readerShutdown.Done()
 
+	encReader, err := crypto.NewReader(ctl.conn, []byte(config.ServerCommonCfg.PrivilegeToken))
+	if err != nil {
+		ctl.conn.Error("crypto new reader error: %v", err)
+		return
+	}
 	for {
-		if m, err := msg.ReadMsg(ctl.conn); err != nil {
+		if m, err := msg.ReadMsg(encReader); err != nil {
 			if err == io.EOF {
 				ctl.conn.Debug("control connection closed")
 				return
diff --git a/server/proxy.go b/server/proxy.go
index 108a8a9c..e751789c 100644
--- a/server/proxy.go
+++ b/server/proxy.go
@@ -165,6 +165,7 @@ func (pxy *UdpProxy) Close() {
 }
 
 // HandleUserTcpConnection is used for incoming tcp user connections.
+// It can be used for tcp, http, https type.
 func HandleUserTcpConnection(pxy Proxy, userConn net.Conn) {
 	defer userConn.Close()
 	ctl := pxy.GetControl()
diff --git a/utils/version/version_test.go b/utils/version/version_test.go
index a77bf421..ce70dd94 100644
--- a/utils/version/version_test.go
+++ b/utils/version/version_test.go
@@ -61,5 +61,5 @@ func TestCompact(t *testing.T) {
 	assert.True(ok)
 
 	ok, _ = Compat("0.10.0")
-	assert.False(ok)
+	assert.True(ok)
 }