From 18dd871282982edf449d5bc9fc8350a1bdd61afb Mon Sep 17 00:00:00 2001
From: Spike Curtis <spike@coder.com>
Date: Fri, 2 Jun 2023 08:26:19 +0000
Subject: [PATCH] fix: select on context to avoid leak in logFollower

Signed-off-by: Spike Curtis <spike@coder.com>
---
 coderd/provisionerjobs.go | 20 +++++++++++++++++---
 1 file changed, 17 insertions(+), 3 deletions(-)

diff --git a/coderd/provisionerjobs.go b/coderd/provisionerjobs.go
index ba6e8cb5ce8f9..f667c2ec65d72 100644
--- a/coderd/provisionerjobs.go
+++ b/coderd/provisionerjobs.go
@@ -300,6 +300,9 @@ func newLogFollower(
 }
 
 func (f *logFollower) follow() {
+	var cancel context.CancelFunc
+	f.ctx, cancel = context.WithCancel(f.ctx)
+	defer cancel()
 	// note that we only need to subscribe to updates if the job is not yet
 	// complete.
 	if !f.complete {
@@ -405,17 +408,28 @@ func (f *logFollower) follow() {
 }
 
 func (f *logFollower) listener(_ context.Context, message []byte, err error) {
+	// in this function we always pair writes to channels with a select on the context
+	// otherwise we could block a goroutine if the follow() method exits.
 	if err != nil {
-		f.errors <- err
+		select {
+		case <-f.ctx.Done():
+		case f.errors <- err:
+		}
 		return
 	}
 	var n provisionersdk.ProvisionerJobLogsNotifyMessage
 	err = json.Unmarshal(message, &n)
 	if err != nil {
-		f.errors <- err
+		select {
+		case <-f.ctx.Done():
+		case f.errors <- err:
+		}
 		return
 	}
-	f.notifications <- n
+	select {
+	case <-f.ctx.Done():
+	case f.notifications <- n:
+	}
 }
 
 // query fetches the latest job logs from the database and writes them to the