add SSH master connection feature

By default, sshcode will now start a master connection with no command
so that users only need to authenticate once and multiple connections
don't need to be established. This speeds up load times significantly as
there are less handshakes required.

To disable this behaviour you can use `--no-reuse-connection`.

Author: deansheather

main.go

@@ -34,11 +34,12 @@ var _ interface {
 } = new(rootCmd)
 
 type rootCmd struct {
-	skipSync     bool
-	syncBack     bool
-	printVersion bool
-	bindAddr     string
-	sshFlags     string
+	skipSync          bool
+	syncBack          bool
+	printVersion      bool
+	noReuseConnection bool
+	bindAddr          string
+	sshFlags          string
 }
 
 func (c *rootCmd) Spec() cli.CommandSpec {
@@ -53,6 +54,7 @@ func (c *rootCmd) RegisterFlags(fl *flag.FlagSet) {
 	fl.BoolVar(&c.skipSync, "skipsync", false, "skip syncing local settings and extensions to remote host")
 	fl.BoolVar(&c.syncBack, "b", false, "sync extensions back on termination")
 	fl.BoolVar(&c.printVersion, "version", false, "print version information and exit")
+	fl.BoolVar(&c.noReuseConnection, "no-reuse-connection", false, "do not reuse SSH connection via control socket")
 	fl.StringVar(&c.bindAddr, "bind", "", "local bind address for ssh tunnel")
 	fl.StringVar(&c.sshFlags, "ssh-flags", "", "custom SSH flags")
 }
@@ -76,10 +78,11 @@ func (c *rootCmd) Run(fl *flag.FlagSet) {
 	}
 
 	err := sshCode(host, dir, options{
-		skipSync: c.skipSync,
-		sshFlags: c.sshFlags,
-		bindAddr: c.bindAddr,
-		syncBack: c.syncBack,
+		skipSync:          c.skipSync,
+		sshFlags:          c.sshFlags,
+		bindAddr:          c.bindAddr,
+		syncBack:          c.syncBack,
+		noReuseConnection: c.noReuseConnection,
 	})
 
 	if err != nil {
@@ -101,7 +104,7 @@ Environment variables:
 More info: https://github.com/cdr/sshcode
 
 Arguments:
-%vHOST is passed into the ssh command. Valid formats are '<ip-address>' or 'gcp:<instance-name>'. 
+%vHOST is passed into the ssh command. Valid formats are '<ip-address>' or 'gcp:<instance-name>'.
 %vDIR is optional.`,
 		helpTab, vsCodeConfigDirEnv,
 		helpTab, vsCodeExtensionsDirEnv,

sshcode.go

@@ -20,14 +20,16 @@ import (
 )
 
 const codeServerPath = "~/.cache/sshcode/sshcode-server"
+const sshControlPath = "~/.ssh/control-%h-%p-%r"
 
 type options struct {
-	skipSync   bool
-	syncBack   bool
-	noOpen     bool
-	bindAddr   string
-	remotePort string
-	sshFlags   string
+	skipSync          bool
+	syncBack          bool
+	noOpen            bool
+	noReuseConnection bool
+	bindAddr          string
+	remotePort        string
+	sshFlags          string
 }
 
 func sshCode(host, dir string, o options) error {
@@ -53,6 +55,41 @@ func sshCode(host, dir string, o options) error {
 		return xerrors.Errorf("failed to find available remote port: %w", err)
 	}
 
+	// Start SSH master connection socket. This prevents multiple password prompts from appearing as authentication
+	// only happens on the initial connection.
+	var sshMasterCmd *exec.Cmd
+	if !o.noReuseConnection {
+		newSSHFlags := fmt.Sprintf(`%v -o "ControlPath=%v"`, o.sshFlags, sshControlPath)
+
+		// -MN means "start a master socket and don't open a session, just connect".
+		sshMasterCmdStr := fmt.Sprintf(`ssh %v -MN %v`, newSSHFlags, host)
+		sshMasterCmd = exec.Command("sh", "-c", sshMasterCmdStr)
+		sshMasterCmd.Stdin = os.Stdin
+		sshMasterCmd.Stdout = os.Stdout
+		sshMasterCmd.Stderr = os.Stderr
+		err = sshMasterCmd.Start()
+		if err != nil {
+			flog.Error("failed to start SSH master connection, disabling connection reuse feature: %v", err)
+			o.noReuseConnection = true
+		} else {
+			// Wait for master to be ready.
+			err = checkSSHMaster(newSSHFlags, host)
+			if err != nil {
+				flog.Error("SSH master failed to start in time, disabling connection reuse feature: %v", err)
+				o.noReuseConnection = true
+				if sshMasterCmd.Process != nil {
+					err = sshMasterCmd.Process.Kill()
+					if err != nil {
+						flog.Error("failed to kill SSH master connection, ignoring: %v", err)
+					}
+				}
+			} else {
+				sshMasterCmd.Stdin = nil
+				o.sshFlags = newSSHFlags
+			}
+		}
+	}
+
 	dlScript := downloadScript(codeServerPath)
 
 	// Downloads the latest code-server and allows it to be executed.
@@ -146,22 +183,39 @@ func sshCode(host, dir string, o options) error {
 	case <-ctx.Done():
 	case <-c:
 	}
+	flog.Info("exiting")
 
-	if !o.syncBack || o.skipSync {
-		flog.Info("shutting down")
-		return nil
-	}
+	if o.syncBack && !o.skipSync {
+		flog.Info("synchronizing VS Code back to local")
 
-	flog.Info("synchronizing VS Code back to local")
+		err = syncExtensions(o.sshFlags, host, true)
+		if err != nil {
+			flog.Error("failed to sync extensions back: %v", err)
+		}
 
-	err = syncExtensions(o.sshFlags, host, true)
-	if err != nil {
-		return xerrors.Errorf("failed to sync extensions back: %w", err)
+		err = syncUserSettings(o.sshFlags, host, true)
+		if err != nil {
+			flog.Error("failed to sync user settings settings back: %v", err)
+		}
 	}
 
-	err = syncUserSettings(o.sshFlags, host, true)
-	if err != nil {
-		return xerrors.Errorf("failed to sync user settings settings back: %w", err)
+	// Kill the master connection if we made one.
+	if !o.noReuseConnection {
+		// Try using the -O exit syntax first before killing the master.
+		sshCmdStr = fmt.Sprintf(`ssh %v -O exit %v`, o.sshFlags, host)
+		sshCmd = exec.Command("sh", "-c", sshCmdStr)
+		sshCmd.Stdout = os.Stdout
+		sshCmd.Stderr = os.Stderr
+		err = sshCmd.Run()
+		if err != nil {
+			flog.Error("failed to gracefully stop SSH master connection, killing: %v", err)
+			if sshMasterCmd.Process != nil {
+				err = sshMasterCmd.Process.Kill()
+				if err != nil {
+					flog.Error("failed to kill SSH master connection, ignoring: %v", err)
+				}
+			}
+		}
 	}
 
 	return nil
@@ -263,6 +317,26 @@ func randomPort() (string, error) {
 	return "", xerrors.Errorf("max number of tries exceeded: %d", maxTries)
 }
 
+// checkSSHMaster polls every second for 30 seconds to check if the SSH master
+// is ready.
+func checkSSHMaster(sshFlags string, host string) (err error) {
+	maxTries := 30
+	check := func() error {
+		sshCmdStr := fmt.Sprintf(`ssh %v -O check %v`, sshFlags, host)
+		sshCmd := exec.Command("sh", "-c", sshCmdStr)
+		return sshCmd.Run()
+	}
+
+	for i := 0; i < maxTries; i++ {
+		err = check()
+		if err == nil {
+			return nil
+		}
+		time.Sleep(time.Second)
+	}
+	return err
+}
+
 func syncUserSettings(sshFlags string, host string, back bool) error {
 	localConfDir, err := configDir()
 	if err != nil {