From a84dd05351e33b1be90ccf20ae22a4b0abacece1 Mon Sep 17 00:00:00 2001
From: fatedier <fatedier@gmail.com>
Date: Wed, 10 May 2017 00:46:42 +0800
Subject: [PATCH] add connection read timeout

---
 client/control.go | 14 +++++++++++++-
 server/service.go |  6 ++++++
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/client/control.go b/client/control.go
index d618680b..928bb64a 100644
--- a/client/control.go
+++ b/client/control.go
@@ -30,6 +30,10 @@ import (
 	"github.com/fatedier/frp/utils/version"
 )
 
+const (
+	connReadTimeout time.Duration = 10 * time.Second
+)
+
 type Control struct {
 	// frpc service
 	svr *Service
@@ -144,7 +148,7 @@ func (ctl *Control) NewWorkConn() {
 
 	// dispatch this work connection to related proxy
 	if pxy, ok := ctl.proxies[startMsg.ProxyName]; ok {
-		workConn.Info("start a new work connection")
+		workConn.Info("start a new work connection, localAddr: %s remoteAddr: %s", workConn.LocalAddr().String(), workConn.RemoteAddr().String())
 		go pxy.InWorkConn(workConn)
 	} else {
 		workConn.Close()
@@ -168,6 +172,12 @@ func (ctl *Control) login() (err error) {
 		return err
 	}
 
+	defer func() {
+		if err != nil {
+			conn.Close()
+		}
+	}()
+
 	now := time.Now().Unix()
 	ctl.loginMsg.PrivilegeKey = util.GetAuthKey(config.ClientCommonCfg.PrivilegeToken, now)
 	ctl.loginMsg.Timestamp = now
@@ -178,9 +188,11 @@ func (ctl *Control) login() (err error) {
 	}
 
 	var loginRespMsg msg.LoginResp
+	conn.SetReadDeadline(time.Now().Add(connReadTimeout))
 	if err = msg.ReadMsgInto(conn, &loginRespMsg); err != nil {
 		return err
 	}
+	conn.SetReadDeadline(time.Time{})
 
 	if loginRespMsg.Error != "" {
 		err = fmt.Errorf("%s", loginRespMsg.Error)
diff --git a/server/service.go b/server/service.go
index b11c44c9..39adf20c 100644
--- a/server/service.go
+++ b/server/service.go
@@ -28,6 +28,10 @@ import (
 	"github.com/fatedier/frp/utils/vhost"
 )
 
+const (
+	connReadTimeout time.Duration = 10 * time.Second
+)
+
 var ServerService *Service
 
 // Server service.
@@ -121,11 +125,13 @@ func (svr *Service) Run() {
 		// Start a new goroutine for dealing connections.
 		go func(frpConn net.Conn) {
 			var rawMsg msg.Message
+			frpConn.SetReadDeadline(time.Now().Add(connReadTimeout))
 			if rawMsg, err = msg.ReadMsg(frpConn); err != nil {
 				log.Warn("Failed to read message: %v", err)
 				frpConn.Close()
 				return
 			}
+			frpConn.SetReadDeadline(time.Time{})
 
 			switch m := rawMsg.(type) {
 			case *msg.Login: