@@ -387,12 +387,12 @@ func (api *API) workspaceAgentPTY(rw http.ResponseWriter, r *http.Request) {
387387})
388388return
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 )
396396if err != nil {
397397_ = conn .Close (websocket .StatusInternalError ,httpapi .WebsocketCloseSprintf ("dial workspace agent: %s" ,err ))
398398return
@@ -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 ) {
416418client ,server := provisionersdk .TransportPipe ()
417419go func () {
418- _ = peerbroker .ProxyListen (r . Context () ,server , peerbroker.ProxyOptions {
420+ _ = peerbroker .ProxyListen (ctx ,server , peerbroker.ProxyOptions {
419421ChannelID :agentID .String (),
420422Logger :api .Logger .Named ("peerbroker-proxy-dial" ),
421423Pubsub :api .Pubsub ,
@@ -425,7 +427,7 @@ func (api *API) dialWorkspaceAgent(r *http.Request, agentID uuid.UUID) (*agent.C
425427}()
426428
427429peerClient := proto .NewDRPCPeerBrokerClient (provisionersdk .Conn (client ))
428- stream ,err := peerClient .NegotiateConnection (r . Context () )
430+ stream ,err := peerClient .NegotiateConnection (ctx )
429431if err != nil {
430432return nil ,xerrors .Errorf ("negotiate: %w" ,err )
431433}
@@ -437,7 +439,7 @@ func (api *API) dialWorkspaceAgent(r *http.Request, agentID uuid.UUID) (*agent.C
437439options .SettingEngine .SetICEProxyDialer (turnconn .ProxyDialer (func () (c net.Conn ,err error ) {
438440clientPipe ,serverPipe := net .Pipe ()
439441go func () {
440- <- r . Context () .Done ()
442+ <- ctx .Done ()
441443_ = clientPipe .Close ()
442444_ = serverPipe .Close ()
443445}()