Add --user flag to more fly ssh commands

similar concept to --user flag that was added to fly ssh issue: #1572
superfly · Feb 8, 2023 · 51ea3c5 · 51ea3c5
1 parent 8c59e3d
commit 51ea3c5
Show file tree

Hide file tree

Showing 6 changed files with 47 additions and 31 deletions.
diff --git a/flypg/cmd.go b/flypg/cmd.go
@@ -61,7 +61,7 @@ func (pc *Command) UpdateSettings(ctx context.Context, leaderIp string, config m
 	subCmd := fmt.Sprintf("update --patch '%s'", string(configBytes))
 	cmd := fmt.Sprintf("stolonctl-run %s", encodeCommand(subCmd))
 
-	resp, err := ssh.RunSSHCommand(ctx, pc.app, pc.dialer, leaderIp, cmd)
+	resp, err := ssh.RunSSHCommand(ctx, pc.app, pc.dialer, leaderIp, cmd, "root")
 	if err != nil {
 		return err
 	}
@@ -82,7 +82,7 @@ func (pc *Command) UnregisterMember(ctx context.Context, leaderIP string, standb
 	payload := encodeCommand(standbyIP)
 	cmd := fmt.Sprintf("pg_unregister %s", payload)
 
-	resp, err := ssh.RunSSHCommand(ctx, pc.app, pc.dialer, leaderIP, cmd)
+	resp, err := ssh.RunSSHCommand(ctx, pc.app, pc.dialer, leaderIP, cmd, "root")
 	if err != nil {
 		return err
 	}

diff --git a/internal/command/postgres/connect.go b/internal/command/postgres/connect.go
@@ -19,6 +19,8 @@ import (
 	"github.com/superfly/flyctl/internal/flag"
 )
 
+const defaultSshUsername = "root"
+
 func newConnect() *cobra.Command {
 	const (
 		short = "Connect to the Postgres console"
@@ -114,14 +116,15 @@ func runMachineConnect(ctx context.Context, app *api.AppCompact) error {
 		return err
 	}
 	return ssh.SSHConnect(&ssh.SSHParams{
-		Ctx:    ctx,
-		Org:    app.Organization,
-		Dialer: agent.DialerFromContext(ctx),
-		App:    app.Name,
-		Cmd:    fmt.Sprintf("connect %s %s %s", database, user, password),
-		Stdin:  os.Stdin,
-		Stdout: ioutils.NewWriteCloserWrapper(colorable.NewColorableStdout(), func() error { return nil }),
-		Stderr: ioutils.NewWriteCloserWrapper(colorable.NewColorableStderr(), func() error { return nil }),
+		Ctx:      ctx,
+		Org:      app.Organization,
+		Dialer:   agent.DialerFromContext(ctx),
+		App:      app.Name,
+		Username: defaultSshUsername,
+		Cmd:      fmt.Sprintf("connect %s %s %s", database, user, password),
+		Stdin:    os.Stdin,
+		Stdout:   ioutils.NewWriteCloserWrapper(colorable.NewColorableStdout(), func() error { return nil }),
+		Stderr:   ioutils.NewWriteCloserWrapper(colorable.NewColorableStderr(), func() error { return nil }),
 	}, leader.PrivateIP)
 }
 
@@ -159,13 +162,14 @@ func runNomadConnect(ctx context.Context, app *api.AppCompact) error {
 	}
 
 	return ssh.SSHConnect(&ssh.SSHParams{
-		Ctx:    ctx,
-		Org:    app.Organization,
-		Dialer: agent.DialerFromContext(ctx),
-		App:    app.Name,
-		Cmd:    fmt.Sprintf("connect %s %s %s", database, user, password),
-		Stdin:  os.Stdin,
-		Stdout: ioutils.NewWriteCloserWrapper(colorable.NewColorableStdout(), func() error { return nil }),
-		Stderr: ioutils.NewWriteCloserWrapper(colorable.NewColorableStderr(), func() error { return nil }),
+		Ctx:      ctx,
+		Org:      app.Organization,
+		Dialer:   agent.DialerFromContext(ctx),
+		App:      app.Name,
+		Username: defaultSshUsername,
+		Cmd:      fmt.Sprintf("connect %s %s %s", database, user, password),
+		Stdin:    os.Stdin,
+		Stdout:   ioutils.NewWriteCloserWrapper(colorable.NewColorableStdout(), func() error { return nil }),
+		Stderr:   ioutils.NewWriteCloserWrapper(colorable.NewColorableStderr(), func() error { return nil }),
 	}, leaderIP)
 }
diff --git a/internal/command/ssh/console.go b/internal/command/ssh/console.go
@@ -59,6 +59,12 @@ func stdArgsSSH(cmd *cobra.Command) {
 			Shorthand:   "A",
 			Description: "Address of VM to connect to",
 		},
+		flag.String{
+			Name:        "user",
+			Shorthand:   "u",
+			Description: "Unix username to connect as",
+			Default:     defaultSshUsername,
+		},
 	)
 }
 
