mirror of
https://github.com/coder/coder.git
synced 2025-03-14 10:09:57 +00:00
2c7f8ac65f
Migrates us to `coder/websocket` v1.8.12 rather than `nhooyr/websocket` on an older version. Works around https://github.com/coder/websocket/issues/504 by adding an explicit test for `xerrors.Is(err, io.EOF)` where we were previously getting `io.EOF` from the netConn.
54 lines
1.4 KiB
Go
54 lines
1.4 KiB
Go
package codersdk
|
|
|
|
import (
|
|
"context"
|
|
"net"
|
|
|
|
"github.com/coder/websocket"
|
|
)
|
|
|
|
// wsNetConn wraps net.Conn created by websocket.NetConn(). Cancel func
|
|
// is called if a read or write error is encountered.
|
|
// @typescript-ignore wsNetConn
|
|
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 {
|
|
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,
|
|
}
|
|
}
|