From ed7921ce166e742cbe184d515fdecf8ebc6b2685 Mon Sep 17 00:00:00 2001
From: Nikolay Edigaryev <edigaryev@gmail.com>
Date: Fri, 11 Jul 2025 13:23:50 +0200
Subject: [PATCH] Fix websocket.(*Conn).timeoutLoop goroutine leak (#329)

---
 internal/command/worker/run.go             | 13 +++++++++++++
 internal/controller/api_rpc_portforward.go |  7 +++++++
 internal/controller/api_rpc_watch.go       |  7 +++++++
 internal/controller/api_vms_portforward.go |  7 +++++++
 internal/worker/rpcv2.go                   |  7 +++++++
 5 files changed, 41 insertions(+)

diff --git a/internal/command/worker/run.go b/internal/command/worker/run.go
index c96feef..f704d14 100644
--- a/internal/command/worker/run.go
+++ b/internal/command/worker/run.go
@@ -17,6 +17,8 @@ import (
 	"go.uber.org/zap/zapcore"
 	"gopkg.in/natefinch/lumberjack.v2"
 	"io"
+	"net/http"
+	_ "net/http/pprof"
 	"os"
 	"strings"
 )
@@ -37,6 +39,7 @@ var noPKI bool
 var defaultCPU uint64
 var defaultMemory uint64
 var username string
+var addressPprof string
 var debug bool
 
 func newRunCommand() *cobra.Command {
@@ -71,6 +74,8 @@ func newRunCommand() *cobra.Command {
 		"(\"Local Network\" permission workaround: requires starting \"orchard worker run\" as \"root\", "+
 		"the privileges will be then dropped to the specified user after starting the \"orchard localnetworkhelper\" "+
 		"helper process)")
+	cmd.Flags().StringVar(&addressPprof, "listen-pprof", "",
+		"start pprof HTTP server on localhost:6060 for diagnostic purposes (e.g. \"localhost:6060\")")
 	cmd.Flags().BoolVar(&debug, "debug", false, "enable debug logging")
 
 	return cmd
@@ -154,6 +159,14 @@ func runWorker(cmd *cobra.Command, args []string) (err error) {
 	}()
 	workerOpts = append(workerOpts, worker.WithLogger(logger))
 
+	if addressPprof != "" {
+		go func() {
+			if err := http.ListenAndServe(addressPprof, nil); err != nil {
+				logger.Sugar().Errorf("pprof server failed: %v", err)
+			}
+		}()
+	}
+
 	workerInstance, err := worker.New(
 		controllerClient,
 		workerOpts...,
diff --git a/internal/controller/api_rpc_portforward.go b/internal/controller/api_rpc_portforward.go
index af746cf..a29f0bd 100644
--- a/internal/controller/api_rpc_portforward.go
+++ b/internal/controller/api_rpc_portforward.go
@@ -27,6 +27,13 @@ func (controller *Controller) rpcPortForward(ctx *gin.Context) responder.Respond
 	if err != nil {
 		return responder.Error(err)
 	}
+	defer func() {
+		// Ensure that we always close the accepted WebSocket connection,
+		// otherwise resource leak is possible[1]
+		//
+		// [1]: https://github.com/coder/websocket/issues/445#issuecomment-2053792044
+		_ = wsConn.CloseNow()
+	}()
 
 	// Respond with the established connection
 	proxyCtx, err := controller.connRendezvous.Respond(session, rendezvous.ResultWithErrorMessage[net.Conn]{
diff --git a/internal/controller/api_rpc_watch.go b/internal/controller/api_rpc_watch.go
index c46c162..d7d0d18 100644
--- a/internal/controller/api_rpc_watch.go
+++ b/internal/controller/api_rpc_watch.go
@@ -37,6 +37,13 @@ func (controller *Controller) rpcWatch(ctx *gin.Context) responder.Responder {
 	if err != nil {
 		return responder.Error(err)
 	}
+	defer func() {
+		// Ensure that we always close the accepted WebSocket connection,
+		// otherwise resource leak is possible[1]
+		//
+		// [1]: https://github.com/coder/websocket/issues/445#issuecomment-2053792044
+		_ = wsConn.CloseNow()
+	}()
 
 	// Ensure that pongs will be processed by reading
 	// from the connection in the background
diff --git a/internal/controller/api_vms_portforward.go b/internal/controller/api_vms_portforward.go
index 1a94532..deeb2e6 100644
--- a/internal/controller/api_vms_portforward.go
+++ b/internal/controller/api_vms_portforward.go
@@ -102,6 +102,13 @@ func (controller *Controller) portForward(
 		if err != nil {
 			return responder.Error(err)
 		}
+		defer func() {
+			// Ensure that we always close the accepted WebSocket connection,
+			// otherwise resource leak is possible[1]
+			//
+			// [1]: https://github.com/coder/websocket/issues/445#issuecomment-2053792044
+			_ = wsConn.CloseNow()
+		}()
 
 		expectedMsgType := websocket.MessageBinary
 
diff --git a/internal/worker/rpcv2.go b/internal/worker/rpcv2.go
index a264746..bed413d 100644
--- a/internal/worker/rpcv2.go
+++ b/internal/worker/rpcv2.go
@@ -55,6 +55,13 @@ func (worker *Worker) handlePortForwardV2(ctx context.Context, portForward *v1.P
 
 		return
 	}
+	defer func() {
+		// Ensure that we always close the accepted WebSocket connection,
+		// otherwise resource leak is possible[1]
+		//
+		// [1]: https://github.com/coder/websocket/issues/445#issuecomment-2053792044
+		_ = netConn.Close()
+	}()
 
 	// Proxy bytes if the connection was established without errors
 	if errorMessage == "" {