From dee5898e3698d846f0239225c9753e204632316a Mon Sep 17 00:00:00 2001
From: Cesaryuan <35998162+cesaryuan@users.noreply.github.com>
Date: Sat, 15 Mar 2025 13:27:29 +0800
Subject: [PATCH] fix: memory leak due to unclosed session (#1908)

---
 transport/gost-plugin/websocket.go | 23 ++++++++++++++++++++++-
 1 file changed, 22 insertions(+), 1 deletion(-)

diff --git a/transport/gost-plugin/websocket.go b/transport/gost-plugin/websocket.go
index 40bf08f5..e60af595 100644
--- a/transport/gost-plugin/websocket.go
+++ b/transport/gost-plugin/websocket.go
@@ -23,6 +23,23 @@ type Option struct {
 	Mux            bool
 }
 
+// muxConn is a wrapper around smux.Stream that also closes the session when closed
+type muxConn struct {
+	net.Conn
+	session *smux.Session
+}
+
+func (m *muxConn) Close() error {
+	streamErr := m.Conn.Close()
+	sessionErr := m.session.Close()
+
+	// Return stream error if there is one, otherwise return session error
+	if streamErr != nil {
+		return streamErr
+	}
+	return sessionErr
+}
+
 // NewGostWebsocket return a gost websocket
 func NewGostWebsocket(ctx context.Context, conn net.Conn, option *Option) (net.Conn, error) {
 	header := http.Header{}
@@ -72,10 +89,14 @@ func NewGostWebsocket(ctx context.Context, conn net.Conn, option *Option) (net.C
 
 		stream, err := session.OpenStream()
 		if err != nil {
+			session.Close()
 			return nil, err
 		}
 
-		conn = stream
+		return &muxConn{
+			Conn:    stream,
+			session: session,
+		}, nil
 	}
 	return conn, nil
 }