diff --git a/cli/configssh.go b/cli/configssh.go index 000c0f8965..090489c2ba 100644 --- 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 index 9852d933eb..88de43b73a 100644 --- 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 index ceb7e42317..046f763285 100644 --- 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