From 66ca0497dfb58c4daaf663f53fc8bef7ded309be Mon Sep 17 00:00:00 2001
From: ymmt2005 <ymmt2005@gmail.com>
Date: Tue, 23 Feb 2016 20:23:36 +0900
Subject: [PATCH] Let proxy pass down EOF (FIN) to the other end.

This PR improves proxy() implementation so that it passes down EOF
(i.e. TCP FIN) from the read end to the write end by using
TCPConn.CloseWrite().  This is a must for TCP proxy as TCP sockets
can be half-closed.

Moreover, Server.handleConnect() should wait for two acks from
errCh to guarantee that the communication is completely over.
This is also fixed.

Lastly, noisy debug log is removed.
---
 request.go | 31 ++++++++++++++-----------------
 1 file changed, 14 insertions(+), 17 deletions(-)

diff --git a/request.go b/request.go
index 5929ba5..7e713f5 100644
--- a/request.go
+++ b/request.go
@@ -3,10 +3,8 @@ package socks5
 import (
 	"fmt"
 	"io"
-	"log"
 	"net"
 	"strings"
-	"time"
 )
 
 const (
@@ -183,14 +181,18 @@ func (s *Server) handleConnect(conn conn, req *Request) error {
 
 	// Start proxying
 	errCh := make(chan error, 2)
-	go proxy("target", target, req.bufConn, errCh, s.config.Logger)
-	go proxy("client", conn, target, errCh, s.config.Logger)
+	go proxy(target, req.bufConn, errCh)
+	go proxy(conn, target, errCh)
 
 	// Wait
-	select {
-	case e := <-errCh:
-		return e
+	for i := 0; i < 2; i++ {
+		e := <-errCh
+		if e != nil {
+			// return from this function closes target (and conn).
+			return e
+		}
 	}
+	return nil
 }
 
 // handleBind is used to handle a connect command
@@ -327,15 +329,10 @@ func sendReply(w io.Writer, resp uint8, addr *AddrSpec) error {
 
 // proxy is used to suffle data from src to destination, and sends errors
 // down a dedicated channel
-func proxy(name string, dst io.Writer, src io.Reader, errCh chan error, logger *log.Logger) {
-	// Copy
-	n, err := io.Copy(dst, src)
-
-	// Log, and sleep. This is jank but allows the otherside
-	// to finish a pending copy
-	logger.Printf("[DEBUG] socks: Copied %d bytes to %s", n, name)
-	time.Sleep(10 * time.Millisecond)
-
-	// Send any errors
+func proxy(dst io.Writer, src io.Reader, errCh chan error) {
+	_, err := io.Copy(dst, src)
+	if tcpConn, ok := dst.(*net.TCPConn); ok {
+		tcpConn.CloseWrite()
+	}
 	errCh <- err
 }