From 667b7187d8974abe4c21bf823db9b7cae64db5b7 Mon Sep 17 00:00:00 2001
From: Spike Curtis <spike@coder.com>
Date: Fri, 3 May 2024 13:35:48 +0400
Subject: [PATCH] fix: use a native websocket.NetConn for agent RPC client

---
 codersdk/agentsdk/agentsdk.go | 22 ++++------------------
 1 file changed, 4 insertions(+), 18 deletions(-)

diff --git a/codersdk/agentsdk/agentsdk.go b/codersdk/agentsdk/agentsdk.go
index 75bec0047eb5b..5dcccca09e350 100644
--- 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()
-}