synced/coder
1
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2025-07-15 22:20:27 +00:00
Code Issues Packages Projects Releases Wiki Activity

chore: consolidate websocketNetConn implementations (#12065)

Browse Source 
Consolidates websocketNetConn from multiple packages in favor of a central one in codersdk
This commit is contained in:
Spike Curtis
2024-02-09 11:39:08 +04:00
committed by GitHub
parent ec8e41f516
commit 1f5a6d59ba
10 changed files with 145 additions and 184 deletions
Download Patch File Download Diff File Expand all files Collapse all files

46
codersdk/agentsdk/agentsdk.go
View File

@ -203,7 +203,7 @@ func (c *Client) ConnectRPC(ctx context.Context) (drpc.Conn, error) {
return nil, codersdk.ReadBodyAsError(res)
}
_, wsNetConn := websocketNetConn(ctx, conn, websocket.MessageBinary)
_, wsNetConn := codersdk.WebsocketNetConn(ctx, conn, websocket.MessageBinary)
netConn := &closeNetConn{
Conn: wsNetConn,
@ -596,50 +596,6 @@ func (c *Client) ExternalAuth(ctx context.Context, req ExternalAuthRequest) (Ext
return authResp, json.NewDecoder(res.Body).Decode(&authResp)
}
// wsNetConn wraps net.Conn created by websocket.NetConn(). Cancel func
// is called if a read or write error is encountered.
type wsNetConn struct {
cancel context.CancelFunc
net.Conn
}
func (c *wsNetConn) Read(b []byte) (n int, err error) {
n, err = c.Conn.Read(b)
if err != nil {
c.cancel()
}
return n, err
}
func (c *wsNetConn) Write(b []byte) (n int, err error) {
n, err = c.Conn.Write(b)
if err != nil {
c.cancel()
}
return n, err
}
func (c *wsNetConn) Close() error {
defer c.cancel()
return c.Conn.Close()
}
// websocketNetConn wraps websocket.NetConn and returns a context that
// is tied to the parent context and the lifetime of the conn. Any error
// during read or write will cancel the context, but not close the
// conn. Close should be called to release context resources.
func websocketNetConn(ctx context.Context, conn *websocket.Conn, msgType websocket.MessageType) (context.Context, net.Conn) {
// 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)
ctx, cancel := context.WithCancel(ctx)
nc := websocket.NetConn(ctx, conn, msgType)
return ctx, &wsNetConn{
cancel: cancel,
Conn: nc,
}
}
// LogsNotifyChannel returns the channel name responsible for notifying
// of new logs.
func LogsNotifyChannel(agentID uuid.UUID) string {