@@ -387,12 +387,12 @@ func (api *API) workspaceAgentPTY(rw http.ResponseWriter, r *http.Request) {
387
387
})
388
388
return
389
389
}
390
- defer func () {
391
- _ = conn .Close (websocket .StatusNormalClosure , "ended" )
392
- }()
390
+
393
391
// 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 )
396
396
if err != nil {
397
397
_ = conn .Close (websocket .StatusInternalError , httpapi .WebsocketCloseSprintf ("dial workspace agent: %s" , err ))
398
398
return
@@ -411,11 +411,13 @@ func (api *API) workspaceAgentPTY(rw http.ResponseWriter, r *http.Request) {
411
411
_ , _ = io .Copy (ptNetConn , wsNetConn )
412
412
}
413
413
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 ) {
416
418
client , server := provisionersdk .TransportPipe ()
417
419
go func () {
418
- _ = peerbroker .ProxyListen (r . Context () , server , peerbroker.ProxyOptions {
420
+ _ = peerbroker .ProxyListen (ctx , server , peerbroker.ProxyOptions {
419
421
ChannelID : agentID .String (),
420
422
Logger : api .Logger .Named ("peerbroker-proxy-dial" ),
421
423
Pubsub : api .Pubsub ,
@@ -425,7 +427,7 @@ func (api *API) dialWorkspaceAgent(r *http.Request, agentID uuid.UUID) (*agent.C
425
427
}()
426
428
427
429
peerClient := proto .NewDRPCPeerBrokerClient (provisionersdk .Conn (client ))
428
- stream , err := peerClient .NegotiateConnection (r . Context () )
430
+ stream , err := peerClient .NegotiateConnection (ctx )
429
431
if err != nil {
430
432
return nil , xerrors .Errorf ("negotiate: %w" , err )
431
433
}
@@ -437,7 +439,7 @@ func (api *API) dialWorkspaceAgent(r *http.Request, agentID uuid.UUID) (*agent.C
437
439
options .SettingEngine .SetICEProxyDialer (turnconn .ProxyDialer (func () (c net.Conn , err error ) {
438
440
clientPipe , serverPipe := net .Pipe ()
439
441
go func () {
440
- <- r . Context () .Done ()
442
+ <- ctx .Done ()
441
443
_ = clientPipe .Close ()
442
444
_ = serverPipe .Close ()
443
445
}()
0 commit comments