package agentsdk import ( "context" "io" "net/http" "net/http/cookiejar" "github.com/hashicorp/yamux" "golang.org/x/xerrors" "nhooyr.io/websocket" agentproto "github.com/coder/coder/v2/agent/proto" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/codersdk/drpc" ) func (c *Client) RPC(ctx context.Context) (agentproto.DRPCAgentClient, error) { rpcURL, err := c.SDK.URL.Parse("/api/v2/workspaceagents/me/rpc") if err != nil { return nil, xerrors.Errorf("parse url: %w", err) } jar, err := cookiejar.New(nil) if err != nil { return nil, xerrors.Errorf("create cookie jar: %w", err) } jar.SetCookies(rpcURL, []*http.Cookie{{ Name: codersdk.SessionTokenCookie, Value: c.SDK.SessionToken(), }}) httpClient := &http.Client{ Jar: jar, Transport: c.SDK.HTTPClient.Transport, } // nolint:bodyclose conn, res, err := websocket.Dial(ctx, rpcURL.String(), &websocket.DialOptions{ HTTPClient: httpClient, }) if err != nil { if res == nil { return nil, err } return nil, codersdk.ReadBodyAsError(res) } ctx, cancelFunc := context.WithCancel(ctx) ctx, wsNetConn := websocketNetConn(ctx, conn, websocket.MessageBinary) pingClosed := pingWebSocket(ctx, c.SDK.Logger(), conn, "RPC") nconn := &closeNetConn{ Conn: wsNetConn, closeFunc: func() { cancelFunc() _ = conn.Close(websocket.StatusGoingAway, "Listen closed") <-pingClosed }, } config := yamux.DefaultConfig() config.LogOutput = io.Discard mux, err := yamux.Client(nconn, config) if err != nil { return nil, xerrors.Errorf("create yamux client: %w", err) } dconn := drpc.MultiplexedConn(mux) return agentproto.NewDRPCAgentClient(dconn), nil }