Skip to content

Commit e90a0c0

Browse files
committed
fix: Use of r.Context() in workspaceAgentPTY
1 parent e0fb30c commit e90a0c0

File tree

1 file changed

+12
-10
lines changed

1 file changed

+12
-10
lines changed

coderd/workspaceagents.go

Lines changed: 12 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -387,12 +387,12 @@ func (api *API) workspaceAgentPTY(rw http.ResponseWriter, r *http.Request) {
387387
})
388388
return
389389
}
390-
defer func() {
391-
_ = conn.Close(websocket.StatusNormalClosure, "ended")
392-
}()
390+
393391
// Accept text connections, because it's more developer friendly.
394-
wsNetConn := websocket.NetConn(r.Context(), conn, websocket.MessageBinary)
395-
agentConn, err := api.dialWorkspaceAgent(r, workspaceAgent.ID)
392+
ctx, wsNetConn := websocketNetConn(r.Context(), conn, websocket.MessageBinary)
393+
defer wsNetConn.Close() // Also closes conn.
394+
395+
agentConn, err := api.dialWorkspaceAgent(ctx, r, workspaceAgent.ID)
396396
if err != nil {
397397
_ = conn.Close(websocket.StatusInternalError, httpapi.WebsocketCloseSprintf("dial workspace agent: %s", err))
398398
return
@@ -411,11 +411,13 @@ func (api *API) workspaceAgentPTY(rw http.ResponseWriter, r *http.Request) {
411411
_, _ = io.Copy(ptNetConn, wsNetConn)
412412
}
413413

414-
// dialWorkspaceAgent connects to a workspace agent by ID.
415-
func (api *API) dialWorkspaceAgent(r *http.Request, agentID uuid.UUID) (*agent.Conn, error) {
414+
// dialWorkspaceAgent connects to a workspace agent by ID. Only rely on
415+
// r.Context() for cancellation if it's use is safe or r.Hijack() has
416+
// not been performed.
417+
func (api *API) dialWorkspaceAgent(ctx context.Context, r *http.Request, agentID uuid.UUID) (*agent.Conn, error) {
416418
client, server := provisionersdk.TransportPipe()
417419
go func() {
418-
_ = peerbroker.ProxyListen(r.Context(), server, peerbroker.ProxyOptions{
420+
_ = peerbroker.ProxyListen(ctx, server, peerbroker.ProxyOptions{
419421
ChannelID: agentID.String(),
420422
Logger: api.Logger.Named("peerbroker-proxy-dial"),
421423
Pubsub: api.Pubsub,
@@ -425,7 +427,7 @@ func (api *API) dialWorkspaceAgent(r *http.Request, agentID uuid.UUID) (*agent.C
425427
}()
426428

427429
peerClient := proto.NewDRPCPeerBrokerClient(provisionersdk.Conn(client))
428-
stream, err := peerClient.NegotiateConnection(r.Context())
430+
stream, err := peerClient.NegotiateConnection(ctx)
429431
if err != nil {
430432
return nil, xerrors.Errorf("negotiate: %w", err)
431433
}
@@ -437,7 +439,7 @@ func (api *API) dialWorkspaceAgent(r *http.Request, agentID uuid.UUID) (*agent.C
437439
options.SettingEngine.SetICEProxyDialer(turnconn.ProxyDialer(func() (c net.Conn, err error) {
438440
clientPipe, serverPipe := net.Pipe()
439441
go func() {
440-
<-r.Context().Done()
442+
<-ctx.Done()
441443
_ = clientPipe.Close()
442444
_ = serverPipe.Close()
443445
}()

0 commit comments

Comments
 (0)