Fix WebSocket not closing

kylecarbs · kylecarbs · commit 0b6b47072ce4 · 2022-09-19T23:49:54.000Z
diff --git a/agent/agent_test.go b/agent/agent_test.go
@@ -522,7 +522,7 @@ func setupSSHCommand(t *testing.T, beforeArgs []string, afterArgs []string) *exe
 				return
 			}
 			ssh, err := agentConn.SSH()
-			if !assert.NoError(t, err) {
+			if err != nil {
 				_ = conn.Close()
 				return
 			}
@@ -581,11 +581,16 @@ func setupAgent(t *testing.T, metadata agent.Metadata, ptyTimeout time.Duration)
 		},
 		CoordinatorDialer: func(ctx context.Context) (net.Conn, error) {
 			clientConn, serverConn := net.Pipe()
+			closed := make(chan struct{})
 			t.Cleanup(func() {
 				_ = serverConn.Close()
 				_ = clientConn.Close()
+				<-closed
 			})
-			go coordinator.ServeAgent(serverConn, agentID)
+			go func() {
+				_ = coordinator.ServeAgent(serverConn, agentID)
+				close(closed)
+			}()
 			return clientConn, nil
 		},
 		Logger:                 slogtest.Make(t, nil).Leveled(slog.LevelDebug),
diff --git a/codersdk/provisionerdaemons.go b/codersdk/provisionerdaemons.go
@@ -135,6 +135,7 @@ func (c *Client) provisionerJobLogsAfter(ctx context.Context, path string, after
 	decoder := json.NewDecoder(websocket.NetConn(ctx, conn, websocket.MessageText))
 	go func() {
 		defer close(logs)
+		defer conn.Close(websocket.StatusGoingAway, "")
 		var log ProvisionerJobLog
 		for {
 			err = decoder.Decode(&log)
diff --git a/codersdk/workspaceagents.go b/codersdk/workspaceagents.go
@@ -281,10 +281,12 @@ func (c *Client) DialWorkspaceAgentTailnet(ctx context.Context, logger slog.Logg
 				CompressionMode: websocket.CompressionDisabled,
 			})
 			if errors.Is(err, context.Canceled) {
+				_ = ws.Close(websocket.StatusAbnormalClosure, "")
 				return
 			}
 			if err != nil {
 				logger.Debug(ctx, "failed to dial", slog.Error(err))
+				_ = ws.Close(websocket.StatusAbnormalClosure, "")
 				continue
 			}
 			sendNode, errChan := tailnet.ServeCoordinator(websocket.NetConn(ctx, ws, websocket.MessageBinary), func(node []*tailnet.Node) error {
@@ -294,12 +296,15 @@ func (c *Client) DialWorkspaceAgentTailnet(ctx context.Context, logger slog.Logg
 			logger.Debug(ctx, "serving coordinator")
 			err = <-errChan
 			if errors.Is(err, context.Canceled) {
+				_ = ws.Close(websocket.StatusAbnormalClosure, "")
 				return
 			}
 			if err != nil {
 				logger.Debug(ctx, "error serving coordinator", slog.Error(err))
+				_ = ws.Close(websocket.StatusAbnormalClosure, "")
 				continue
 			}
+			_ = ws.Close(websocket.StatusAbnormalClosure, "")
 		}
 	}()
 	return &agent.Conn{
@@ -423,6 +428,7 @@ func (c *Client) AgentReportStats(
 					var req AgentStatsReportRequest
 					err := wsjson.Read(ctx, conn, &req)
 					if err != nil {
+						_ = conn.Close(websocket.StatusAbnormalClosure, "")
 						return err
 					}
 
@@ -436,6 +442,7 @@ func (c *Client) AgentReportStats(
 
 					err = wsjson.Write(ctx, conn, resp)
 					if err != nil {
+						_ = conn.Close(websocket.StatusAbnormalClosure, "")
 						return err
 					}
 				}

Original file line number	Diff line number	Diff line change
`@@ -281,10 +281,12 @@ func (c *Client) DialWorkspaceAgentTailnet(ctx context.Context, logger slog.Logg`
`281`	`281`	`CompressionMode: websocket.CompressionDisabled,`
`282`	`282`	`})`
`283`	`283`	`if errors.Is(err, context.Canceled) {`
	`284`	`+ _ = ws.Close(websocket.StatusAbnormalClosure, "")`
`284`	`285`	`return`
`285`	`286`	`}`
`286`	`287`	`if err != nil {`
`287`	`288`	`logger.Debug(ctx, "failed to dial", slog.Error(err))`
	`289`	`+ _ = ws.Close(websocket.StatusAbnormalClosure, "")`
`288`	`290`	`continue`
`289`	`291`	`}`
`290`	`292`	`sendNode, errChan := tailnet.ServeCoordinator(websocket.NetConn(ctx, ws, websocket.MessageBinary), func(node []*tailnet.Node) error {`
`@@ -294,12 +296,15 @@ func (c *Client) DialWorkspaceAgentTailnet(ctx context.Context, logger slog.Logg`
`294`	`296`	`logger.Debug(ctx, "serving coordinator")`
`295`	`297`	`err = <-errChan`
`296`	`298`	`if errors.Is(err, context.Canceled) {`
	`299`	`+ _ = ws.Close(websocket.StatusAbnormalClosure, "")`
`297`	`300`	`return`
`298`	`301`	`}`
`299`	`302`	`if err != nil {`
`300`	`303`	`logger.Debug(ctx, "error serving coordinator", slog.Error(err))`
	`304`	`+ _ = ws.Close(websocket.StatusAbnormalClosure, "")`
`301`	`305`	`continue`
`302`	`306`	`}`
	`307`	`+ _ = ws.Close(websocket.StatusAbnormalClosure, "")`
`303`	`308`	`}`
`304`	`309`	`}()`
`305`	`310`	`return &agent.Conn{`
`@@ -423,6 +428,7 @@ func (c *Client) AgentReportStats(`
`423`	`428`	`var req AgentStatsReportRequest`
`424`	`429`	`err := wsjson.Read(ctx, conn, &req)`
`425`	`430`	`if err != nil {`
	`431`	`+ _ = conn.Close(websocket.StatusAbnormalClosure, "")`
`426`	`432`	`return err`
`427`	`433`	`}`
`428`	`434`
`@@ -436,6 +442,7 @@ func (c *Client) AgentReportStats(`
`436`	`442`
`437`	`443`	`err = wsjson.Write(ctx, conn, resp)`
`438`	`444`	`if err != nil {`
	`445`	`+ _ = conn.Close(websocket.StatusAbnormalClosure, "")`
`439`	`446`	`return err`
`440`	`447`	`}`
`441`	`448`	`}`