From 17cc0735d1bce68ebe73a26c24475c1a854dd334 Mon Sep 17 00:00:00 2001
From: fatedier <fatedier@gmail.com>
Date: Fri, 12 Jul 2019 16:53:21 +0800
Subject: [PATCH] add read timeout for TLS check operation

---
 server/service.go        | 11 ++++++++++-
 utils/net/tls.go         | 14 +++++++++++---
 utils/version/version.go |  2 +-
 3 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/server/service.go b/server/service.go
index 6cd8e50..d1207f4 100644
--- a/server/service.go
+++ b/server/service.go
@@ -259,7 +259,16 @@ func (svr *Service) HandleListener(l frpNet.Listener) {
 			log.Warn("Listener for incoming connections from client closed")
 			return
 		}
-		c = frpNet.CheckAndEnableTLSServerConn(c, svr.tlsConfig)
+
+		log.Trace("start check TLS connection...")
+		originConn := c
+		c, err = frpNet.CheckAndEnableTLSServerConnWithTimeout(c, svr.tlsConfig, connReadTimeout)
+		if err != nil {
+			log.Warn("CheckAndEnableTLSServerConnWithTimeout error: %v", err)
+			originConn.Close()
+			continue
+		}
+		log.Trace("success check TLS connection")
 
 		// Start a new goroutine for dealing connections.
 		go func(frpConn frpNet.Conn) {
diff --git a/utils/net/tls.go b/utils/net/tls.go
index ae1bfc7..4ac51d5 100644
--- a/utils/net/tls.go
+++ b/utils/net/tls.go
@@ -17,6 +17,7 @@ package net
 import (
 	"crypto/tls"
 	"net"
+	"time"
 
 	gnet "github.com/fatedier/golib/net"
 )
@@ -31,10 +32,17 @@ func WrapTLSClientConn(c net.Conn, tlsConfig *tls.Config) (out Conn) {
 	return
 }
 
-func CheckAndEnableTLSServerConn(c net.Conn, tlsConfig *tls.Config) (out Conn) {
-	sc, r := gnet.NewSharedConnSize(c, 1)
+func CheckAndEnableTLSServerConnWithTimeout(c net.Conn, tlsConfig *tls.Config, timeout time.Duration) (out Conn, err error) {
+	sc, r := gnet.NewSharedConnSize(c, 2)
 	buf := make([]byte, 1)
-	n, _ := r.Read(buf)
+	var n int
+	c.SetReadDeadline(time.Now().Add(timeout))
+	n, err = r.Read(buf)
+	c.SetReadDeadline(time.Time{})
+	if err != nil {
+		return
+	}
+
 	if n == 1 && int(buf[0]) == FRP_TLS_HEAD_BYTE {
 		out = WrapConn(tls.Server(c, tlsConfig))
 	} else {
diff --git a/utils/version/version.go b/utils/version/version.go
index 9bc4934..da23fbe 100644
--- a/utils/version/version.go
+++ b/utils/version/version.go
@@ -19,7 +19,7 @@ import (
 	"strings"
 )
 
-var version string = "0.27.0"
+var version string = "0.27.1"
 
 func Full() string {
 	return version