@@ -376,12 +376,12 @@ func (api *API) workspaceAgentPTY(rw http.ResponseWriter, r *http.Request) {
376376})
377377return
378378}
379- defer func () {
380- _ = conn .Close (websocket .StatusNormalClosure ,"ended" )
381- }()
379+
382380// Accept text connections, because it's more developer friendly.
383- wsNetConn := websocket .NetConn (r .Context (),conn ,websocket .MessageBinary )
384- agentConn ,err := api .dialWorkspaceAgent (r ,workspaceAgent .ID )
381+ ctx ,wsNetConn := websocketNetConn (r .Context (),conn ,websocket .MessageBinary )
382+ defer wsNetConn .Close ()// Also closes conn.
383+
384+ agentConn ,err := api .dialWorkspaceAgent (ctx ,r ,workspaceAgent .ID )
385385if err != nil {
386386_ = conn .Close (websocket .StatusInternalError ,httpapi .WebsocketCloseSprintf ("dial workspace agent: %s" ,err ))
387387return
@@ -400,11 +400,13 @@ func (api *API) workspaceAgentPTY(rw http.ResponseWriter, r *http.Request) {
400400_ ,_ = io .Copy (ptNetConn ,wsNetConn )
401401}
402402
403- // dialWorkspaceAgent connects to a workspace agent by ID.
404- func (api * API )dialWorkspaceAgent (r * http.Request ,agentID uuid.UUID ) (* agent.Conn ,error ) {
403+ // dialWorkspaceAgent connects to a workspace agent by ID. Only rely on
404+ // r.Context() for cancellation if it's use is safe or r.Hijack() has
405+ // not been performed.
406+ func (api * API )dialWorkspaceAgent (ctx context.Context ,r * http.Request ,agentID uuid.UUID ) (* agent.Conn ,error ) {
405407client ,server := provisionersdk .TransportPipe ()
406408go func () {
407- _ = peerbroker .ProxyListen (r . Context () ,server , peerbroker.ProxyOptions {
409+ _ = peerbroker .ProxyListen (ctx ,server , peerbroker.ProxyOptions {
408410ChannelID :agentID .String (),
409411Logger :api .Logger .Named ("peerbroker-proxy-dial" ),
410412Pubsub :api .Pubsub ,
@@ -414,7 +416,7 @@ func (api *API) dialWorkspaceAgent(r *http.Request, agentID uuid.UUID) (*agent.C
414416}()
415417
416418peerClient := proto .NewDRPCPeerBrokerClient (provisionersdk .Conn (client ))
417- stream ,err := peerClient .NegotiateConnection (r . Context () )
419+ stream ,err := peerClient .NegotiateConnection (ctx )
418420if err != nil {
419421return nil ,xerrors .Errorf ("negotiate: %w" ,err )
420422}
@@ -426,7 +428,7 @@ func (api *API) dialWorkspaceAgent(r *http.Request, agentID uuid.UUID) (*agent.C
426428options .SettingEngine .SetICEProxyDialer (turnconn .ProxyDialer (func () (c net.Conn ,err error ) {
427429clientPipe ,serverPipe := net .Pipe ()
428430go func () {
429- <- r . Context () .Done ()
431+ <- ctx .Done ()
430432_ = clientPipe .Close ()
431433_ = serverPipe .Close ()
432434}()