coder · mafredri · May 27, 2022 · kylecarbs · May 27, 2022 · mafredri
diff --git a/cli/configssh.go b/cli/configssh.go
@@ -1,12 +1,13 @@
 package cli
 
 import (
+	"bufio"
 	"fmt"
+	"io/fs"
 	"os"
 	"path/filepath"
 	"runtime"
 	"strings"
-	"sync"
 
 	"github.com/cli/safeexec"
 	"github.com/spf13/cobra"
@@ -18,15 +19,19 @@ import (
 	"github.com/coder/coder/codersdk"
 )
 
-const sshStartToken = "# ------------START-CODER-----------"
-const sshStartMessage = `# This was generated by "coder config-ssh".
-#
-# 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------------"
+const (
+	// Include path is relative to `~/.ssh` and each workspace will
+	// have a separate file (e.g. `~/.ssh/coder.d/host-my-workspace`).
+	// By prefixing hosts as `host-` we give ourselves the flexibility
+	// to manage other files in this folder as well, e.g. keys, vscode
+	// specific config (i.e. for only listing coder files in vscode),
+	// etc.
+	sshEnabledLine = "Include coder.d/host-*"
+	// TODO(mafredri): Does this hold on Windows?
+	sshCoderConfigd = "~/.ssh/coder.d"
+	// TODO(mafredri): Write a README to the folder?
+	// sshCoderConfigdReadme = `Information, tricks, removal, etc.`
+)
 
 func configSSH() *cobra.Command {
 	var (
@@ -52,18 +57,11 @@ func configSSH() *cobra.Command {
 			if err != nil {
 				return err
 			}
+			dirname, _ := os.UserHomeDir()
 			if strings.HasPrefix(sshConfigFile, "~/") {
-				dirname, _ := os.UserHomeDir()
 				sshConfigFile = filepath.Join(dirname, sshConfigFile[2:])
 			}
-			// Doesn't matter if this fails, because we write the file anyways.
-			sshConfigContentRaw, _ := os.ReadFile(sshConfigFile)
-			sshConfigContent := string(sshConfigContentRaw)
-			startIndex := strings.Index(sshConfigContent, sshStartToken)
-			endIndex := strings.Index(sshConfigContent, sshEndToken)
-			if startIndex != -1 && endIndex != -1 {
-				sshConfigContent = sshConfigContent[:startIndex-1] + sshConfigContent[endIndex+len(sshEndToken):]
-			}
+			confd := filepath.Join(dirname, sshCoderConfigd[2:])
 
 			workspaces, err := client.WorkspacesByOwner(cmd.Context(), organization.ID, codersdk.Me)
 			if err != nil {
@@ -78,10 +76,74 @@ func configSSH() *cobra.Command {
 				return err
 			}
 
+			enabled := false
+			err = func() error {
+				exists := true
+				configRaw, err := os.Open(sshConfigFile)
+				if err != nil && !xerrors.Is(err, fs.ErrNotExist) {
+					return err
+				} else if xerrors.Is(err, fs.ErrNotExist) {
+					exists = false
+				}
+				defer configRaw.Close()
+
+				if exists {
+					s := bufio.NewScanner(configRaw)
+					for s.Scan() {
+						if strings.HasPrefix(s.Text(), sshEnabledLine) {
+							enabled = true
+							break
+						}
+					}
+					if s.Err() != nil {
+						return err
+					}
+				}
+				return nil
+			}()
+			if err != nil {
+				return err
+			}
+
+			if !enabled {
+				_, err = cliui.Prompt(cmd, cliui.PromptOptions{
+					Text:      fmt.Sprintf("The following line will be added to %s:\n\n    %s\n\n  And configuration files will be stored in ~/.ssh/coder.d\n\n  Continue?", sshConfigFile, sshEnabledLine),
+					IsConfirm: true,
+				})
+				if err != nil {
+					return err
+				}
+				_, _ = fmt.Fprintf(cmd.OutOrStdout(), "\n")
+
+				// Create directory first in case of error since we do not check for existence.
+				err = os.Mkdir(confd, 0o700)
+				if err != nil && !xerrors.Is(err, fs.ErrExist) {
+					return err
+				}
+
+				f, err := os.OpenFile(sshConfigFile, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0o644)
+				if err != nil {
+					return err
+				}
+
+				// TODO(mafredri): Only add newline if necessary.
+				_, err = f.WriteString("\n" + sshEnabledLine)
+				if err != nil {
+					return err
+				}
+
+				err = f.Close()
+				if err != nil {
+					return err
+				}
+
+				_, _ = fmt.Fprintf(cmd.OutOrStdout(), "* Added Include directive to %s\n", sshConfigFile)
+				_, _ = fmt.Fprintf(cmd.OutOrStdout(), "* Created configuration directory %s\n", confd)
+			}
+
 			root := createConfig(cmd)
-			sshConfigContent += "\n" + sshStartToken + "\n" + sshStartMessage + "\n\n"
-			sshConfigContentMutex := sync.Mutex{}
 			var errGroup errgroup.Group
+			// TODO(mafredri): Delete configurations that no longer exist.
 			for _, workspace := range workspaces {
 				workspace := workspace
 				errGroup.Go(func() error {
@@ -94,7 +156,6 @@ func configSSH() *cobra.Command {
 							continue
 						}
 						for _, agent := range resource.Agents {
-							sshConfigContentMutex.Lock()
 							hostname := workspace.Name
 							if len(resource.Agents) > 1 {
 								hostname += "." + agent.Name
@@ -119,8 +180,13 @@ func configSSH() *cobra.Command {
 							if !skipProxyCommand {
 								configOptions = append(configOptions, fmt.Sprintf("\tProxyCommand %q --global-config %q ssh --stdio %s", binaryFile, root, hostname))
 							}
-							sshConfigContent += strings.Join(configOptions, "\n") + "\n"
-							sshConfigContentMutex.Unlock()
+
+							dest := filepath.Join(confd, fmt.Sprintf("host-%s", hostname))
+							// TODO(mafredri): Avoid re-write if files match.
+							err := os.WriteFile(dest, []byte(strings.Join(configOptions, "\n")), 0o600)
+							if err != nil {
+								return err
+							}
 						}
 					}
 					return nil
@@ -130,18 +196,10 @@ func configSSH() *cobra.Command {
 			if err != nil {
 				return err
 			}
-			sshConfigContent += "\n" + sshEndToken
-			err = os.MkdirAll(filepath.Dir(sshConfigFile), os.ModePerm)
-			if err != nil {
-				return err
-			}
-			err = os.WriteFile(sshConfigFile, []byte(sshConfigContent), os.ModePerm)
-			if err != nil {
-				return err
-			}
-			_, _ = fmt.Fprintf(cmd.OutOrStdout(), "An auto-generated ssh config was written to %q\n", sshConfigFile)
-			_, _ = fmt.Fprintln(cmd.OutOrStdout(), "You should now be able to ssh into your workspace")
-			_, _ = fmt.Fprintf(cmd.OutOrStdout(), "For example, try running\n\n\t$ ssh coder.%s\n\n", workspaces[0].Name)
+
+			_, _ = fmt.Fprintf(cmd.OutOrStdout(), "* Created workspace configurations in %s\n\n", confd)
+			_, _ = fmt.Fprintln(cmd.OutOrStdout(), "You should now be able to ssh into your workspace.")
+			_, _ = fmt.Fprintf(cmd.OutOrStdout(), "For example, try running:\n\n\t$ ssh coder.%s\n\n", workspaces[0].Name)
 			return nil
 		},
 	}