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  }