- Notifications
You must be signed in to change notification settings - Fork928
chore: consolidate websocketNetConn implementations#12065
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to ourterms of service andprivacy statement. We’ll occasionally send you account related emails.
Already on GitHub?Sign in to your account
Uh oh!
There was an error while loading.Please reload this page.
Changes fromall commits
File filter
Filter by extension
Conversations
Uh oh!
There was an error while loading.Please reload this page.
Jump to
Uh oh!
There was an error while loading.Please reload this page.
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,53 @@ | ||
| package codersdk | ||
| import ( | ||
| "context" | ||
| "net" | ||
| "nhooyr.io/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 { | ||
| 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, | ||
| } | ||
| } |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,80 @@ | ||
| package codersdk_test | ||
| import ( | ||
| "crypto/rand" | ||
| "net/http" | ||
| "net/http/httptest" | ||
| "testing" | ||
| "github.com/stretchr/testify/assert" | ||
| "github.com/stretchr/testify/require" | ||
| "nhooyr.io/websocket" | ||
| "github.com/coder/coder/v2/codersdk" | ||
| "github.com/coder/coder/v2/testutil" | ||
| ) | ||
| // TestWebsocketNetConn_LargeWrites tests that we can write large amounts of data thru the netconn | ||
| // in a single write. Without specifically setting the read limit, the websocket library limits | ||
| // the amount of data that can be read in a single message to 32kiB. Even after raising the limit, | ||
| // curiously, it still only reads 32kiB per Read(), but allows the large write to go thru. | ||
| func TestWebsocketNetConn_LargeWrites(t *testing.T) { | ||
| t.Parallel() | ||
| ctx := testutil.Context(t, testutil.WaitShort) | ||
| n := 4 * 1024 * 1024 // 4 MiB | ||
| svr := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { | ||
| sws, err := websocket.Accept(w, r, nil) | ||
| if !assert.NoError(t, err) { | ||
| return | ||
| } | ||
| _, nc := codersdk.WebsocketNetConn(r.Context(), sws, websocket.MessageBinary) | ||
| defer nc.Close() | ||
| // Although the writes are all in one go, the reads get broken up by | ||
| // the library. | ||
| j := 0 | ||
| b := make([]byte, n) | ||
| for j < n { | ||
| k, err := nc.Read(b[j:]) | ||
| if !assert.NoError(t, err) { | ||
| return | ||
| } | ||
| j += k | ||
| t.Logf("server read %d bytes, total %d", k, j) | ||
| } | ||
| assert.Equal(t, n, j) | ||
| j, err = nc.Write(b) | ||
| assert.Equal(t, n, j) | ||
| if !assert.NoError(t, err) { | ||
| return | ||
| } | ||
| })) | ||
| // use of random data is worst case scenario for compression | ||
| cb := make([]byte, n) | ||
| rk, err := rand.Read(cb) | ||
| require.NoError(t, err) | ||
| require.Equal(t, n, rk) | ||
| // nolint: bodyclose | ||
| cws, _, err := websocket.Dial(ctx, svr.URL, nil) | ||
| require.NoError(t, err) | ||
| _, cnc := codersdk.WebsocketNetConn(ctx, cws, websocket.MessageBinary) | ||
| ck, err := cnc.Write(cb) | ||
| require.NoError(t, err) | ||
| require.Equal(t, n, ck) | ||
| cb2 := make([]byte, n) | ||
| j := 0 | ||
| for j < n { | ||
| k, err := cnc.Read(cb2[j:]) | ||
| if !assert.NoError(t, err) { | ||
| return | ||
| } | ||
| j += k | ||
| t.Logf("client read %d bytes, total %d", k, j) | ||
| } | ||
| require.NoError(t, err) | ||
| require.Equal(t, n, j) | ||
| require.Equal(t, cb, cb2) | ||
| } |
| Original file line number | Diff line number | Diff line change | ||||
|---|---|---|---|---|---|---|
| @@ -50,7 +50,7 @@ func TestTailnetAPIConnector_Disconnects(t *testing.T) { | ||||||
| if !assert.NoError(t, err) { | ||||||
| return | ||||||
| } | ||||||
| ctx, nc :=WebsocketNetConn(r.Context(), sws, websocket.MessageBinary) | ||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others.Learn more. I'd prefer to fix this bad use of Suggested change
I know some other places are using Perhaps this is actually a case for moving There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others.Learn more.
I don't understand what it means to "invalidate" a context cancelation. Are you referring to the idea that once we have a websocket, passing Are you suggesting we don't accept a context at all? Member There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others.Learn more. The context is actually not even tied to the TCP connection, it is tied to the handler which creates a scenario where it’s unlikely to ever be cancelled in a way that matters. https://pkg.go.dev/net/http#Hijacker Functionally, you’re right, there’s no difference but this behavior can easily trip anyone up so its usage should be discouraged. | ||||||
| err = svc.ServeConnV2(ctx, nc, tailnet.StreamID{ | ||||||
| Name: "client", | ||||||
| ID: clientID, | ||||||
Uh oh!
There was an error while loading.Please reload this page.