fix: use a native websocket.NetConn for agent RPC client

spikecurtis · spikecurtis · commit f68443946951 · 2024-05-03T14:28:39.000+04:00
diff --git a/codersdk/agentsdk/agentsdk.go b/codersdk/agentsdk/agentsdk.go
@@ -206,14 +206,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)
