From 7bc08271e89b45eb656941fd7ed32498756432de Mon Sep 17 00:00:00 2001
From: Jason <xjasonlyu@gmail.com>
Date: Sun, 11 Aug 2019 20:35:49 +0800
Subject: [PATCH] [context] ctxCopy

---
 proxy/socks/tcp.go | 46 ++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 44 insertions(+), 2 deletions(-)

diff --git a/proxy/socks/tcp.go b/proxy/socks/tcp.go
index 7d712fe..2cef1bf 100644
--- a/proxy/socks/tcp.go
+++ b/proxy/socks/tcp.go
@@ -1,6 +1,7 @@
 package socks
 
 import (
+	"context"
 	"io"
 	"net"
 	"strconv"
@@ -33,6 +34,47 @@ func NewTCPHandler(proxyHost string, proxyPort uint16, fakeDns dns.FakeDns, sess
 	}
 }
 
+func ctxCopy(ctx context.Context, dst, src net.Conn) (written int64, err error) {
+	buf := core.NewBytes(core.BufSize)
+	defer core.FreeBytes(buf)
+
+	for {
+		select {
+		case <-ctx.Done():
+			return written, err
+		default:
+		}
+		src.SetReadDeadline(time.Now().Add(30*time.Second))
+		nr, er := src.Read(buf)
+		if nr > 0 {
+			dst.SetWriteDeadline(time.Now().Add(30*time.Second))
+			nw, ew := dst.Write(buf[0:nr])
+			if nw > 0 {
+				written += int64(nw)
+			}
+			if ew != nil {
+				if ew, ok := ew.(net.Error); !ok || !ew.Timeout() {
+					err = ew
+					break
+				}
+			}
+			if nr != nw {
+				err = io.ErrShortWrite
+				break
+			}
+		}
+		if er != nil {
+			if er, ok := er.(net.Error); !ok || !er.Timeout() {
+				if er != io.EOF {
+					err = er
+				}
+				break
+			}
+		}
+	}
+	return written, err
+}
+
 func (h *tcpHandler) relay(localConn, remoteConn net.Conn) {
 	var once sync.Once
 	closeOnce := func() {
@@ -65,8 +107,8 @@ func (h *tcpHandler) relay(localConn, remoteConn net.Conn) {
 	}()
 
 	select {
-	case <-up: // Wait for UpLink done.
-	case <-down:
+	case <-up: // Wait for Up Link done
+	case <-down: // Wait for Down Link done
 	}
 
 	if h.sessionStater != nil {