fix: use a native websocket.NetConn for agent RPC client (#13142)

spikecurtis · web-flow · commit e76b5950527a · 2024-05-06T15:00:34.000+04:00
One cause of #13139 is a peculiar failure mode of `WebsocketNetConn` which causes it to return `context.Canceled` in some circumstances when the underlying websocket fails. We have special processing for that error in the `agent.run()` routine, which is erroneously being triggered. Since we don't actually need the returned context from `WebsocketNetConn`, we can simplify and just use the netConn from the `websocket` library directly.
diff --git a/codersdk/agentsdk/agentsdk.go b/codersdk/agentsdk/agentsdk.go
@@ -5,7 +5,6 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
-	"net"
 	"net/http"
 	"net/http/cookiejar"
 	"net/url"
@@ -206,14 +205,11 @@ func (c *Client) ConnectRPC(ctx context.Context) (drpc.Conn, error) {
 		return nil, codersdk.ReadBodyAsError(res)
 	}
 
-	_, wsNetConn := codersdk.WebsocketNetConn(ctx, conn, websocket.MessageBinary)
+	// Set the read limit to 4 MiB -- about the limit for protobufs.  This needs to be larger than
+	// the default because some of our protocols can include large messages like startup scripts.
+	conn.SetReadLimit(1 << 22)
+	netConn := websocket.NetConn(ctx, conn, websocket.MessageBinary)
 
-	netConn := &closeNetConn{
-		Conn: wsNetConn,
-		closeFunc: func() {
-			_ = conn.Close(websocket.StatusGoingAway, "ConnectRPC closed")
-		},
-	}
 	config := yamux.DefaultConfig()
 	config.LogOutput = nil
 	config.Logger = slog.Stdlib(ctx, c.SDK.Logger(), slog.LevelInfo)
@@ -618,13 +614,3 @@ func LogsNotifyChannel(agentID uuid.UUID) string {
 type LogsNotifyMessage struct {
 	CreatedAfter int64 `json:"created_after"`
 }
-
-type closeNetConn struct {
-	net.Conn
-	closeFunc func()
-}
-
-func (c *closeNetConn) Close() error {
-	c.closeFunc()
-	return c.Conn.Close()
-}
