Merge branch '91-identity-file-support' into 'master'

# Description
- Support an identity file in CLI for ssh authentication.
- Use an identity key if specified, otherwise use an ssh-agent

See merge request postgres-ai/database-lab!121

Author: agneum

cmd/cli/commands/client.go

@@ -21,6 +21,7 @@ const (
 	InsecureKey      = "insecure"
 	FwServerURLKey   = "forwarding-server-url"
 	FwLocalPortKey   = "forwarding-local-port"
+	IdentityFileKey  = "identity-file"
 )
 
 // ClientByCLIContext creates a new Database Lab API client.

cmd/cli/commands/clone/actions.go

@@ -261,7 +261,7 @@ func forward(cliCtx *cli.Context) error {
 
 	wg := &sync.WaitGroup{}
 
-	port, err := retrieveClonePort(cliCtx, wg, remoteURL.Host)
+	port, err := retrieveClonePort(cliCtx, wg, remoteURL)
 	if err != nil {
 		return err
 	}
@@ -270,7 +270,9 @@ func forward(cliCtx *cli.Context) error {
 
 	log.Dbg(fmt.Sprintf("The clone port has been retrieved: %s", port))
 
-	tunnel, err := commands.BuildTunnel(cliCtx, commands.BuildHostname(remoteURL.Hostname(), port))
+	remoteURL.Host = commands.BuildHostname(remoteURL.Hostname(), port)
+
+	tunnel, err := commands.BuildTunnel(cliCtx, remoteURL)
 	if err != nil {
 		return err
 	}
@@ -288,7 +290,7 @@ func forward(cliCtx *cli.Context) error {
 	return nil
 }
 
-func retrieveClonePort(cliCtx *cli.Context, wg *sync.WaitGroup, remoteHost string) (string, error) {
+func retrieveClonePort(cliCtx *cli.Context, wg *sync.WaitGroup, remoteHost *url.URL) (string, error) {
 	tunnel, err := commands.BuildTunnel(cliCtx, remoteHost)
 	if err != nil {
 		return "", err

cmd/cli/commands/config/command_list.go

@@ -45,6 +45,10 @@ func CommandList() []*cli.Command {
 							Name:  "forwarding-local-port",
 							Usage: "local port for forwarding to the Database Lab instance",
 						},
+						&cli.StringFlag{
+							Name:  "identity-file",
+							Usage: "select a file from which the identity (private key) for public key authentication is read",
+						},
 					},
 				},
 				{
@@ -73,6 +77,10 @@ func CommandList() []*cli.Command {
 							Name:  "forwarding-local-port",
 							Usage: "local port for forwarding to the Database Lab instance",
 						},
+						&cli.StringFlag{
+							Name:  "identity-file",
+							Usage: "select a file from which the identity (private key) for public key authentication is read",
+						},
 					},
 				},
 				{

cmd/cli/commands/config/environment.go

@@ -28,8 +28,9 @@ type Environment struct {
 
 // Forwarding defines configuration for port forwarding.
 type Forwarding struct {
-	ServerURL string `yaml:"server_url" json:"server_url"`
-	LocalPort string `yaml:"local_port" json:"local_port"`
+	ServerURL    string `yaml:"server_url" json:"server_url"`
+	LocalPort    string `yaml:"local_port" json:"local_port"`
+	IdentityFile string `yaml:"identity_file" json:"identity_file"`
 }
 
 // AddEnvironmentToConfig adds a new environment to CLIConfig.
@@ -47,8 +48,9 @@ func AddEnvironmentToConfig(c *cli.Context, cfg *CLIConfig, environmentID string
 		Token:    c.String(commands.TokenKey),
 		Insecure: c.Bool(commands.InsecureKey),
 		Forwarding: Forwarding{
-			ServerURL: c.String(commands.FwServerURLKey),
-			LocalPort: c.String(commands.FwLocalPortKey),
+			ServerURL:    c.String(commands.FwServerURLKey),
+			LocalPort:    c.String(commands.FwLocalPortKey),
+			IdentityFile: c.String(commands.IdentityFileKey),
 		},
 	}
 
@@ -99,6 +101,10 @@ func updateEnvironmentInConfig(c *cli.Context, cfg *CLIConfig, environmentID str
 		newEnvironment.Forwarding.LocalPort = c.String(commands.FwLocalPortKey)
 	}
 
+	if c.IsSet(commands.IdentityFileKey) {
+		newEnvironment.Forwarding.IdentityFile = c.String(commands.IdentityFileKey)
+	}
+
 	if newEnvironment == environment {
 		return errors.New("config unchanged. Set different option values to update.") // nolint
 	}

cmd/cli/commands/global/actions.go

@@ -60,7 +60,7 @@ func forward(cliCtx *cli.Context) error {
 		return err
 	}
 
-	tunnel, err := commands.BuildTunnel(cliCtx, remoteURL.Host)
+	tunnel, err := commands.BuildTunnel(cliCtx, remoteURL)
 	if err != nil {
 		return err
 	}

cmd/cli/commands/global/command_list.go

@@ -46,6 +46,10 @@ func List() []*cli.Command {
 					Name:  "forwarding-local-port",
 					Usage: "local port for forwarding to the Database Lab instance",
 				},
+				&cli.StringFlag{
+					Name:  "identity-file",
+					Usage: "select a file from which the identity (private key) for public key authentication is read",
+				},
 			},
 			Action: initCLI,
 		},

cmd/cli/commands/port_forwarding.go

@@ -17,31 +17,31 @@ import (
 )
 
 // BuildTunnel creates a new instance of SSH tunnel.
-func BuildTunnel(cliCtx *cli.Context, remoteHost string) (*portfwd.SSHTunnel, error) {
-	remoteURL, err := url.Parse(cliCtx.String(URLKey))
+func BuildTunnel(cliCtx *cli.Context, remoteHost *url.URL) (*portfwd.SSHTunnel, error) {
+	localEndpoint := forwardingLocalEndpoint(remoteHost, cliCtx.String(FwLocalPortKey))
+
+	serverURL, err := url.Parse(cliCtx.String(FwServerURLKey))
 	if err != nil {
 		return nil, err
 	}
 
-	localEndpoint := forwardingLocalEndpoint(remoteURL, cliCtx.String(FwLocalPortKey))
-
-	serverURL, err := url.Parse(cliCtx.String(FwServerURLKey))
+	authMethod, err := getAuthMethod(cliCtx)
 	if err != nil {
 		return nil, err
 	}
 
 	sshConfig := &ssh.ClientConfig{
 		User: serverURL.User.Username(),
 		Auth: []ssh.AuthMethod{
-			portfwd.SSHAgent(),
+			authMethod,
 		},
 		HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
 			// Always accept key.
 			return nil
 		},
 	}
 
-	tunnel := portfwd.NewTunnel(localEndpoint, serverURL.Host, remoteHost, sshConfig)
+	tunnel := portfwd.NewTunnel(localEndpoint, serverURL.Host, remoteHost.Host, sshConfig)
 
 	return tunnel, nil
 }
@@ -54,6 +54,19 @@ func forwardingLocalEndpoint(remoteURL *url.URL, localPort string) string {
 	return fmt.Sprintf("%s:%s", "127.0.0.1", localPort)
 }
 
+func getAuthMethod(cliCtx *cli.Context) (ssh.AuthMethod, error) {
+	if cliCtx.String(IdentityFileKey) != "" {
+		authMethod, err := portfwd.ReadAuthFromIdentityFile(cliCtx.String(IdentityFileKey))
+		if err != nil {
+			return nil, err
+		}
+
+		return authMethod, nil
+	}
+
+	return portfwd.SSHAgent(), nil
+}
+
 // BuildHostname builds a hostname string.
 func BuildHostname(host, port string) string {
 	return fmt.Sprintf("%s:%s", host, port)

cmd/cli/main.go

@@ -58,12 +58,17 @@ func main() {
 			&cli.StringFlag{
 				Name:    "forwarding-server-url",
 				Usage:   "forwarding server URL of Database Lab instance",
-				EnvVars: []string{"DBLAB_FORWARDING_SERVER_URL"},
+				EnvVars: []string{"DBLAB_CLI_FORWARDING_SERVER_URL"},
 			},
 			&cli.StringFlag{
 				Name:    "forwarding-local-port",
 				Usage:   "local port for forwarding to the Database Lab instance",
-				EnvVars: []string{"DBLAB_FORWARDING_LOCAL_PORT"},
+				EnvVars: []string{"DBLAB_CLI_FORWARDING_LOCAL_PORT"},
+			},
+			&cli.StringFlag{
+				Name:    "identity-file",
+				Usage:   "select a file from which the identity (private key) for public key authentication is read",
+				EnvVars: []string{"DBLAB_CLI_IDENTITY_FILE"},
 			},
 			&cli.BoolFlag{
 				Name:    "debug",
@@ -126,6 +131,12 @@ func loadEnvironmentParams(c *cli.Context) error {
 				return err
 			}
 		}
+
+		if !c.IsSet(commands.IdentityFileKey) {
+			if err := c.Set(commands.IdentityFileKey, env.Forwarding.IdentityFile); err != nil {
+				return err
+			}
+		}
 	}
 
 	return nil

pkg/portfwd/sshtunnel.go

@@ -7,7 +7,9 @@ package portfwd
 
 import (
 	"context"
+	"fmt"
 	"io"
+	"io/ioutil"
 	"net"
 	"os"
 
@@ -159,3 +161,20 @@ func SSHAgent() ssh.AuthMethod {
 
 	return nil
 }
+
+// ReadAuthFromIdentityFile reads an identity file and returns an AuthMethod that uses the given key pairs.
+func ReadAuthFromIdentityFile(identityFilename string) (ssh.AuthMethod, error) {
+	log.Dbg(fmt.Sprintf("Read identity file %q", identityFilename))
+
+	privateKey, err := ioutil.ReadFile(identityFilename)
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to read an identity file")
+	}
+
+	signer, err := ssh.ParsePrivateKey(privateKey)
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to parse a private key")
+	}
+
+	return ssh.PublicKeys(signer), nil
+}