coder · Emyrk · Feb 3, 2023 · Feb 2, 2023 · Feb 2, 2023 · Feb 2, 2023
diff --git a/cli/configssh.go b/cli/configssh.go
@@ -206,7 +206,11 @@ func configSSH() *cobra.Command {
 			// Parse the previous configuration only if config-ssh
 			// has been run previously.
 			var lastConfig *sshConfigOptions
-			if section, ok := sshConfigGetCoderSection(configRaw); ok {
+			section, ok, err := sshConfigGetCoderSection(configRaw)
+			if err != nil {
+				return err
+			}
+			if ok {
 				c := sshConfigParseLastOptions(bytes.NewReader(section))
 				lastConfig = &c
 			}
@@ -249,7 +253,10 @@ func configSSH() *cobra.Command {
 			configModified := configRaw
 
 			buf := &bytes.Buffer{}
-			before, after := sshConfigSplitOnCoderSection(configModified)
+			before, _, after, err := sshConfigSplitOnCoderSection(configModified)
+			if err != nil {
+				return err
+			}
 			// Write the first half of the users config file to buf.
 			_, _ = buf.Write(before)
 
@@ -418,22 +425,39 @@ func sshConfigParseLastOptions(r io.Reader) (o sshConfigOptions) {
 	return o
 }
 
-func sshConfigGetCoderSection(data []byte) (section []byte, ok bool) {
-	startIndex := bytes.Index(data, []byte(sshStartToken))
-	endIndex := bytes.Index(data, []byte(sshEndToken))
-	if startIndex != -1 && endIndex != -1 {
-		return data[startIndex : endIndex+len(sshEndToken)], true
+// sshConfigGetCoderSection is a helper function that only returns the coder
+// section of the SSH config and a boolean if it exists.
+func sshConfigGetCoderSection(data []byte) (section []byte, ok bool, err error) {
+	_, section, _, err = sshConfigSplitOnCoderSection(data)
+	if err != nil {
+		return nil, false, err
 	}
-	return nil, false
+
+	return section, len(section) > 0, nil
 }
 
-// sshConfigSplitOnCoderSection splits the SSH config into two sections,
-// before contains the lines before sshStartToken and after contains the
-// lines after sshEndToken.
-func sshConfigSplitOnCoderSection(data []byte) (before, after []byte) {
+// sshConfigSplitOnCoderSection splits the SSH config into 3 sections.
+// All lines before sshStartToken, the coder section, and all lines after
+// sshEndToken.
+func sshConfigSplitOnCoderSection(data []byte) (before, section []byte, after []byte, err error) {
+	startCount := bytes.Count(data, []byte(sshStartToken))
+	endCount := bytes.Count(data, []byte(sshEndToken))
+	if startCount > 1 || endCount > 1 {
+		return nil, nil, nil, xerrors.New("Malformed config: ssh config has multiple coder sections, please remove all but one")
+	}
+
 	startIndex := bytes.Index(data, []byte(sshStartToken))
 	endIndex := bytes.Index(data, []byte(sshEndToken))
+	if startIndex == -1 && endIndex != -1 {
+		return nil, nil, nil, xerrors.New("Malformed config: ssh config has end header, but missing start header")
+	}
+	if startIndex != -1 && endIndex == -1 {
+		return nil, nil, nil, xerrors.New("Malformed config: ssh config has start header, but missing end header")
+	}
 	if startIndex != -1 && endIndex != -1 {
+		if startIndex > endIndex {
+			return nil, nil, nil, xerrors.New("Malformed config: ssh config has coder section, but it is malformed and the END header is before the START header")
+		}
 		// We use -1 and +1 here to also include the preceding
 		// and trailing newline, where applicable.
 		start := startIndex
@@ -444,10 +468,10 @@ func sshConfigSplitOnCoderSection(data []byte) (before, after []byte) {
 		if end < len(data) {
 			end++
 		}
-		return data[:start], data[end:]
+		return data[:start], data[start:end], data[end:], nil
 	}
 
-	return data, nil
+	return data, nil, nil, nil
 }
 
 // writeWithTempFileAndMove writes to a temporary file in the same

diff --git a/cli/configssh_internal_test.go b/cli/configssh_internal_test.go
@@ -11,6 +11,125 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+func Test_sshConfigSplitOnCoderSection(t *testing.T) {
+	t.Parallel()
+
+	testCases := []struct {
+		Name    string
+		Input   string
+		Before  string
+		Section string
+		After   string
+		Err     bool
+	}{
+		{
+			Name:    "Empty",
+			Input:   "",
+			Before:  "",
+			Section: "",
+			After:   "",
+			Err:     false,
+		},
+		{
+			Name:    "JustSection",
+			Input:   strings.Join([]string{sshStartToken, sshEndToken}, "\n"),
+			Before:  "",
+			Section: strings.Join([]string{sshStartToken, sshEndToken}, "\n"),
+			After:   "",
+			Err:     false,
+		},
+		{
+			Name:    "NoSection",
+			Input:   strings.Join([]string{"# Some content"}, "\n"),
+			Before:  "# Some content",
+			Section: "",
+			After:   "",
+			Err:     false,
+		},
+		{
+			Name: "Normal",
+			Input: strings.Join([]string{
+				"# Content before the section",
+				sshStartToken,
+				sshEndToken,
+				"# Content after the section",
+			}, "\n"),
+			Before:  "# Content before the section",
+			Section: strings.Join([]string{"", sshStartToken, sshEndToken, ""}, "\n"),
+			After:   "# Content after the section",
+			Err:     false,
+		},
+		{
+			Name: "OutOfOrder",
+			Input: strings.Join([]string{
+				"# Content before the section",
+				sshEndToken,
+				sshStartToken,
+				"# Content after the section",
+			}, "\n"),
+			Err: true,
+		},
+		{
+			Name: "MissingStart",
+			Input: strings.Join([]string{
+				"# Content before the section",
+				sshEndToken,
+				"# Content after the section",
+			}, "\n"),
+			Err: true,
+		},
+		{
+			Name: "MissingEnd",
+			Input: strings.Join([]string{
+				"# Content before the section",
+				sshEndToken,
+				"# Content after the section",
+			}, "\n"),
+			Err: true,
+		},
+		{
+			Name: "ExtraStart",
+			Input: strings.Join([]string{
+				"# Content before the section",
+				sshStartToken,
+				sshEndToken,
+				sshStartToken,
+				"# Content after the section",
+			}, "\n"),
+			Err: true,
+		},
+		{
+			Name: "ExtraEnd",
+			Input: strings.Join([]string{
+				"# Content before the section",
+				sshStartToken,
+				sshEndToken,
+				sshEndToken,
+				"# Content after the section",
+			}, "\n"),
+			Err: true,
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+		t.Run(tc.Name, func(t *testing.T) {
+			t.Parallel()
+
+			before, section, after, err := sshConfigSplitOnCoderSection([]byte(tc.Input))
+			if tc.Err {
+				require.Error(t, err)
+				return
+			}
+
+			require.NoError(t, err)
+			require.Equal(t, tc.Before, string(before), "before")
+			require.Equal(t, tc.Section, string(section), "section")
+			require.Equal(t, tc.After, string(after), "after")
+		})
+	}
+}
+
 // This test tries to mimic the behavior of OpenSSH
 // when executing e.g. a ProxyCommand.
 func Test_sshConfigExecEscape(t *testing.T) {

diff --git a/cli/configssh_test.go b/cli/configssh_test.go
@@ -529,6 +529,36 @@ func TestConfigSSH_FileWriteAndOptionsFlow(t *testing.T) {
 				"--yes",
 			},
 		},
+		{
+			name:    "Start/End out of order",
+			matches: []match{
+				//{match: "Continue?", write: "yes"},
+			},
+			writeConfig: writeConfig{
+				ssh: strings.Join([]string{
+					"# Content before coder block",
+					headerEnd,
+					headerStart,
+					"# Content after coder block",
+				}, "\n"),
+			},
+			wantErr: true,
+		},
+		{
+			name:    "Multiple sections",
+			matches: []match{
+				//{match: "Continue?", write: "yes"},
+			},
+			writeConfig: writeConfig{
+				ssh: strings.Join([]string{
+					headerStart,
+					headerEnd,
+					headerStart,
+					headerEnd,
+				}, "\n"),
+			},
+			wantErr: true,
+		},
 	}
 	for _, tt := range tests {
 		tt := tt