Porting ss-redir to go

cmd/shadowsocks-local/local.go

@@ -18,7 +18,7 @@ import (
 	"syscall"
 	"time"
 
-	ss "github.com/shadowsocks/shadowsocks-go/shadowsocks"
+	ss "github.com/bonafideyan/shadowsocks-go/shadowsocks"
 )
 
 const SO_ORIGINAL_DST = 80

shadowsocks/leakybuf.go

@@ -20,6 +20,11 @@ func NewLeakyBuf(n, bufSize int) *LeakyBuf {
 	}
 }
 
+// GetLeakyBuf is invoked outside of the package.
+func GetLeakyBuf() []byte {
+	return leakyBuf.Get()
+}
+
 // Get returns a buffer from the leaky buffer or create a new buffer.
 func (lb *LeakyBuf) Get() (b []byte) {
 	select {

shadowsocks/pipe.go

@@ -12,16 +12,15 @@ func SetReadTimeout(c net.Conn) {
 }
 
 // PipeThenClose copies data from src to dst, closes dst when done.
-func PipeThenClose(src, dst net.Conn, addTraffic func(int)) {
+func PipeThenClose(src, dst net.Conn, addTraffic func(int), buf []byte, n int) {
 	defer dst.Close()
-	buf := leakyBuf.Get()
+	if buf == nil {
+		buf = leakyBuf.Get()
+	}
 	defer leakyBuf.Put(buf)
+
+	var err error
 	for {
-		SetReadTimeout(src)
-		n, err := src.Read(buf)
-		if addTraffic != nil {
-			addTraffic(n)
-		}
 		// read may return EOF with n > 0
 		// should always process n > 0 bytes before handling error
 		if n > 0 {
@@ -31,6 +30,13 @@ func PipeThenClose(src, dst net.Conn, addTraffic func(int)) {
 				break
 			}
 		}
+
+		SetReadTimeout(src)
+		n, err = src.Read(buf)
+		if addTraffic != nil {
+			addTraffic(n)
+		}
+
 		if err != nil {
 			// Always "use of closed network connection", but no easy way to
 			// identify this specific error. So just leave the error along for now.