Create only one additional goroutine when piping data.

cmd/shadowsocks-local/local.go

@@ -226,7 +226,12 @@ func handleConnection(conn net.Conn) {
 	if debug {
 		debug.Printf("socks connect from %s\n", conn.RemoteAddr().String())
 	}
-	defer conn.Close()
+	closed := false
+	defer func() {
+		if !closed {
+			conn.Close()
+		}
+	}()
 
 	var err error = nil
 	if err = handShake(conn); err != nil {
@@ -254,13 +259,15 @@ func handleConnection(conn net.Conn) {
 		}
 		return
 	}
-	defer remote.Close()
+	defer func() {
+		if !closed {
+			remote.Close()
+		}
+	}()
 
-	c := make(chan byte, 2)
-	go ss.Pipe(conn, remote, c)
-	go ss.Pipe(remote, conn, c)
-	<-c // close the other connection whenever one connection is closed
-	debug.Println("closing connection to", addr)
+	ss.Pipe(conn, remote)
+	closed = true
+	debug.Println("closed connection to", addr)
 }
 
 func run(port string) {

cmd/shadowsocks-server/server.go

@@ -103,12 +103,15 @@ func handleConnection(conn *ss.Conn) {
 	if debug {
 		debug.Printf("new client %s->%s\n", conn.RemoteAddr().String(), conn.LocalAddr())
 	}
+	closed := false
 	defer func() {
 		if debug {
-			debug.Printf("closing pipe %s<->%s\n", conn.RemoteAddr(), host)
+			debug.Printf("closed pipe %s<->%s\n", conn.RemoteAddr(), host)
 		}
 		atomic.AddInt32(&connCnt, -1)
-		conn.Close()
+		if !closed {
+			conn.Close()
+		}
 	}()
 
 	host, extra, err := getRequest(conn)
@@ -128,7 +131,11 @@ func handleConnection(conn *ss.Conn) {
 		}
 		return
 	}
-	defer remote.Close()
+	defer func() {
+		if !closed {
+			remote.Close()
+		}
+	}()
 	// write extra bytes read from
 	if extra != nil {
 		// debug.Println("getRequest read extra data, writing to remote, len", len(extra))
@@ -140,10 +147,8 @@ func handleConnection(conn *ss.Conn) {
 	if debug {
 		debug.Printf("piping %s<->%s", conn.RemoteAddr(), host)
 	}
-	c := make(chan byte, 2)
-	go ss.Pipe(conn, remote, c)
-	go ss.Pipe(remote, conn, c)
-	<-c // close the other connection whenever one connection is closed
+	ss.Pipe(conn, remote)
+	closed = true
 	return
 }
 

shadowsocks/pipe.go

@@ -12,12 +12,15 @@ func SetReadTimeout(c net.Conn) {
 	}
 }
 
-func Pipe(src, dst net.Conn, end chan byte) {
-	// Should not use io.Copy here.
-	// io.Copy will try to use the ReadFrom interface of TCPConn, but the src
-	// here is not a regular file, so sendfile is not applicable.
-	// io.Copy will fallback to the normal copy after discovering this,
-	// introducing unnecessary overhead.
+// Pipe copies data between c1 and c2. Closes c1 and c2 when done.
+func Pipe(c1, c2 net.Conn) {
+	go pipeClose(c1, c2)
+	pipeClose(c2, c1)
+}
+
+// pipeClose copies data from src to dst. Closes dst when done.
+func pipeClose(src, dst net.Conn) {
+	defer dst.Close()
 	buf := make([]byte, 4096)
 	for {
 		SetReadTimeout(src)
@@ -42,5 +45,4 @@ func Pipe(src, dst net.Conn, end chan byte) {
 			break
 		}
 	}
-	end <- 1
 }