From 58a055609dea29874e5a4e1ba00a56897a1599a2 Mon Sep 17 00:00:00 2001 From: Jose Garcia Date: Thu, 29 Jul 2021 10:57:51 -0400 Subject: [PATCH] logs cmd spike and refactor of ssh tunnel methods --- cmd/ghcs/logs.go | 95 ++++++++++++++++++++++++ cmd/ghcs/ssh.go | 87 +++++----------------- internal/codespaces/codespaces.go | 32 +++++++++ internal/codespaces/ssh.go | 116 ++++++++++++++++++++++++++++++ 4 files changed, 260 insertions(+), 70 deletions(-) create mode 100644 cmd/ghcs/logs.go create mode 100644 internal/codespaces/ssh.go diff --git a/cmd/ghcs/logs.go b/cmd/ghcs/logs.go new file mode 100644 index 000000000..3696999d5 --- /dev/null +++ b/cmd/ghcs/logs.go @@ -0,0 +1,95 @@ +package main + +import ( + "bufio" + "context" + "fmt" + "os" + + "github.com/github/ghcs/api" + "github.com/github/ghcs/internal/codespaces" + "github.com/spf13/cobra" +) + +func NewLogsCmd() *cobra.Command { + return &cobra.Command{ + Use: "logs", + Short: "Access Codespace logs", + RunE: func(cmd *cobra.Command, args []string) error { + var codespaceName string + if len(args) > 0 { + codespaceName = args[0] + } + return Logs(codespaceName) + }, + } +} + +func init() { + rootCmd.AddCommand(NewLogsCmd()) +} + +func Logs(codespaceName string) error { + apiClient := api.New(os.Getenv("GITHUB_TOKEN")) + ctx := context.Background() + + user, err := apiClient.GetUser(ctx) + if err != nil { + return fmt.Errorf("getting user: %v", err) + } + + codespace, token, err := codespaces.GetOrChooseCodespace(ctx, apiClient, user, codespaceName) + if err != nil { + return fmt.Errorf("get or choose codespace: %v", err) + } + + lsclient, err := codespaces.ConnectToLiveshare(ctx, apiClient, token, codespace) + if err != nil { + return fmt.Errorf("connecting to liveshare: %v", err) + } + + tunnelPort, connClosed, err := codespaces.MakeSSHTunnel(ctx, lsclient, 0) + if err != nil { + return fmt.Errorf("make ssh tunnel: %v", err) + } + + dst := fmt.Sprintf("%s@localhost", getSSHUser(codespace)) + stdout, err := codespaces.RunCommand( + ctx, tunnelPort, dst, "cat /workspaces/.codespaces/.persistedshare/creation.log", + ) + if err != nil { + return fmt.Errorf("run command: %v", err) + } + + done := make(chan error) + go func() { + scanner := bufio.NewScanner(stdout) + for scanner.Scan() { + fmt.Println(scanner.Text()) + } + + if err := scanner.Err(); err != nil { + done <- fmt.Errorf("error scanning: %v", err) + return + } + + if err := stdout.Close(); err != nil { + done <- fmt.Errorf("close stdout: %v", err) + return + } + done <- nil + }() + + select { + case err := <-connClosed: + if err != nil { + return fmt.Errorf("connection closed: %v", err) + } + case err := <-done: + if err != nil { + return err + } + } + + return nil +} diff --git a/cmd/ghcs/ssh.go b/cmd/ghcs/ssh.go index ef03ba946..23a4c2ca0 100644 --- a/cmd/ghcs/ssh.go +++ b/cmd/ghcs/ssh.go @@ -4,10 +4,7 @@ import ( "bufio" "context" "fmt" - "math/rand" "os" - "os/exec" - "strconv" "strings" "time" @@ -49,37 +46,9 @@ func SSH(sshProfile, codespaceName string, sshServerPort int) error { return fmt.Errorf("error getting user: %v", err) } - var ( - codespace *api.Codespace - token string - ) - - if codespaceName == "" { - codespace, err = codespaces.ChooseCodespace(ctx, apiClient, user) - if err != nil { - if err == codespaces.ErrNoCodespaces { - fmt.Println(err.Error()) - return nil - } - - return fmt.Errorf("error choosing codespace: %v", err) - } - codespaceName = codespace.Name - - token, err = apiClient.GetCodespaceToken(ctx, user.Login, codespaceName) - if err != nil { - return fmt.Errorf("error getting codespace token: %v", err) - } - } else { - token, err = apiClient.GetCodespaceToken(ctx, user.Login, codespaceName) - if err != nil { - return fmt.Errorf("error getting codespace token: %v", err) - } - - codespace, err = apiClient.GetCodespace(ctx, token, user.Login, codespaceName) - if err != nil { - return fmt.Errorf("error getting full codespace details: %v", err) - } + codespace, token, err := codespaces.GetOrChooseCodespace(ctx, apiClient, user, codespaceName) + if err != nil { + return fmt.Errorf("get or choose codespace: %v") } lsclient, err := codespaces.ConnectToLiveshare(ctx, apiClient, token, codespace) @@ -106,56 +75,34 @@ func SSH(sshProfile, codespaceName string, sshServerPort int) error { fmt.Printf("\n") } - server, err := liveshare.NewServer(lsclient) + tunnelPort, tunnelClosed, err := codespaces.MakeSSHTunnel(ctx, lsclient, sshServerPort) if err != nil { - return fmt.Errorf("error creating server: %v", err) + return fmt.Errorf("make ssh tunnel: %v", err) } - rand.Seed(time.Now().Unix()) - port := rand.Intn(9999-2000) + 2000 // improve this obviously - if sshServerPort != 0 { - port = sshServerPort - } - - if err := server.StartSharing(ctx, "sshd", 2222); err != nil { - return fmt.Errorf("error sharing sshd port: %v", err) - } - - portForwarder := liveshare.NewPortForwarder(lsclient, server, port) - go func() { - if err := portForwarder.Start(ctx); err != nil { - panic(fmt.Errorf("error forwarding port: %v", err)) - } - }() - connectDestination := sshProfile if connectDestination == "" { connectDestination = fmt.Sprintf("%s@localhost", getSSHUser(codespace)) } + usingCustomPort := tunnelPort == sshServerPort + connClosed := codespaces.ConnectToTunnel(ctx, tunnelPort, connectDestination, usingCustomPort) + fmt.Println("Ready...") - if err := connect(ctx, port, connectDestination, port == sshServerPort); err != nil { - return fmt.Errorf("error connecting via SSH: %v", err) + select { + case err := <-tunnelClosed: + if err != nil { + return fmt.Errorf("tunnel closed: %v", err) + } + case err := <-connClosed: + if err != nil { + return fmt.Errorf("connection closed: %v", err) + } } return nil } -func connect(ctx context.Context, port int, destination string, setServerPort bool) error { - connectionDetailArgs := []string{"-p", strconv.Itoa(port), "-o", "NoHostAuthenticationForLocalhost=yes"} - - if setServerPort { - fmt.Println("Connection Details: ssh " + destination + " " + strings.Join(connectionDetailArgs, " ")) - } - - args := []string{destination, "-X", "-Y", "-C"} // X11, X11Trust, Compression - cmd := exec.CommandContext(ctx, "ssh", append(args, connectionDetailArgs...)...) - cmd.Stdout = os.Stdout - cmd.Stdin = os.Stdin - cmd.Stderr = os.Stderr - return cmd.Run() -} - func getContainerID(ctx context.Context, terminal *liveshare.Terminal) (string, error) { fmt.Print(".") cmd := terminal.NewCommand( diff --git a/internal/codespaces/codespaces.go b/internal/codespaces/codespaces.go index 6c3517f39..4c62d9aff 100644 --- a/internal/codespaces/codespaces.go +++ b/internal/codespaces/codespaces.go @@ -111,3 +111,35 @@ func ConnectToLiveshare(ctx context.Context, apiClient *api.API, token string, c return lsclient, nil } + +func GetOrChooseCodespace(ctx context.Context, apiClient *api.API, user *api.User, codespaceName string) (codespace *api.Codespace, token string, err error) { + if codespaceName == "" { + codespace, err = ChooseCodespace(ctx, apiClient, user) + if err != nil { + if err == ErrNoCodespaces { + fmt.Println(err.Error()) + return nil, "", nil + } + + return nil, "", fmt.Errorf("choosing codespace: %v", err) + } + codespaceName = codespace.Name + + token, err = apiClient.GetCodespaceToken(ctx, user.Login, codespaceName) + if err != nil { + return nil, "", fmt.Errorf("getting codespace token: %v", err) + } + } else { + token, err = apiClient.GetCodespaceToken(ctx, user.Login, codespaceName) + if err != nil { + return nil, "", fmt.Errorf("getting codespace token for given codespace: %v", err) + } + + codespace, err = apiClient.GetCodespace(ctx, token, user.Login, codespaceName) + if err != nil { + return nil, "", fmt.Errorf("getting full codespace details: %v", err) + } + } + + return codespace, token, nil +} diff --git a/internal/codespaces/ssh.go b/internal/codespaces/ssh.go new file mode 100644 index 000000000..2bb661086 --- /dev/null +++ b/internal/codespaces/ssh.go @@ -0,0 +1,116 @@ +package codespaces + +import ( + "context" + "fmt" + "io" + "math/rand" + "os" + "os/exec" + "strconv" + "strings" + "time" + + "github.com/github/go-liveshare" +) + +func MakeSSHTunnel(ctx context.Context, lsclient *liveshare.Client, serverPort int) (int, <-chan error, error) { + tunnelClosed := make(chan error) + + server, err := liveshare.NewServer(lsclient) + if err != nil { + return 0, nil, fmt.Errorf("new liveshare server: %v", err) + } + + rand.Seed(time.Now().Unix()) + port := rand.Intn(9999-2000) + 2000 // improve this obviously + if serverPort != 0 { + port = serverPort + } + + // TODO(josebalius): This port won't always be 2222 + if err := server.StartSharing(ctx, "sshd", 2222); err != nil { + return 0, nil, fmt.Errorf("sharing sshd port: %v", err) + } + + go func() { + portForwarder := liveshare.NewPortForwarder(lsclient, server, port) + if err := portForwarder.Start(ctx); err != nil { + tunnelClosed <- fmt.Errorf("forwarding port: %v", err) + return + } + tunnelClosed <- nil + }() + + return port, tunnelClosed, nil +} + +func makeSSHArgs(port int, dst, cmd string) ([]string, []string) { + connArgs := []string{"-p", strconv.Itoa(port), "-o", "NoHostAuthenticationForLocalhost=yes"} + cmdArgs := append([]string{dst, "-X", "-Y", "-C"}, connArgs...) // X11, X11Trust, Compression + + if cmd != "" { + cmdArgs = append(cmdArgs, cmd) + } + + return cmdArgs, connArgs +} + +func ConnectToTunnel(ctx context.Context, port int, destination string, usingCustomPort bool) <-chan error { + connClosed := make(chan error) + args, connArgs := makeSSHArgs(port, destination, "") + + if usingCustomPort { + fmt.Println("Connection Details: ssh " + destination + " " + strings.Join(connArgs, " ")) + } + + cmd := exec.CommandContext(ctx, "ssh", args...) + cmd.Stdout = os.Stdout + cmd.Stdin = os.Stdin + cmd.Stderr = os.Stderr + + go func() { + connClosed <- cmd.Run() + }() + + return connClosed +} + +type command struct { + Cmd *exec.Cmd + StdoutPipe io.ReadCloser +} + +func newCommand(cmd *exec.Cmd) (*command, error) { + stdoutPipe, err := cmd.StdoutPipe() + if err != nil { + return nil, fmt.Errorf("create stdout pipe: %v", err) + } + + if err := cmd.Start(); err != nil { + return nil, fmt.Errorf("cmd start: %v", err) + } + + return &command{ + Cmd: cmd, + StdoutPipe: stdoutPipe, + }, nil +} + +func (c *command) Read(p []byte) (int, error) { + return c.StdoutPipe.Read(p) +} + +func (c *command) Close() error { + if err := c.StdoutPipe.Close(); err != nil { + return fmt.Errorf("close stdout: %v", err) + } + + return c.Cmd.Wait() +} + +func RunCommand(ctx context.Context, tunnelPort int, destination, cmdString string) (io.ReadCloser, error) { + args, _ := makeSSHArgs(tunnelPort, destination, cmdString) + cmd := exec.CommandContext(ctx, "ssh", args...) + return newCommand(cmd) +}