@@ -177,14 +183,15 @@ func runConsole(ctx context.Context) error {
 
 	// BUG(tqbf): many of these are no longer really params
 	params := &SSHParams{
-		Ctx:    ctx,
-		Org:    app.Organization,
-		Dialer: dialer,
-		App:    appName,
-		Cmd:    flag.GetString(ctx, "command"),
-		Stdin:  os.Stdin,
-		Stdout: ioutils.NewWriteCloserWrapper(colorable.NewColorableStdout(), func() error { return nil }),
-		Stderr: ioutils.NewWriteCloserWrapper(colorable.NewColorableStderr(), func() error { return nil }),
+		Ctx:      ctx,
+		Org:      app.Organization,
+		Dialer:   dialer,
+		App:      appName,
+		Username: flag.GetString(ctx, "user"),
+		Cmd:      flag.GetString(ctx, "command"),
+		Stdin:    os.Stdin,
+		Stdout:   ioutils.NewWriteCloserWrapper(colorable.NewColorableStdout(), func() error { return nil }),
+		Stderr:   ioutils.NewWriteCloserWrapper(colorable.NewColorableStderr(), func() error { return nil }),
 	}
 
 	if quiet(ctx) {
@@ -226,7 +233,7 @@ func sshConnect(p *SSHParams, addr string) (*ssh.Client, error) {
 
 	sshClient := &ssh.Client{
 		Addr: net.JoinHostPort(addr, "22"),
-		User: "root",
+		User: p.Username,
 
 		Dial: p.Dialer.DialContext,
 

diff --git a/internal/command/ssh/issue.go b/internal/command/ssh/issue.go
@@ -46,7 +46,7 @@ validity.`
 			Name:        "username",
 			Shorthand:   "u",
 			Description: "Unix usernames the SSH cert can authenticate as",
-			Default:     []string{"root", "fly"},
+			Default:     []string{defaultSshUsername, "fly"},
 		},
 		flag.Int{
 			Name:        "hours",

diff --git a/internal/command/ssh/sftp.go b/internal/command/ssh/sftp.go
@@ -110,6 +110,7 @@ func newSFTPConnection(ctx context.Context) (*sftp.Client, error) {
 		Org:            app.Organization,
 		Dialer:         dialer,
 		App:            appName,
+		Username:       defaultSshUsername,
 		Stdin:          os.Stdin,
 		Stdout:         os.Stdout,
 		Stderr:         os.Stderr,

diff --git a/internal/command/ssh/ssh_terminal.go b/internal/command/ssh/ssh_terminal.go
@@ -46,10 +46,13 @@ func spin(in, out string) context.CancelFunc {
 	return cancel
 }
 
+const defaultSshUsername = "root"
+
 type SSHParams struct {
 	Ctx            context.Context
 	Org            api.OrganizationImpl
 	App            string
+	Username       string
 	Dialer         agent.Dialer
 	Cmd            string
 	Stdin          io.Reader
@@ -58,7 +61,7 @@ type SSHParams struct {
 	DisableSpinner bool
 }
 
-func RunSSHCommand(ctx context.Context, app *api.AppCompact, dialer agent.Dialer, addr string, cmd string) ([]byte, error) {
+func RunSSHCommand(ctx context.Context, app *api.AppCompact, dialer agent.Dialer, addr string, cmd string, username string) ([]byte, error) {
 	var inBuf bytes.Buffer
 	var errBuf bytes.Buffer
 	var outBuf bytes.Buffer
@@ -71,6 +74,7 @@ func RunSSHCommand(ctx context.Context, app *api.AppCompact, dialer agent.Dialer
 		Org:            app.Organization,
 		Dialer:         dialer,
 		App:            app.Name,
+		Username:       username,
 		Cmd:            cmd,
 		Stdin:          inReader,
 		Stdout:         stdoutWriter,
@@ -102,7 +106,7 @@ func SSHConnect(p *SSHParams, addr string) error {
 
 	sshClient := &ssh.Client{
 		Addr: net.JoinHostPort(addr, "22"),
-		User: "root",
+		User: p.Username,
 
 		Dial: p.Dialer.DialContext,
 
@@ -151,7 +155,7 @@ func singleUseSSHCertificate(ctx context.Context, org api.OrganizationImpl) (*ap
 		return nil, nil, err
 	}
 
-	icert, err := client.IssueSSHCertificate(ctx, org, []string{"root", "fly"}, nil, &hours, pub)
+	icert, err := client.IssueSSHCertificate(ctx, org, []string{defaultSshUsername, "fly"}, nil, &hours, pub)
 	if err != nil {
 		return nil, nil, err
 	}