Implement correct half-close sequence for the connections.

Fixes inetaf#21

Signed-off-by: Andrey Smirnov <[email protected]>

Author: smira

go.mod

@@ -2,4 +2,7 @@ module inet.af/tcpproxy
 
 go 1.16
 
-require github.com/armon/go-proxyproto v0.0.0-20210323213023-7e956b284f0a
+require (
+	github.com/armon/go-proxyproto v0.0.0-20210323213023-7e956b284f0a
+	golang.org/x/sys v0.8.0
+)

go.sum

@@ -1,2 +1,4 @@
 github.com/armon/go-proxyproto v0.0.0-20210323213023-7e956b284f0a h1:AP/vsCIvJZ129pdm9Ek7bH7yutN3hByqsMoNrWAxRQc=
 github.com/armon/go-proxyproto v0.0.0-20210323213023-7e956b284f0a/go.mod h1:QmP9hvJ91BbJmGVGSbutW19IC0Q9phDCLGaomwTJbgU=
+golang.org/x/sys v0.8.0 h1:EBmGv8NaZBZTWvrbjNoL6HVt+IVy3QDQpJs7VRIw3tU=
+golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=

tcpproxy.go

@@ -349,8 +349,6 @@ func UnderlyingConn(c net.Conn) net.Conn {
 	return c
 }
 
-func goCloseConn(c net.Conn) { go c.Close() }
-
 // HandleConn implements the Target interface.
 func (dp *DialProxy) HandleConn(src net.Conn) {
 	ctx := context.Background()
@@ -366,13 +364,13 @@ func (dp *DialProxy) HandleConn(src net.Conn) {
 		dp.onDialError()(src, err)
 		return
 	}
-	defer goCloseConn(dst)
+	defer dst.Close()
 
 	if err = dp.sendProxyHeader(dst, src); err != nil {
 		dp.onDialError()(src, err)
 		return
 	}
-	defer goCloseConn(src)
+	defer src.Close()
 
 	if ka := dp.keepAlivePeriod(); ka > 0 {
 		if c, ok := UnderlyingConn(src).(*net.TCPConn); ok {
@@ -393,7 +391,12 @@ func (dp *DialProxy) HandleConn(src net.Conn) {
 	errc := make(chan error, 1)
 	go proxyCopy(errc, src, dst)
 	go proxyCopy(errc, dst, src)
-	<-errc
+
+	for i := 0; i < 2; i++ {
+		if err = <-errc; err != nil {
+			return
+		}
+	}
 }
 
 func (dp *DialProxy) sendProxyHeader(w io.Writer, src net.Conn) error {
@@ -444,6 +447,14 @@ func proxyCopy(errc chan<- error, dst, src net.Conn) {
 	dst = UnderlyingConn(dst)
 
 	_, err := io.Copy(dst, src)
+
+	if tcpConn, ok := dst.(*net.TCPConn); ok {
+		tcpConn.CloseWrite()
+	}
+	if tcpConn, ok := src.(*net.TCPConn); ok {
+		tcpConn.CloseRead()
+	}
+
 	errc <- err
 }