Update: wait both direction and close

- origin code channel no closed
- io.Copy will never return EOF error
This commit is contained in:
liuyang 2018-09-11 17:42:26 +08:00
parent 84dc5a612e
commit 9f72d6b1d3

View File

@ -3,6 +3,7 @@ package gost
import ( import (
"io" "io"
"net" "net"
"sync"
"time" "time"
"github.com/go-log/log" "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 { func transport(rw1, rw2 io.ReadWriter) error {
errc := make(chan error, 1) var err1 error
go func() { var err2 error
_, err := io.Copy(rw1, rw2) wg := sync.WaitGroup{}
errc <- err wg.Add(2)
}()
go func() { go func() {
_, err := io.Copy(rw2, rw1) _, err1 = io.Copy(rw1, rw2)
errc <- err wg.Done()
}()
go func() {
_, err2 = io.Copy(rw2, rw1)
wg.Done()
}() }()
err := <-errc // wait for both direction complete
if err != nil && err == io.EOF { wg.Wait()
err = nil if err1 != nil {
return err1
} }
return err return err2
} }