Skip to content

Commit 89920d9

Browse files
spikecurtisspikecurtis
committed
chore: consolidate websocketNetConn implementations
1 parent 151aaad commit 89920d9

File tree

8 files changed

+17
-143
lines changed

8 files changed

+17
-143
lines changed

‎coderd/workspaceagents.go

Lines changed: 4 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@ import (
77
"errors"
88
"fmt"
99
"io"
10-
"net"
1110
"net/http"
1211
"net/url"
1312
"sort"
@@ -544,7 +543,7 @@ func (api *API) workspaceAgentLogs(rw http.ResponseWriter, r *http.Request) {
544543
}
545544
go httpapi.Heartbeat(ctx, conn)
546545

547-
ctx, wsNetConn := websocketNetConn(ctx, conn, websocket.MessageText)
546+
ctx, wsNetConn := codersdk.WebsocketNetConn(ctx, conn, websocket.MessageText)
548547
defer wsNetConn.Close() // Also closes conn.
549548

550549
// The Go stdlib JSON encoder appends a newline character after message write.
@@ -881,7 +880,7 @@ func (api *API) derpMapUpdates(rw http.ResponseWriter, r *http.Request) {
881880
})
882881
return
883882
}
884-
ctx, nconn := websocketNetConn(ctx, ws, websocket.MessageBinary)
883+
ctx, nconn := codersdk.WebsocketNetConn(ctx, ws, websocket.MessageBinary)
885884
defer nconn.Close()
886885

887886
// Slurp all packets from the connection into io.Discard so pongs get sent
@@ -990,7 +989,7 @@ func (api *API) workspaceAgentCoordinate(rw http.ResponseWriter, r *http.Request
990989
return
991990
}
992991

993-
ctx, wsNetConn := websocketNetConn(ctx, conn, websocket.MessageBinary)
992+
ctx, wsNetConn := codersdk.WebsocketNetConn(ctx, conn, websocket.MessageBinary)
994993
defer wsNetConn.Close()
995994

996995
closeCtx, closeCtxCancel := context.WithCancel(ctx)
@@ -1077,7 +1076,7 @@ func (api *API) workspaceAgentClientCoordinate(rw http.ResponseWriter, r *http.R
10771076
})
10781077
return
10791078
}
1080-
ctx, wsNetConn := websocketNetConn(ctx, conn, websocket.MessageBinary)
1079+
ctx, wsNetConn := codersdk.WebsocketNetConn(ctx, conn, websocket.MessageBinary)
10811080
defer wsNetConn.Close()
10821081

10831082
go httpapi.Heartbeat(ctx, conn)
@@ -2104,47 +2103,6 @@ func createExternalAuthResponse(typ, token string, extra pqtype.NullRawMessage)
21042103
return resp, err
21052104
}
21062105

