From 6f85d8e975e316d2e825db5c349c33eb8eb627d2 Mon Sep 17 00:00:00 2001
From: Andrey Smirnov <smirnov.andrey@gmail.com>
Date: Thu, 15 Oct 2020 00:51:38 +0300
Subject: [PATCH] Implement correct half-close sequence for the connections.

Fixes #21

Signed-off-by: Andrey Smirnov <smirnov.andrey@gmail.com>
---
 go.mod      |  5 ++++-
 go.sum      |  2 ++
 tcpproxy.go | 21 ++++++++++++++++-----
 3 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/go.mod b/go.mod
index 9c8ce9f..4d3d20c 100644
--- a/go.mod
+++ b/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
+)
diff --git a/go.sum b/go.sum
index de51fb1..8f848e3 100644
--- a/go.sum
+++ b/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=
diff --git a/tcpproxy.go b/tcpproxy.go
index df13942..3e0c59f 100644
--- a/tcpproxy.go
+++ b/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
 }