From 749cfaea65f8a74163f67a8e6c368aa8bad172d1 Mon Sep 17 00:00:00 2001
From: nange <cctonange@gmail.com>
Date: Tue, 14 Mar 2023 14:53:00 +0800
Subject: [PATCH] fix: memory leak on http tunnel

---
 httptunnel/server.go      |  9 ++++++++-
 httptunnel/server_conn.go | 16 +++++++++++++---
 2 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/httptunnel/server.go b/httptunnel/server.go
index c91a8496..d0b1c9f5 100644
--- a/httptunnel/server.go
+++ b/httptunnel/server.go
@@ -141,7 +141,7 @@ func (s *Server) push(w http.ResponseWriter, r *http.Request) {
 	s.Lock()
 	conn, ok := s.connMap[reqID]
 	if !ok {
-		conn = NewServerConn()
+		conn = NewServerConn(reqID, s.CloseConn)
 		s.connMap[reqID] = conn
 		s.connCh <- conn
 	}
@@ -162,6 +162,13 @@ func (s *Server) push(w http.ResponseWriter, r *http.Request) {
 	writeNoContent(w)
 }
 
+func (s *Server) CloseConn(reqID string) {
+	s.Lock()
+	defer s.Unlock()
+	s.connMap[reqID] = nil
+	delete(s.connMap, reqID)
+}
+
 func writeNotFoundError(w http.ResponseWriter) {
 	http.Error(w, "404 NOT FOUND", http.StatusNotFound)
 }
diff --git a/httptunnel/server_conn.go b/httptunnel/server_conn.go
index 497176fa..770375f5 100644
--- a/httptunnel/server_conn.go
+++ b/httptunnel/server_conn.go
@@ -8,15 +8,20 @@ import (
 var _ net.Conn = (*ServerConn)(nil)
 
 type ServerConn struct {
+	reqID     string
+	closeConn func(reqID string)
+
 	local  net.Conn
 	remote net.Conn
 }
 
-func NewServerConn() *ServerConn {
+func NewServerConn(reqID string, closeConn func(reqID string)) *ServerConn {
 	read, write := net.Pipe()
 	return &ServerConn{
-		local:  read,
-		remote: write,
+		reqID:     reqID,
+		closeConn: closeConn,
+		local:     read,
+		remote:    write,
 	}
 }
 
@@ -37,6 +42,11 @@ func (c *ServerConn) Write(b []byte) (n int, err error) {
 }
 
 func (c *ServerConn) Close() error {
+	if c.closeConn != nil {
+		c.closeConn(c.reqID)
+	}
+	c.closeConn = nil
+
 	return c.remote.Close()
 }