Use env.ssh_available field in config-ssh (#197)

deansheather · web-flow · commit 44f4a0442a56 · 2020-12-04T05:26:16.000+10:00
diff --git a/coder-sdk/env.go b/coder-sdk/env.go
@@ -32,6 +32,7 @@ type Environment struct {
 	LastOpenedAt     time.Time        `json:"last_opened_at"     table:"-"`
 	LastConnectionAt time.Time        `json:"last_connection_at" table:"-"`
 	AutoOffThreshold Duration         `json:"auto_off_threshold" table:"-"`
+	SSHAvailable     bool             `json:"ssh_available"      table:"-"`
 }
 
 // RebuildMessage defines the message shown when an Environment requires a rebuild for it can be accessed.
diff --git a/internal/cmd/configssh.go b/internal/cmd/configssh.go
@@ -20,6 +20,17 @@ import (
 	"golang.org/x/xerrors"
 )
 
+const sshStartToken = "# ------------START-CODER-ENTERPRISE-----------"
+const sshStartMessage = `# The following has been auto-generated by "coder config-ssh"
+# to make accessing your Coder Enterprise environments easier.
+#
+# To remove this blob, run:
+#
+#    coder config-ssh --remove
+#
+# You should not hand-edit this section, unless you are deleting it.`
+const sshEndToken = "# ------------END-CODER-ENTERPRISE------------"
+
 func configSSHCmd() *cobra.Command {
 	var (
 		configpath string
@@ -39,17 +50,6 @@ func configSSHCmd() *cobra.Command {
 }
 
 func configSSH(configpath *string, remove *bool) func(cmd *cobra.Command, _ []string) error {
-	const startToken = "# ------------START-CODER-ENTERPRISE-----------"
-	startMessage := `# The following has been auto-generated by "coder config-ssh"
-# to make accessing your Coder Enterprise environments easier.
-#
-# To remove this blob, run:
-#
-#    coder config-ssh --remove
-#
-# You should not hand-edit this section, unless you are deleting it.`
-	const endToken = "# ------------END-CODER-ENTERPRISE------------"
-
 	return func(cmd *cobra.Command, _ []string) error {
 		ctx := cmd.Context()
 		usr, err := user.Current()
@@ -71,14 +71,11 @@ func configSSH(configpath *string, remove *bool) func(cmd *cobra.Command, _ []st
 			return xerrors.Errorf("read ssh config file %q: %w", *configpath, err)
 		}
 
-		startIndex := strings.Index(currentConfig, startToken)
-		endIndex := strings.Index(currentConfig, endToken)
-
+		currentConfig, didRemoveConfig := removeOldConfig(currentConfig)
 		if *remove {
-			if startIndex == -1 || endIndex == -1 {
+			if !didRemoveConfig {
 				return xerrors.Errorf("the Coder Enterprise ssh configuration section could not be safely deleted or does not exist")
 			}
-			currentConfig = currentConfig[:startIndex-1] + currentConfig[endIndex+len(endToken)+1:]
 
 			err = writeStr(*configpath, currentConfig)
 			if err != nil {
@@ -93,10 +90,6 @@ func configSSH(configpath *string, remove *bool) func(cmd *cobra.Command, _ []st
 			return err
 		}
 
-		if !isSSHAvailable(ctx) {
-			return xerrors.New("SSH is disabled or not available for your Coder Enterprise deployment.")
-		}
-
 		user, err := client.Me(ctx)
 		if err != nil {
 			return xerrors.Errorf("fetch username: %w", err)
@@ -109,14 +102,19 @@ func configSSH(configpath *string, remove *bool) func(cmd *cobra.Command, _ []st
 		if len(envs) < 1 {
 			return xerrors.New("no environments found")
 		}
-		newConfig, err := makeNewConfigs(user.Username, envs, startToken, startMessage, endToken, privateKeyFilepath)
+
+		if !sshAvailable(envs) {
+			return xerrors.New("SSH is disabled or not available for any environments in your Coder Enterprise deployment.")
+		}
+
+		err = canConnectSSH(ctx)
 		if err != nil {
-			return xerrors.Errorf("make new ssh configurations: %w", err)
+			return xerrors.Errorf("check if SSH is available: unable to connect to SSH endpoint: %w", err)
 		}
 
-		// if we find the old config, remove those chars from the string
-		if startIndex != -1 && endIndex != -1 {
-			currentConfig = currentConfig[:startIndex-1] + currentConfig[endIndex+len(endToken)+1:]
+		newConfig, err := makeNewConfigs(user.Username, envs, privateKeyFilepath)
+		if err != nil {
+			return xerrors.Errorf("make new ssh configurations: %w", err)
 		}
 
 		err = os.MkdirAll(filepath.Dir(*configpath), os.ModePerm)
@@ -145,6 +143,57 @@ func configSSH(configpath *string, remove *bool) func(cmd *cobra.Command, _ []st
 	}
 }
 
+// removeOldConfig removes the old ssh configuration from the user's sshconfig.
+// Returns true if the config was modified.
+func removeOldConfig(config string) (string, bool) {
+	startIndex := strings.Index(config, sshStartToken)
+	endIndex := strings.Index(config, sshEndToken)
+
+	if startIndex == -1 || endIndex == -1 {
+		return config, false
+	}
+	config = config[:startIndex-1] + config[endIndex+len(sshEndToken)+1:]
+
+	return config, true
+}
+
+// sshAvailable returns true if SSH is available for at least one environment.
+func sshAvailable(envs []coder.Environment) bool {
+	for _, env := range envs {
+		if env.SSHAvailable {
+			return true
+		}
+	}
+
+	return false
+}
+
+// canConnectSSH returns an error if we cannot dial the SSH port.
+func canConnectSSH(ctx context.Context) error {
+	ctx, cancel := context.WithTimeout(ctx, 3*time.Second)
+	defer cancel()
+
+	host, err := configuredHostname()
+	if err != nil {
+		return xerrors.Errorf("get configured manager hostname: %w", err)
+	}
+
+	var (
+		dialer   net.Dialer
+		hostPort = net.JoinHostPort(host, "22")
+	)
+	conn, err := dialer.DialContext(ctx, "tcp", hostPort)
+	if err != nil {
+		if err == context.DeadlineExceeded {
+			err = xerrors.New("timed out after 3 seconds")
+		}
+		return xerrors.Errorf("dial tcp://%v: %w", hostPort, err)
+	}
+	conn.Close()
+
+	return nil
+}
+
 func writeSSHKey(ctx context.Context, client *coder.Client, privateKeyPath string) error {
 	key, err := client.SSHKey(ctx)
 	if err != nil {
@@ -153,17 +202,21 @@ func writeSSHKey(ctx context.Context, client *coder.Client, privateKeyPath strin
 	return ioutil.WriteFile(privateKeyPath, []byte(key.PrivateKey), 0400)
 }
 
-func makeNewConfigs(userName string, envs []coder.Environment, startToken, startMsg, endToken, privateKeyFilepath string) (string, error) {
+func makeNewConfigs(userName string, envs []coder.Environment, privateKeyFilepath string) (string, error) {
 	hostname, err := configuredHostname()
 	if err != nil {
 		return "", err
 	}
 
-	newConfig := fmt.Sprintf("\n%s\n%s\n\n", startToken, startMsg)
+	newConfig := fmt.Sprintf("\n%s\n%s\n\n", sshStartToken, sshStartMessage)
 	for _, env := range envs {
+		if !env.SSHAvailable {
+			continue
+		}
+
 		newConfig += makeSSHConfig(hostname, userName, env.Name, privateKeyFilepath)
 	}
-	newConfig += fmt.Sprintf("\n%s\n", endToken)
+	newConfig += fmt.Sprintf("\n%s\n", sshEndToken)
 
 	return newConfig, nil
 }
@@ -181,20 +234,6 @@ func makeSSHConfig(host, userName, envName, privateKeyFilepath string) string {
 `, envName, host, userName, envName, privateKeyFilepath)
 }
 
-func isSSHAvailable(ctx context.Context) bool {
-	ctx, cancel := context.WithTimeout(ctx, 3*time.Second)
-	defer cancel()
-
-	host, err := configuredHostname()
-	if err != nil {
-		return false
-	}
-
-	var dialer net.Dialer
-	_, err = dialer.DialContext(ctx, "tcp", net.JoinHostPort(host, "22"))
-	return err == nil
-}
-
 func configuredHostname() (string, error) {
 	u, err := config.URL.Read()
 	if err != nil {

Original file line number	Diff line number	Diff line change
`@@ -32,6 +32,7 @@ type Environment struct {`
`32`	`32`	LastOpenedAt time.Time `json:"last_opened_at" table:"-"`
`33`	`33`	LastConnectionAt time.Time `json:"last_connection_at" table:"-"`
`34`	`34`	AutoOffThreshold Duration `json:"auto_off_threshold" table:"-"`
	`35`	+ SSHAvailable bool `json:"ssh_available" table:"-"`
`35`	`36`	`}`
`36`	`37`
`37`	`38`	`// RebuildMessage defines the message shown when an Environment requires a rebuild for it can be accessed.`