-
Notifications
You must be signed in to change notification settings - Fork 886
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 our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
+145
−184
Merged
Changes from all commits
Commits
File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| 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, | ||
| } | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| 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) | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I'd prefer to fix this bad use of
r.Context()since someone could look at this and get the wrong idea. (Thewebsocket.Acceptinvalidatesr.Context()cancellation.)I know some other places are using
r.Context()viactx := r.Context()too, but this one just seems too explicitly wrong. 😄Perhaps this is actually a case for moving
websocket.Acceptintocodersdk.WebsocketNetConnas well (or creating a unified function,codersdk.WebsocketAcceptNetConn).There was a problem hiding this comment.
Choose a reason for hiding this comment
The 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
r.Context()is a bit pointless because the context won't get canceled before the underlying TCP connection is closed? I guess that's true enough, but changing it tocontext.Background()doesn't change anything from a functional standpoint.Are you suggesting we don't accept a context at all?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The 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.