From 824f5bd731e48f843d9f0dcc3b9035563405ac69 Mon Sep 17 00:00:00 2001
From: Dreamacro <8615343+Dreamacro@users.noreply.github.com>
Date: Fri, 7 May 2021 11:08:46 +0800
Subject: [PATCH] Fix: reuse http connection broken on previous commit

---
 tunnel/connection.go | 32 +++++++++++++++++++++++++++++++-
 1 file changed, 31 insertions(+), 1 deletion(-)

diff --git a/tunnel/connection.go b/tunnel/connection.go
index fe1e72b2..24314fb6 100644
--- a/tunnel/connection.go
+++ b/tunnel/connection.go
@@ -21,6 +21,9 @@ func handleHTTP(ctx *context.HTTPContext, outbound net.Conn) {
 	req := ctx.Request()
 	conn := ctx.Conn()
 
+	// make outbound close after inbound error or close
+	conn = &connLinker{conn, outbound}
+
 	inboundReader := bufio.NewReader(conn)
 	outboundReader := bufio.NewReader(outbound)
 
@@ -31,7 +34,6 @@ func handleHTTP(ctx *context.HTTPContext, outbound net.Conn) {
 
 		req.RequestURI = ""
 		inbound.RemoveHopByHopHeaders(req.Header)
-		req.Header.Set("Connection", "close")
 		err := req.Write(outbound)
 		if err != nil {
 			break
@@ -160,3 +162,31 @@ func relay(leftConn, rightConn net.Conn) {
 	rightConn.SetReadDeadline(time.Now())
 	<-ch
 }
+
+// connLinker make the two net.Conn correlated, for temporary resolution of leaks.
+// There is no better way to do this for now.
+type connLinker struct {
+	net.Conn
+	linker net.Conn
+}
+
+func (conn *connLinker) Read(b []byte) (n int, err error) {
+	n, err = conn.Conn.Read(b)
+	if err != nil {
+		conn.linker.Close()
+	}
+	return n, err
+}
+
+func (conn *connLinker) Write(b []byte) (n int, err error) {
+	n, err = conn.Conn.Write(b)
+	if err != nil {
+		conn.linker.Close()
+	}
+	return n, err
+}
+
+func (conn *connLinker) Close() error {
+	conn.linker.Close()
+	return conn.Conn.Close()
+}