From e76b5950527aa7beeaffc8a87a2a83f1c078cd46 Mon Sep 17 00:00:00 2001 From: Spike Curtis Date: Mon, 6 May 2024 15:00:34 +0400 Subject: [PATCH] fix: use a native websocket.NetConn for agent RPC client (#13142) 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. --- 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 75bec0047e..5dcccca09e 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() -}