From 9f72d6b1d33312d00e7ee0609dcb19f0b831dbd3 Mon Sep 17 00:00:00 2001
From: liuyang <x@icexz.com>
Date: Tue, 11 Sep 2018 17:42:26 +0800
Subject: [PATCH] Update: wait both direction and close

- origin code channel no closed
- io.Copy will never return EOF error
---
 server.go | 27 ++++++++++++++++-----------
 1 file changed, 16 insertions(+), 11 deletions(-)

diff --git a/server.go b/server.go
index 174fe60..f4dfd88 100644
--- a/server.go
+++ b/server.go
@@ -3,6 +3,7 @@ package gost
 import (
 	"io"
 	"net"
+	"sync"
 	"time"
 
 	"github.com/go-log/log"
@@ -99,20 +100,24 @@ func (ln tcpKeepAliveListener) Accept() (c net.Conn, err error) {
 }
 
 func transport(rw1, rw2 io.ReadWriter) error {
-	errc := make(chan error, 1)
-	go func() {
-		_, err := io.Copy(rw1, rw2)
-		errc <- err
-	}()
+	var err1 error
+	var err2 error
+	wg := sync.WaitGroup{}
+	wg.Add(2)
 
 	go func() {
-		_, err := io.Copy(rw2, rw1)
-		errc <- err
+		_, err1 = io.Copy(rw1, rw2)
+		wg.Done()
+	}()
+	go func() {
+		_, err2 = io.Copy(rw2, rw1)
+		wg.Done()
 	}()
 
-	err := <-errc
-	if err != nil && err == io.EOF {
-		err = nil
+	// wait for both direction complete
+	wg.Wait()
+	if err1 != nil {
+		return err1
 	}
-	return err
+	return err2
 }