From e63edf083bf578a66fde719aa462c6bec5430f26 Mon Sep 17 00:00:00 2001
From: Riobard <me@riobard.com>
Date: Sat, 11 Mar 2017 13:14:45 +0800
Subject: [PATCH] Fix socket leaks

No need to half close any more.
---
 shadowaead/stream.go   | 22 ---------------------
 shadowstream/stream.go | 22 ---------------------
 tcp.go                 | 44 ++++++++++++++----------------------------
 3 files changed, 14 insertions(+), 74 deletions(-)

diff --git a/shadowaead/stream.go b/shadowaead/stream.go
index 6c931552..a8dce31a 100644
--- a/shadowaead/stream.go
+++ b/shadowaead/stream.go
@@ -181,14 +181,6 @@ func increment(b []byte) {
 	}
 }
 
-type closeWriter interface {
-	CloseWrite() error
-}
-
-type closeReader interface {
-	CloseRead() error
-}
-
 type streamConn struct {
 	net.Conn
 	Cipher
@@ -264,19 +256,5 @@ func (c *streamConn) ReadFrom(r io.Reader) (int64, error) {
 	return c.w.ReadFrom(r)
 }
 
-func (c *streamConn) CloseRead() error {
-	if c, ok := c.Conn.(closeReader); ok {
-		return c.CloseRead()
-	}
-	return nil
-}
-
-func (c *streamConn) CloseWrite() error {
-	if c, ok := c.Conn.(closeWriter); ok {
-		return c.CloseWrite()
-	}
-	return nil
-}
-
 // NewConn wraps a stream-oriented net.Conn with cipher.
 func NewConn(c net.Conn, ciph Cipher) net.Conn { return &streamConn{Conn: c, Cipher: ciph} }
diff --git a/shadowstream/stream.go b/shadowstream/stream.go
index 9e85fefc..eb4d9679 100644
--- a/shadowstream/stream.go
+++ b/shadowstream/stream.go
@@ -169,25 +169,3 @@ func (c *conn) ReadFrom(r io.Reader) (int64, error) {
 	}
 	return c.w.ReadFrom(r)
 }
-
-type closeWriter interface {
-	CloseWrite() error
-}
-
-type closeReader interface {
-	CloseRead() error
-}
-
-func (c *conn) CloseRead() error {
-	if c, ok := c.Conn.(closeReader); ok {
-		return c.CloseRead()
-	}
-	return nil
-}
-
-func (c *conn) CloseWrite() error {
-	if c, ok := c.Conn.(closeWriter); ok {
-		return c.CloseWrite()
-	}
-	return nil
-}
diff --git a/tcp.go b/tcp.go
index 50654da9..c90445ed 100644
--- a/tcp.go
+++ b/tcp.go
@@ -3,6 +3,7 @@ package main
 import (
 	"io"
 	"net"
+	"time"
 
 	"github.com/riobard/go-shadowsocks2/core"
 	"github.com/riobard/go-shadowsocks2/socks"
@@ -69,8 +70,10 @@ func tcpLocalHandle(c net.Conn, server string, target socks.Addr, ciph core.Stre
 
 	_, _, err = relay(sc, c)
 	if err != nil {
+		if err, ok := err.(net.Error); ok && err.Timeout() {
+			return // ignore i/o timeout
+		}
 		logf("relay error: %v", err)
-		return
 	}
 }
 
@@ -112,14 +115,16 @@ func tcpRemoteHandle(c net.Conn) {
 
 	_, _, err = relay(c, conn)
 	if err != nil {
+		if err, ok := err.(net.Error); ok && err.Timeout() {
+			return // ignore i/o timeout
+		}
 		logf("relay error: %v", err)
-		return
 	}
 }
 
 // relay copies between left and right bidirectionally. Returns number of
 // bytes copied from right to left, from left to right, and any error occurred.
-func relay(left, right io.ReadWriter) (int64, int64, error) {
+func relay(left, right net.Conn) (int64, int64, error) {
 	type res struct {
 		N   int64
 		Err error
@@ -127,11 +132,15 @@ func relay(left, right io.ReadWriter) (int64, int64, error) {
 	ch := make(chan res)
 
 	go func() {
-		n, err := copyHalfClose(right, left)
+		n, err := io.Copy(right, left)
+		right.SetDeadline(time.Now()) // wake up the other goroutine blocking on right
+		left.SetDeadline(time.Now())  // wake up the other goroutine blocking on left
 		ch <- res{n, err}
 	}()
 
-	n, err := copyHalfClose(left, right)
+	n, err := io.Copy(left, right)
+	right.SetDeadline(time.Now()) // wake up the other goroutine blocking on right
+	left.SetDeadline(time.Now())  // wake up the other goroutine blocking on left
 	rs := <-ch
 
 	if err == nil {
@@ -139,28 +148,3 @@ func relay(left, right io.ReadWriter) (int64, int64, error) {
 	}
 	return n, rs.N, err
 }
-
-type closeWriter interface {
-	CloseWrite() error
-}
-
-type closeReader interface {
-	CloseRead() error
-}
-
-// copyHalfClose copies to dst from src and optionally closes dst for writing and src for reading.
-func copyHalfClose(dst io.Writer, src io.Reader) (int64, error) {
-	defer func() {
-		// half-close to wake up other goroutines blocking on dst and src
-
-		if c, ok := dst.(closeWriter); ok {
-			c.CloseWrite()
-		}
-
-		if c, ok := src.(closeReader); ok {
-			c.CloseRead()
-		}
-	}()
-
-	return io.Copy(dst, src) // will use io.ReaderFrom or io.WriterTo shortcut if possible
-}