2107-
// wsNetConn wraps net.Conn created by websocket.NetConn(). Cancel func
2108-
// is called if a read or write error is encountered.
2109-
type wsNetConn struct {
2110-
cancel context.CancelFunc
2111-
net.Conn
2112-
}
2113-
2114-
func (c *wsNetConn) Read(b []byte) (n int, err error) {
2115-
n, err = c.Conn.Read(b)
2116-
if err != nil {
2117-
c.cancel()
2118-
}
2119-
return n, err
2120-
}
2121-
2122-
func (c *wsNetConn) Write(b []byte) (n int, err error) {
2123-
n, err = c.Conn.Write(b)
2124-
if err != nil {
2125-
c.cancel()
2126-
}
2127-
return n, err
2128-
}
2129-
2130-
func (c *wsNetConn) Close() error {
2131-
defer c.cancel()
2132-
return c.Conn.Close()
2133-
}
2134-
2135-
// websocketNetConn wraps websocket.NetConn and returns a context that
2136-
// is tied to the parent context and the lifetime of the conn. Any error
2137-
// during read or write will cancel the context, but not close the
2138-
// conn. Close should be called to release context resources.
2139-
func websocketNetConn(ctx context.Context, conn *websocket.Conn, msgType websocket.MessageType) (context.Context, net.Conn) {
2140-
ctx, cancel := context.WithCancel(ctx)
2141-
nc := websocket.NetConn(ctx, conn, msgType)
2142-
return ctx, &wsNetConn{
2143-
cancel: cancel,
2144-
Conn: nc,
2145-
}
2146-
}
2147-
21482106
func convertWorkspaceAgentLogs(logs []database.WorkspaceAgentLog) []codersdk.WorkspaceAgentLog {
21492107
sdk := make([]codersdk.WorkspaceAgentLog, 0, len(logs))
21502108
for _, logEntry := range logs {

‎coderd/workspaceagentsrpc.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -100,7 +100,7 @@ func (api *API) workspaceAgentRPC(rw http.ResponseWriter, r *http.Request) {
100100
return
101101
}
102102

103-
ctx, wsNetConn := websocketNetConn(ctx, conn, websocket.MessageBinary)
103+
ctx, wsNetConn := codersdk.WebsocketNetConn(ctx, conn, websocket.MessageBinary)
104104
defer wsNetConn.Close()
105105

106106
ycfg := yamux.DefaultConfig()

‎codersdk/agentsdk/agentsdk.go

Lines changed: 1 addition & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -203,7 +203,7 @@ func (c *Client) ConnectRPC(ctx context.Context) (drpc.Conn, error) {
203203
return nil, codersdk.ReadBodyAsError(res)
204204
}
205205

206-
_, wsNetConn := websocketNetConn(ctx, conn, websocket.MessageBinary)
206+
_, wsNetConn := codersdk.WebsocketNetConn(ctx, conn, websocket.MessageBinary)
207207

208208
netConn := &closeNetConn{
209209
Conn: wsNetConn,
@@ -596,50 +596,6 @@ func (c *Client) ExternalAuth(ctx context.Context, req ExternalAuthRequest) (Ext
596596
return authResp, json.NewDecoder(res.Body).Decode(&authResp)
597597
}
598598

599-
// wsNetConn wraps net.Conn created by websocket.NetConn(). Cancel func
600-
// is called if a read or write error is encountered.
601-
type wsNetConn struct {
602-
cancel context.CancelFunc
603-
net.Conn
604-
}
605-
606-
func (c *wsNetConn) Read(b []byte) (n int, err error) {
607-
n, err = c.Conn.Read(b)
608-
if err != nil {
609-
c.cancel()
610-
}
611-
return n, err
612-
}
613-
614-
func (c *wsNetConn) Write(b []byte) (n int, err error) {
615-
n, err = c.Conn.Write(b)
616-
if err != nil {
617-
c.cancel()
618-
}
619-
return n, err
620-
}
621-
622-
func (c *wsNetConn) Close() error {
623-
defer c.cancel()
624-
return c.Conn.Close()
625-
}
626-
627-
// websocketNetConn wraps websocket.NetConn and returns a context that
628-
// is tied to the parent context and the lifetime of the conn. Any error
629-
// during read or write will cancel the context, but not close the
630-
// conn. Close should be called to release context resources.
631-
func websocketNetConn(ctx context.Context, conn *websocket.Conn, msgType websocket.MessageType) (context.Context, net.Conn) {
632-
// Set the read limit to 4 MiB -- about the limit for protobufs. This needs to be larger than
633-
// the default because some of our protocols can include large messages like startup scripts.
634-
conn.SetReadLimit(1 << 22)
635-
ctx, cancel := context.WithCancel(ctx)
636-
nc := websocket.NetConn(ctx, conn, msgType)
637-
return ctx, &wsNetConn{
638-
cancel: cancel,
639-
Conn: nc,
640-
}
641-
}
642-
643599
// LogsNotifyChannel returns the channel name responsible for notifying
644600
// of new logs.
645601
func LogsNotifyChannel(agentID uuid.UUID) string {

‎codersdk/provisionerdaemons.go

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -248,7 +248,7 @@ func (c *Client) ServeProvisionerDaemon(ctx context.Context, req ServeProvisione
248248
config := yamux.DefaultConfig()
249249
config.LogOutput = io.Discard
250250
// Use background context because caller should close the client.
251-
_, wsNetConn := websocketNetConn(context.Background(), conn, websocket.MessageBinary)
251+
_, wsNetConn := WebsocketNetConn(context.Background(), conn, websocket.MessageBinary)
252252
session, err := yamux.Client(wsNetConn, config)
253253
if err != nil {
254254
_ = conn.Close(websocket.StatusGoingAway, "")
@@ -287,11 +287,14 @@ func (c *wsNetConn) Close() error {
287287
return c.Conn.Close()
288288
}
289289

290-
// websocketNetConn wraps websocket.NetConn and returns a context that
290+
// WebsocketNetConn wraps websocket.NetConn and returns a context that
291291
// is tied to the parent context and the lifetime of the conn. Any error
292292
// during read or write will cancel the context, but not close the
293293
// conn. Close should be called to release context resources.
294-
func websocketNetConn(ctx context.Context, conn *websocket.Conn, msgType websocket.MessageType) (context.Context, net.Conn) {
294+
func WebsocketNetConn(ctx context.Context, conn *websocket.Conn, msgType websocket.MessageType) (context.Context, net.Conn) {
295+
// Set the read limit to 4 MiB -- about the limit for protobufs. This needs to be larger than
296+
// the default because some of our protocols can include large messages like startup scripts.
297+
conn.SetReadLimit(1 << 22)
295298
ctx, cancel := context.WithCancel(ctx)
296299
nc := websocket.NetConn(ctx, conn, msgType)
297300
return ctx, &wsNetConn{

‎codersdk/workspaceagents.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -844,7 +844,7 @@ func (c *Client) WorkspaceAgentLogsAfter(ctx context.Context, agentID uuid.UUID,
844844
}
845845
logChunks := make(chan []WorkspaceAgentLog, 1)
846846
closed := make(chan struct{})
847-
ctx, wsNetConn := websocketNetConn(ctx, conn, websocket.MessageText)
847+
ctx, wsNetConn := WebsocketNetConn(ctx, conn, websocket.MessageText)
848848
decoder := json.NewDecoder(wsNetConn)
849849
go func() {
850850
defer close(closed)

‎codersdk/workspaceagents_internal_test.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,7 @@ func TestTailnetAPIConnector_Disconnects(t *testing.T) {
5050
if !assert.NoError(t, err) {
5151
return
5252
}
53-
ctx, nc := websocketNetConn(r.Context(), sws, websocket.MessageBinary)
53+
ctx, nc := WebsocketNetConn(r.Context(), sws, websocket.MessageBinary)
5454
err = svc.ServeConnV2(ctx, nc, tailnet.StreamID{
5555
Name: "client",
5656
ID: clientID,

‎enterprise/coderd/provisionerdaemons.go

Lines changed: 2 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -7,13 +7,10 @@ import (
77
"errors"
88
"fmt"
99
"io"
10-
"net"
1110
"net/http"
1211
"strings"
1312
"time"
1413

15-
"github.com/coder/coder/v2/provisionersdk"
16-
1714
"github.com/google/uuid"
1815
"github.com/hashicorp/yamux"
1916
"github.com/moby/moby/pkg/namesgenerator"
@@ -37,6 +34,7 @@ import (
3734
"github.com/coder/coder/v2/coderd/util/ptr"
3835
"github.com/coder/coder/v2/codersdk"
3936
"github.com/coder/coder/v2/provisionerd/proto"
37+
"github.com/coder/coder/v2/provisionersdk"
4038
)
4139

4240
func (api *API) provisionerDaemonsEnabledMW(next http.Handler) http.Handler {
@@ -297,7 +295,7 @@ func (api *API) provisionerDaemonServe(rw http.ResponseWriter, r *http.Request)
297295
// the same connection.
298296
config := yamux.DefaultConfig()
299297
config.LogOutput = io.Discard
300-
ctx, wsNetConn := websocketNetConn(ctx, conn, websocket.MessageBinary)
298+
ctx, wsNetConn := codersdk.WebsocketNetConn(ctx, conn, websocket.MessageBinary)
301299
defer wsNetConn.Close()
302300
session, err := yamux.Server(wsNetConn, config)
303301
if err != nil {
@@ -360,44 +358,3 @@ func (api *API) provisionerDaemonServe(rw http.ResponseWriter, r *http.Request)
360358
}
361359
_ = conn.Close(websocket.StatusGoingAway, "")
362360
}
363-
364-
// wsNetConn wraps net.Conn created by websocket.NetConn(). Cancel func
365-
// is called if a read or write error is encountered.
366-
type wsNetConn struct {
367-
cancel context.CancelFunc
368-
net.Conn
369-
}
370-
371-
func (c *wsNetConn) Read(b []byte) (n int, err error) {
372-
n, err = c.Conn.Read(b)
373-
if err != nil {
374-
c.cancel()
375-
}
376-
return n, err
377-
}
378-
379-
func (c *wsNetConn) Write(b []byte) (n int, err error) {
380-
n, err = c.Conn.Write(b)
381-
if err != nil {
382-
c.cancel()
383-
}
384-
return n, err
385-
}
386-
387-
func (c *wsNetConn) Close() error {
388-
defer c.cancel()
389-
return c.Conn.Close()
390-
}
391-
392-
// websocketNetConn wraps websocket.NetConn and returns a context that
393-
// is tied to the parent context and the lifetime of the conn. Any error
394-
// during read or write will cancel the context, but not close the
395-
// conn. Close should be called to release context resources.
396-
func websocketNetConn(ctx context.Context, conn *websocket.Conn, msgType websocket.MessageType) (context.Context, net.Conn) {
397-
ctx, cancel := context.WithCancel(ctx)
398-
nc := websocket.NetConn(ctx, conn, msgType)
399-
return ctx, &wsNetConn{
400-
cancel: cancel,
401-
Conn: nc,
402-
}
403-
}

‎enterprise/coderd/workspaceproxycoordinate.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,7 @@ func (api *API) workspaceProxyCoordinate(rw http.ResponseWriter, r *http.Request
5757
return
5858
}
5959

60-
ctx, nc := websocketNetConn(ctx, conn, msgType)
60+
ctx, nc := codersdk.WebsocketNetConn(ctx, conn, msgType)
6161
defer nc.Close()
6262

6363
id := uuid.New()

0 commit comments

Comments
 (0)