github.com/ungtb10d/cli/v2@v2.0.0-20221110210412-98537dd9d6a1/internal/codespaces/grpc/client.go (about) 1 package grpc 2 3 // gRPC client implementation to be able to connect to the gRPC server and perform the following operations: 4 // - Start a remote JupyterLab server 5 6 import ( 7 "context" 8 "fmt" 9 "net" 10 "strconv" 11 "time" 12 13 "github.com/ungtb10d/cli/v2/internal/codespaces/grpc/jupyter" 14 "github.com/ungtb10d/cli/v2/pkg/liveshare" 15 "golang.org/x/crypto/ssh" 16 "google.golang.org/grpc" 17 "google.golang.org/grpc/credentials/insecure" 18 "google.golang.org/grpc/metadata" 19 ) 20 21 const ( 22 ConnectionTimeout = 5 * time.Second 23 RequestTimeout = 30 * time.Second 24 ) 25 26 const ( 27 codespacesInternalPort = 16634 28 codespacesInternalSessionName = "CodespacesInternal" 29 ) 30 31 type Client struct { 32 conn *grpc.ClientConn 33 token string 34 listener net.Listener 35 jupyterClient jupyter.JupyterServerHostClient 36 cancelPF context.CancelFunc 37 } 38 39 type liveshareSession interface { 40 KeepAlive(string) 41 OpenStreamingChannel(context.Context, liveshare.ChannelID) (ssh.Channel, error) 42 StartSharing(context.Context, string, int) (liveshare.ChannelID, error) 43 } 44 45 // Finds a free port to listen on and creates a new gRPC client that connects to that port 46 func Connect(ctx context.Context, session liveshareSession, token string) (*Client, error) { 47 listener, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", 0)) 48 if err != nil { 49 return nil, fmt.Errorf("failed to listen to local port over tcp: %w", err) 50 } 51 localAddress := fmt.Sprintf("127.0.0.1:%d", listener.Addr().(*net.TCPAddr).Port) 52 53 client := &Client{ 54 token: token, 55 listener: listener, 56 } 57 58 // Create a cancelable context to be able to cancel background tasks 59 // if we encounter an error while connecting to the gRPC server 60 connectctx, cancel := context.WithCancel(context.Background()) 61 defer func() { 62 if err != nil { 63 cancel() 64 } 65 }() 66 67 ch := make(chan error, 2) // Buffered channel to ensure we don't block on the goroutine 68 69 // Ensure we close the port forwarder if we encounter an error 70 // or once the gRPC connection is closed. pfcancel is retained 71 // to close the PF whenever we close the gRPC connection. 72 pfctx, pfcancel := context.WithCancel(connectctx) 73 client.cancelPF = pfcancel 74 75 // Tunnel the remote gRPC server port to the local port 76 go func() { 77 fwd := liveshare.NewPortForwarder(session, codespacesInternalSessionName, codespacesInternalPort, true) 78 ch <- fwd.ForwardToListener(pfctx, listener) 79 }() 80 81 var conn *grpc.ClientConn 82 go func() { 83 // Attempt to connect to the port 84 opts := []grpc.DialOption{ 85 grpc.WithTransportCredentials(insecure.NewCredentials()), 86 grpc.WithBlock(), 87 } 88 conn, err = grpc.DialContext(connectctx, localAddress, opts...) 89 ch <- err // nil if we successfully connected 90 }() 91 92 // Wait for the connection to be established or for the context to be cancelled 93 select { 94 case <-ctx.Done(): 95 return nil, ctx.Err() 96 case err := <-ch: 97 if err != nil { 98 return nil, err 99 } 100 } 101 102 client.conn = conn 103 client.jupyterClient = jupyter.NewJupyterServerHostClient(conn) 104 105 return client, nil 106 } 107 108 // Closes the gRPC connection 109 func (g *Client) Close() error { 110 g.cancelPF() 111 112 // Closing the local listener effectively closes the gRPC connection 113 if err := g.listener.Close(); err != nil { 114 g.conn.Close() // If we fail to close the listener, explicitly close the gRPC connection and ignore any error 115 return fmt.Errorf("failed to close local tcp port listener: %w", err) 116 } 117 118 return nil 119 } 120 121 // Appends the authentication token to the gRPC context 122 func (g *Client) appendMetadata(ctx context.Context) context.Context { 123 return metadata.AppendToOutgoingContext(ctx, "Authorization", "Bearer "+g.token) 124 } 125 126 // Starts a remote JupyterLab server to allow the user to connect to the codespace via JupyterLab in their browser 127 func (g *Client) StartJupyterServer(ctx context.Context) (port int, serverUrl string, err error) { 128 ctx = g.appendMetadata(ctx) 129 130 response, err := g.jupyterClient.GetRunningServer(ctx, &jupyter.GetRunningServerRequest{}) 131 if err != nil { 132 return 0, "", fmt.Errorf("failed to invoke JupyterLab RPC: %w", err) 133 } 134 135 if !response.Result { 136 return 0, "", fmt.Errorf("failed to start JupyterLab: %s", response.Message) 137 } 138 139 port, err = strconv.Atoi(response.Port) 140 if err != nil { 141 return 0, "", fmt.Errorf("failed to parse JupyterLab port: %w", err) 142 } 143 144 return port, response.ServerUrl, err 145 }