cli/vscodessh.go
Rama main · réplica de lectura del repositorio original.
package cli
import (
"context"
"fmt"
"io"
"net/http"
"net/url"
"os"
"path/filepath"
"strings"
"time"
"github.com/spf13/afero"
"golang.org/x/xerrors"
"cdr.dev/slog/v3"
"cdr.dev/slog/v3/sloggers/sloghuman"
"github.com/coder/coder/v2/cli/cliui"
"github.com/coder/coder/v2/cli/cliutil"
"github.com/coder/coder/v2/codersdk"
"github.com/coder/coder/v2/codersdk/workspacesdk"
"github.com/coder/serpent"
)
// vscodeSSH is used by the Coder VS Code extension to establish
// a connection to a workspace.
//
// This command needs to remain stable for compatibility with
// various VS Code versions, so it's kept separate from our
// standard SSH command.
func (r *RootCmd) vscodeSSH() *serpent.Command {
var (
sessionTokenFile string
urlFile string
logDir string
networkInfoDir string
networkInfoInterval time.Duration
waitEnum string
)
cmd := &serpent.Command{
// A SSH config entry is added by the VS Code extension that
// passes %h to ProxyCommand. The prefix of `coder-vscode--`
// is a magical string represented in our VS Code extension.
// It's not important here, only the delimiter `--` is.
Use: "vscodessh <coder-vscode--<owner>--<workspace>--<agent?>>",
Hidden: true,
Middleware: serpent.RequireNArgs(1),
Handler: func(inv *serpent.Invocation) error {
if networkInfoDir == "" {
return xerrors.New("network-info-dir must be specified")
}
if sessionTokenFile == "" {
return xerrors.New("session-token-file must be specified")
}
if urlFile == "" {
return xerrors.New("url-file must be specified")
}
fs, ok := inv.Context().Value("fs").(afero.Fs)
if !ok {
fs = afero.NewOsFs()
}
sessionToken, err := afero.ReadFile(fs, sessionTokenFile)
if err != nil {
return xerrors.Errorf("read session token: %w", err)
}
rawURL, err := afero.ReadFile(fs, urlFile)
if err != nil {
return xerrors.Errorf("read url: %w", err)
}
serverURL, err := url.Parse(string(rawURL))
if err != nil {
return xerrors.Errorf("parse url: %w", err)
}
ctx, cancel := context.WithCancel(inv.Context())
defer cancel()
// Configure HTTP client with transport wrappers
httpClient, err := r.createHTTPClient(ctx, serverURL, inv)
if err != nil {
return xerrors.Errorf("create HTTP client: %w", err)
}
client := codersdk.New(serverURL,
codersdk.WithSessionToken(string(sessionToken)),
codersdk.WithHTTPClient(httpClient),
)
parts := strings.Split(inv.Args[0], "--")
if len(parts) < 3 {
return xerrors.Errorf("invalid argument format. must be: coder-vscode--<owner>--<name>--<agent?>")
}
owner := parts[1]
name := parts[2]
if len(parts) > 3 {
name += "." + parts[3]
}
// Set autostart to false because it's assumed the VS Code extension
// will call this command after the workspace is started.
autostart := false
workspace, workspaceAgent, _, err := GetWorkspaceAndAgent(ctx, inv, client, autostart, fmt.Sprintf("%s/%s", owner, name))
if err != nil {
return xerrors.Errorf("find workspace and agent: %w", err)
}
// Select the startup script behavior based on template configuration or flags.
var wait bool
switch waitEnum {
case "yes":
wait = true
case "no":
wait = false
case "auto":
for _, script := range workspaceAgent.Scripts {
if script.StartBlocksLogin {
wait = true
break
}
}
default:
return xerrors.Errorf("unknown wait value %q", waitEnum)
}
appearanceCfg, err := client.Appearance(ctx)
if err != nil {
var sdkErr *codersdk.Error
if !xerrors.As(err, &sdkErr) || sdkErr.StatusCode() != http.StatusNotFound {
return xerrors.Errorf("get appearance config: %w", err)
}
appearanceCfg.DocsURL = codersdk.DefaultDocsURL()
}
err = cliui.Agent(ctx, inv.Stderr, workspaceAgent.ID, cliui.AgentOptions{
Fetch: client.WorkspaceAgent,
FetchLogs: client.WorkspaceAgentLogsAfter,
Wait: wait,
DocsURL: appearanceCfg.DocsURL,
})
if err != nil {
if xerrors.Is(err, context.Canceled) {
return cliui.ErrCanceled
}
}
// Use a stripped down writer that doesn't sync, otherwise you get
// "failed to sync sloghuman: sync /dev/stderr: The handle is
// invalid" on Windows. Syncing isn't required for stdout/stderr
// anyways.
logger := inv.Logger.AppendSinks(sloghuman.Sink(slogWriter{w: inv.Stderr})).Leveled(slog.LevelDebug)
if logDir != "" {
logFilePath := filepath.Join(logDir, fmt.Sprintf("%d.log", os.Getppid()))
logFile, err := fs.OpenFile(logFilePath, os.O_CREATE|os.O_WRONLY, 0o600)
if err != nil {
return xerrors.Errorf("open log file %q: %w", logFilePath, err)
}
dc := cliutil.DiscardAfterClose(logFile)
defer dc.Close()
logger = logger.AppendSinks(sloghuman.Sink(dc))
}
if r.disableDirect {
logger.Info(ctx, "direct connections disabled")
}
agentConn, err := workspacesdk.New(client).
DialAgent(ctx, workspaceAgent.ID, &workspacesdk.DialAgentOptions{
Logger: logger,
BlockEndpoints: r.disableDirect,
})
if err != nil {
return xerrors.Errorf("dial workspace agent: %w", err)
}
defer agentConn.Close()
agentConn.AwaitReachable(ctx)
closeUsage := client.UpdateWorkspaceUsageWithBodyContext(ctx, workspace.ID, codersdk.PostWorkspaceUsageRequest{
AgentID: workspaceAgent.ID,
AppName: string(codersdk.UsageAppNameVscode),
})
defer closeUsage()
rawSSH, err := agentConn.SSHTCPConn(ctx)
if err != nil {
return err
}
defer rawSSH.Close()
// Copy SSH traffic over stdio.
go func() {
_, _ = io.Copy(inv.Stdout, rawSSH)
}()
go func() {
_, _ = io.Copy(rawSSH, inv.Stdin)
}()
errCh, err := setStatsCallback(ctx, agentConn, logger, networkInfoDir, networkInfoInterval)
if err != nil {
return err
}
select {
case <-ctx.Done():
return nil
case err := <-errCh:
return err
}
},
}
cmd.Options = serpent.OptionSet{
{
Flag: "network-info-dir",
Description: "Specifies a directory to write network information periodically.",
Value: serpent.StringOf(&networkInfoDir),
},
{
Flag: "log-dir",
Description: "Specifies a directory to write logs to.",
Value: serpent.StringOf(&logDir),
},
{
Flag: "session-token-file",
Description: "Specifies a file that contains a session token.",
Value: serpent.StringOf(&sessionTokenFile),
},
{
Flag: "url-file",
Description: "Specifies a file that contains the Coder URL.",
Value: serpent.StringOf(&urlFile),
},
{
Flag: "network-info-interval",
Description: "Specifies the interval to update network information.",
Default: "5s",
Value: serpent.DurationOf(&networkInfoInterval),
},
{
Flag: "wait",
Description: "Specifies whether or not to wait for the startup script to finish executing. Auto means that the agent startup script behavior configured in the workspace template is used.",
Default: "auto",
Value: serpent.EnumOf(&waitEnum, "yes", "no", "auto"),
},
}
return cmd
}
// slogWriter wraps an io.Writer and removes all other methods (such as Sync),
// which may cause undesired/broken behavior.
type slogWriter struct {
w io.Writer
}
var _ io.Writer = slogWriter{}
func (s slogWriter) Write(p []byte) (n int, err error) {
return s.w.Write(p)
}