derp/derphttp: wait for websocket connections to exit on close

kylecarbs · kylecarbs · commit a7f70662221a · 2023-03-07T01:47:41.000Z
This was causing a leak in our CI!
diff --git a/derp/derphttp/derphttp_client.go b/derp/derphttp/derphttp_client.go
@@ -56,6 +56,7 @@ type Client struct {
 
 	forcedWebsocket         atomic.Bool // optional; set if the server has failed to upgrade the connection on the DERP server
 	forcedWebsocketCallback atomic.Pointer[func(int, string)]
+	websocketsClosed        sync.WaitGroup
 
 	privateKey key.NodePrivate
 	logf       logger.Logf
@@ -331,6 +332,12 @@ func (c *Client) connect(ctx context.Context, caller string) (client *derp.Clien
 			c.logf("%s: websocket to %v error: %v", caller, urlStr, err)
 			return nil, 0, err
 		}
+		c.websocketsClosed.Add(1)
+		defer func() {
+			conn.Close()
+			c.websocketsClosed.Done()
+		}()
+
 		brw := bufio.NewReadWriter(bufio.NewReader(conn), bufio.NewWriter(conn))
 		derpClient, err := derp.NewClient(c.privateKey, conn, brw, c.logf,
 			derp.MeshKey(c.MeshKey),
@@ -1051,6 +1058,7 @@ func (c *Client) Close() error {
 	if c.netConn != nil {
 		c.netConn.Close()
 	}
+	c.websocketsClosed.Wait()
 	return nil
 }
 
