From 7aed8e85ef555067e861c5e542b665eca92193a4 Mon Sep 17 00:00:00 2001
From: Chen Yufei <cyfdecyf@gmail.com>
Date: Sun, 26 Apr 2015 22:57:15 +0800
Subject: [PATCH] Add -t for local, always set timeout for data transfer.

---
 cmd/shadowsocks-local/local.go   |  7 +++++--
 cmd/shadowsocks-server/server.go |  9 +++------
 shadowsocks/config.go            |  3 +++
 shadowsocks/pipe.go              | 11 ++---------
 4 files changed, 13 insertions(+), 17 deletions(-)

diff --git a/cmd/shadowsocks-local/local.go b/cmd/shadowsocks-local/local.go
index 78d0995..03bd8ef 100644
--- a/cmd/shadowsocks-local/local.go
+++ b/cmd/shadowsocks-local/local.go
@@ -49,6 +49,7 @@ func handShake(conn net.Conn) (err error) {
 	buf := make([]byte, 258)
 
 	var n int
+	ss.SetReadTimeout(conn)
 	// make sure we get the nmethod field
 	if n, err = io.ReadAtLeast(conn, buf, idNmethod+1); err != nil {
 		return
@@ -92,6 +93,7 @@ func getRequest(conn net.Conn) (rawaddr []byte, host string, err error) {
 	// refer to getRequest in server.go for why set buffer size to 263
 	buf := make([]byte, 263)
 	var n int
+	ss.SetReadTimeout(conn)
 	// read till we get possible domain length field
 	if n, err = io.ReadAtLeast(conn, buf, idDmLen+1); err != nil {
 		return
@@ -314,8 +316,8 @@ func handleConnection(conn net.Conn) {
 		}
 	}()
 
-	go ss.PipeThenClose(conn, remote, ss.NO_TIMEOUT)
-	ss.PipeThenClose(remote, conn, ss.NO_TIMEOUT)
+	go ss.PipeThenClose(conn, remote)
+	ss.PipeThenClose(remote, conn)
 	closed = true
 	debug.Println("closed connection to", addr)
 }
@@ -354,6 +356,7 @@ func main() {
 	flag.StringVar(&cmdLocal, "b", "", "local address, listen only to this address if specified")
 	flag.StringVar(&cmdConfig.Password, "k", "", "password")
 	flag.IntVar(&cmdConfig.ServerPort, "p", 0, "server port")
+	flag.IntVar(&cmdConfig.Timeout, "t", 300, "timeout in seconds")
 	flag.IntVar(&cmdConfig.LocalPort, "l", 0, "local socks5 proxy port")
 	flag.StringVar(&cmdConfig.Method, "m", "", "encryption method, default: aes-256-cfb")
 	flag.BoolVar((*bool)(&debug), "d", false, "print debug message")
diff --git a/cmd/shadowsocks-server/server.go b/cmd/shadowsocks-server/server.go
index 7214d95..982fcaf 100644
--- a/cmd/shadowsocks-server/server.go
+++ b/cmd/shadowsocks-server/server.go
@@ -19,8 +19,6 @@ import (
 
 var debug ss.DebugLog
 
-const dnsGoroutineNum = 64
-
 func getRequest(conn *ss.Conn) (host string, extra []byte, err error) {
 	const (
 		idType  = 0 // address type index
@@ -62,7 +60,6 @@ func getRequest(conn *ss.Conn) (host string, extra []byte, err error) {
 	}
 
 	if n < reqLen { // rare case
-		ss.SetReadTimeout(conn)
 		if _, err = io.ReadFull(conn, buf[n:reqLen]); err != nil {
 			return
 		}
@@ -154,8 +151,8 @@ func handleConnection(conn *ss.Conn) {
 	if debug {
 		debug.Printf("piping %s<->%s", conn.RemoteAddr(), host)
 	}
-	go ss.PipeThenClose(conn, remote, ss.SET_TIMEOUT)
-	ss.PipeThenClose(remote, conn, ss.NO_TIMEOUT)
+	go ss.PipeThenClose(conn, remote)
+	ss.PipeThenClose(remote, conn)
 	closed = true
 	return
 }
@@ -321,7 +318,7 @@ func main() {
 	flag.StringVar(&configFile, "c", "config.json", "specify config file")
 	flag.StringVar(&cmdConfig.Password, "k", "", "password")
 	flag.IntVar(&cmdConfig.ServerPort, "p", 0, "server port")
-	flag.IntVar(&cmdConfig.Timeout, "t", 60, "connection timeout (in seconds)")
+	flag.IntVar(&cmdConfig.Timeout, "t", 300, "timeout in seconds")
 	flag.StringVar(&cmdConfig.Method, "m", "", "encryption method, default: aes-256-cfb")
 	flag.IntVar(&core, "core", 0, "maximum number of CPU cores to use, default is determinied by Go runtime")
 	flag.BoolVar((*bool)(&debug), "d", false, "print debug message")
diff --git a/shadowsocks/config.go b/shadowsocks/config.go
index f96ff50..9d3696b 100644
--- a/shadowsocks/config.go
+++ b/shadowsocks/config.go
@@ -127,4 +127,7 @@ func UpdateConfig(old, new *Config) {
 	if old.Method == "table" {
 		old.Method = ""
 	}
+
+	old.Timeout = new.Timeout
+	readTimeout = time.Duration(old.Timeout) * time.Second
 }
diff --git a/shadowsocks/pipe.go b/shadowsocks/pipe.go
index 877853b..515c99f 100644
--- a/shadowsocks/pipe.go
+++ b/shadowsocks/pipe.go
@@ -6,11 +6,6 @@ import (
 	"time"
 )
 
-const (
-	NO_TIMEOUT = iota
-	SET_TIMEOUT
-)
-
 func SetReadTimeout(c net.Conn) {
 	if readTimeout != 0 {
 		c.SetReadDeadline(time.Now().Add(readTimeout))
@@ -18,14 +13,12 @@ func SetReadTimeout(c net.Conn) {
 }
 
 // PipeThenClose copies data from src to dst, closes dst when done.
-func PipeThenClose(src, dst net.Conn, timeoutOpt int) {
+func PipeThenClose(src, dst net.Conn) {
 	defer dst.Close()
 	buf := leakyBuf.Get()
 	defer leakyBuf.Put(buf)
 	for {
-		if timeoutOpt == SET_TIMEOUT {
-			SetReadTimeout(src)
-		}
+		SetReadTimeout(src)
 		n, err := src.Read(buf)
 		// read may return EOF with n > 0
 		// should always process n > 0 bytes before handling error