From a78dbe7133bea24c5daabedfb74e75231cf3fc6f Mon Sep 17 00:00:00 2001
From: V2Ray <admin@v2ray.com>
Date: Thu, 24 Sep 2015 14:51:19 +0200
Subject: [PATCH] Close connection more aggressively.

---
 common/io/encryption.go |  3 ---
 common/net/timed_io.go  | 30 ++++++++++++++++++++++++++++++
 proxy/socks/socks.go    |  7 ++++++-
 proxy/vmess/vmessin.go  | 15 ++++++++-------
 proxy/vmess/vmessout.go |  3 +--
 5 files changed, 45 insertions(+), 13 deletions(-)
 create mode 100644 common/net/timed_io.go

diff --git a/common/io/encryption.go b/common/io/encryption.go
index 409f6b20..eb054ff3 100644
--- a/common/io/encryption.go
+++ b/common/io/encryption.go
@@ -27,9 +27,6 @@ func (reader CryptionReader) Read(blocks []byte) (int, error) {
 	if nBytes > 0 {
 		reader.stream.XORKeyStream(blocks[:nBytes], blocks[:nBytes])
 	}
-	if err != nil && err != io.EOF {
-		log.Error("Error reading blocks: %v", err)
-	}
 	return nBytes, err
 }
 
diff --git a/common/net/timed_io.go b/common/net/timed_io.go
new file mode 100644
index 00000000..cc9fb519
--- /dev/null
+++ b/common/net/timed_io.go
@@ -0,0 +1,30 @@
+package net
+
+import (
+	"net"
+	"time"
+)
+
+var (
+	emptyTime time.Time
+)
+
+type TimeOutReader struct {
+	timeout    int
+	connection net.Conn
+}
+
+func NewTimeOutReader(timeout int, connection net.Conn) *TimeOutReader {
+	return &TimeOutReader{
+		timeout:    timeout,
+		connection: connection,
+	}
+}
+
+func (reader *TimeOutReader) Read(p []byte) (n int, err error) {
+	deadline := time.Duration(reader.timeout) * time.Second
+	reader.connection.SetReadDeadline(time.Now().Add(deadline))
+	n, err = reader.connection.Read(p)
+	reader.connection.SetReadDeadline(emptyTime)
+	return
+}
diff --git a/proxy/socks/socks.go b/proxy/socks/socks.go
index 9286a627..98e88ca9 100644
--- a/proxy/socks/socks.go
+++ b/proxy/socks/socks.go
@@ -7,6 +7,7 @@ import (
 	"net"
 	"strconv"
 	"sync"
+	"time"
 
 	"github.com/v2ray/v2ray-core"
 	"github.com/v2ray/v2ray-core/common/log"
@@ -56,6 +57,10 @@ func (server *SocksServer) AcceptConnections(listener net.Listener) {
 		if err != nil {
 			log.Error("Error on accepting socks connection: %v", err)
 		}
+		if tcpConn, ok := connection.(*net.TCPConn); ok {
+			tcpConn.SetKeepAlive(true)
+			tcpConn.SetKeepAlivePeriod(4 * time.Second)
+		}
 		go server.HandleConnection(connection)
 	}
 }
@@ -63,7 +68,7 @@ func (server *SocksServer) AcceptConnections(listener net.Listener) {
 func (server *SocksServer) HandleConnection(connection net.Conn) error {
 	defer connection.Close()
 
-	reader := connection.(io.Reader)
+	reader := v2net.NewTimeOutReader(4, connection)
 
 	auth, auth4, err := protocol.ReadAuthentication(reader)
 	if err != nil && err != protocol.ErrorSocksVersion4 {
diff --git a/proxy/vmess/vmessin.go b/proxy/vmess/vmessin.go
index 250aa220..30fc9411 100644
--- a/proxy/vmess/vmessin.go
+++ b/proxy/vmess/vmessin.go
@@ -54,6 +54,10 @@ func (handler *VMessInboundHandler) AcceptConnections(listener net.Listener) err
 		if err != nil {
 			return log.Error("Failed to accpet connection: %s", err.Error())
 		}
+		if tcpConn, ok := connection.(*net.TCPConn); ok {
+			tcpConn.SetKeepAlive(true)
+			tcpConn.SetKeepAlivePeriod(4 * time.Second)
+		}
 		go handler.HandleConnection(connection)
 	}
 	return nil
@@ -62,18 +66,15 @@ func (handler *VMessInboundHandler) AcceptConnections(listener net.Listener) err
 func (handler *VMessInboundHandler) HandleConnection(connection net.Conn) error {
 	defer connection.Close()
 
-	reader := protocol.NewVMessRequestReader(handler.clients)
+	connReader := v2net.NewTimeOutReader(4, connection)
+	requestReader := protocol.NewVMessRequestReader(handler.clients)
 
-	// Timeout 4 seconds to prevent DoS attack
-	connection.SetReadDeadline(time.Now().Add(requestReadTimeOut))
-	request, err := reader.Read(connection)
+	request, err := requestReader.Read(connReader)
 	if err != nil {
 		log.Warning("VMessIn: Invalid request from (%s): %v", connection.RemoteAddr().String(), err)
 		return err
 	}
 	log.Debug("VMessIn: Received request for %s", request.Address.String())
-	// Clear read timeout
-	connection.SetReadDeadline(zeroTime)
 
 	ray := handler.vPoint.DispatchToOutbound(v2net.NewTCPPacket(request.Destination()))
 	input := ray.InboundInput()
@@ -82,7 +83,7 @@ func (handler *VMessInboundHandler) HandleConnection(connection net.Conn) error
 	readFinish.Lock()
 	writeFinish.Lock()
 
-	go handleInput(request, connection, input, &readFinish)
+	go handleInput(request, connReader, input, &readFinish)
 
 	responseKey := md5.Sum(request.RequestKey[:])
 	responseIV := md5.Sum(request.RequestIV[:])
diff --git a/proxy/vmess/vmessout.go b/proxy/vmess/vmessout.go
index 1aa97efd..e2f4bf46 100644
--- a/proxy/vmess/vmessout.go
+++ b/proxy/vmess/vmessout.go
@@ -170,8 +170,7 @@ func handleResponse(conn *net.TCPConn, request *protocol.VMessRequest, output ch
 	response := protocol.VMessResponse{}
 	nBytes, err := decryptResponseReader.Read(response[:])
 	if err != nil {
-		log.Error("VMessOut: Failed to read VMess response (%d bytes): %v", nBytes, err)
-		log.Error(InfoTimeNotSync)
+		//log.Error("VMessOut: Failed to read VMess response (%d bytes): %v", nBytes, err)
 		return
 	}
 	if !bytes.Equal(response[:], request.ResponseHeader[:]) {