Movatterモバイル変換

[0]ホーム
URL:

Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

 Sign up
Appearance settings

Commit89920d9

Browse files
committed
chore: consolidate websocketNetConn implementations
1 parent151aaad commit89920d9

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
gohttpapi.Heartbeat(ctx,conn)
546545

547-
ctx,wsNetConn:=websocketNetConn(ctx,conn,websocket.MessageText)
546+
ctx,wsNetConn:=codersdk.WebsocketNetConn(ctx,conn,websocket.MessageText)
548547
deferwsNetConn.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
defernconn.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
deferwsNetConn.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
deferwsNetConn.Close()
10821081

10831082
gohttpapi.Heartbeat(ctx,conn)
@@ -2104,47 +2103,6 @@ func createExternalAuthResponse(typ, token string, extra pqtype.NullRawMessage)
21042103
returnresp,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-
typewsNetConnstruct {
2110-
cancel context.CancelFunc
2111-
net.Conn
2112-
}
2113-
2114-
func (c*wsNetConn)Read(b []byte) (nint,errerror) {
2115-
n,err=c.Conn.Read(b)
2116-
iferr!=nil {
2117-
c.cancel()
2118-
}
2119-
returnn,err
2120-
}
2121-
2122-
func (c*wsNetConn)Write(b []byte) (nint,errerror) {
2123-
n,err=c.Conn.Write(b)
2124-
iferr!=nil {
2125-
c.cancel()
2126-
}
2127-
returnn,err
2128-
}
2129-
2130-
func (c*wsNetConn)Close()error {
2131-
deferc.cancel()
2132-
returnc.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-
funcwebsocketNetConn(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-
returnctx,&wsNetConn{
2143-
cancel:cancel,
2144-
Conn:nc,
2145-
}
2146-
}
2147-
21482106
funcconvertWorkspaceAgentLogs(logs []database.WorkspaceAgentLog) []codersdk.WorkspaceAgentLog {
21492107
sdk:=make([]codersdk.WorkspaceAgentLog,0,len(logs))
21502108
for_,logEntry:=rangelogs {

‎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
deferwsNetConn.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
returnnil,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
returnauthResp,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-
typewsNetConnstruct {
602-
cancel context.CancelFunc
603-
net.Conn
604-
}
605-
606-
func (c*wsNetConn)Read(b []byte) (nint,errerror) {
607-
n,err=c.Conn.Read(b)
608-
iferr!=nil {
609-
c.cancel()
610-
}
611-
returnn,err
612-
}
613-
614-
func (c*wsNetConn)Write(b []byte) (nint,errerror) {
615-
n,err=c.Conn.Write(b)
616-
iferr!=nil {
617-
c.cancel()
618-
}
619-
returnn,err
620-
}
621-
622-
func (c*wsNetConn)Close()error {
623-
deferc.cancel()
624-
returnc.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-
funcwebsocketNetConn(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-
returnctx,&wsNetConn{
638-
cancel:cancel,
639-
Conn:nc,
640-
}
641-
}
642-
643599
// LogsNotifyChannel returns the channel name responsible for notifying
644600
// of new logs.
645601
funcLogsNotifyChannel(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
iferr!=nil {
254254
_=conn.Close(websocket.StatusGoingAway,"")
@@ -287,11 +287,14 @@ func (c *wsNetConn) Close() error {
287287
returnc.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-
funcwebsocketNetConn(ctx context.Context,conn*websocket.Conn,msgType websocket.MessageType) (context.Context, net.Conn) {
294+
funcWebsocketNetConn(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
returnctx,&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(chanstruct{})
847-
ctx,wsNetConn:=websocketNetConn(ctx,conn,websocket.MessageText)
847+
ctx,wsNetConn:=WebsocketNetConn(ctx,conn,websocket.MessageText)
848848
decoder:=json.NewDecoder(wsNetConn)
849849
gofunc() {
850850
deferclose(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
deferwsNetConn.Close()
302300
session,err:=yamux.Server(wsNetConn,config)
303301
iferr!=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-
typewsNetConnstruct {
367-
cancel context.CancelFunc
368-
net.Conn
369-
}
370-
371-
func (c*wsNetConn)Read(b []byte) (nint,errerror) {
372-
n,err=c.Conn.Read(b)
373-
iferr!=nil {
374-
c.cancel()
375-
}
376-
returnn,err
377-
}
378-
379-
func (c*wsNetConn)Write(b []byte) (nint,errerror) {
380-
n,err=c.Conn.Write(b)
381-
iferr!=nil {
382-
c.cancel()
383-
}
384-
returnn,err
385-
}
386-
387-
func (c*wsNetConn)Close()error {
388-
deferc.cancel()
389-
returnc.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-
funcwebsocketNetConn(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-
returnctx,&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
defernc.Close()
6262

6363
id:=uuid.New()

0 commit comments

Comments
 (0)
[8]ページ先頭
©2009-2025 Movatter.jp