From c3de6b385065393b5e02547753b64006970077bc Mon Sep 17 00:00:00 2001
From: Kyle Carberry <kyle@carberry.com>
Date: Thu, 30 Mar 2023 17:12:25 +0000
Subject: [PATCH] chore: ensure agent conn routine is closed before exit

This caused a leak in `main`!
---
 codersdk/agentsdk/agentsdk.go | 21 ++++++++++++++++++++-
 1 file changed, 20 insertions(+), 1 deletion(-)

diff --git a/codersdk/agentsdk/agentsdk.go b/codersdk/agentsdk/agentsdk.go
index 6ebec9a42f26a..1796baea60ac6 100644
--- a/codersdk/agentsdk/agentsdk.go
+++ b/codersdk/agentsdk/agentsdk.go
@@ -161,12 +161,15 @@ func (c *Client) Listen(ctx context.Context) (net.Conn, error) {
 		return nil, codersdk.ReadBodyAsError(res)
 	}
 
+	ctx, cancelFunc := context.WithCancel(ctx)
 	ctx, wsNetConn := websocketNetConn(ctx, conn, websocket.MessageBinary)
 
 	// Ping once every 30 seconds to ensure that the websocket is alive. If we
 	// don't get a response within 30s we kill the websocket and reconnect.
 	// See: https://github.com/coder/coder/pull/5824
+	closed := make(chan struct{})
 	go func() {
+		defer close(closed)
 		tick := 30 * time.Second
 		ticker := time.NewTicker(tick)
 		defer ticker.Stop()
@@ -199,7 +202,13 @@ func (c *Client) Listen(ctx context.Context) (net.Conn, error) {
 		}
 	}()
 
-	return wsNetConn, nil
+	return &closeNetConn{
+		Conn: wsNetConn,
+		closeFunc: func() {
+			cancelFunc()
+			<-closed
+		},
+	}, nil
 }
 
 type PostAppHealthsRequest struct {
@@ -623,3 +632,13 @@ type StartupLogsNotifyMessage struct {
 	CreatedAfter int64 `json:"created_after"`
 	EndOfLogs    bool  `json:"end_of_logs"`
 }
+
+type closeNetConn struct {
+	net.Conn
+	closeFunc func()
+}
+
+func (c *closeNetConn) Close() error {
+	c.closeFunc()
+	return c.Conn.Close()
+}