fix: SSHConfig: check for multiple start/end blocks

src/sshConfig.test.ts

@@ -249,6 +249,189 @@ Host coder-vscode.dev.coder.com--*
   })
 })
 
+it("throws an error if there is a mismatched start and end block count", async () => {
+  // The below config contains two start blocks and one end block.
+  // This is a malformed config and should throw an error.
+  // Previously were were simply taking the first occurrences of the start and
+  // end blocks, which would potentially lead to loss of any content between the
+  // missing end block and the next start block.
+  const existentSSHConfig = `Host beforeconfig
+  HostName before.config.tld
+  User before
+
+# --- START CODER VSCODE dev.coder.com ---
+Host coder-vscode.dev.coder.com--*
+  ConnectTimeout 0
+  LogLevel ERROR
+  ProxyCommand some-command-here
+  StrictHostKeyChecking no
+  UserKnownHostsFile /dev/null
+# missing END CODER VSCODE dev.coder.com ---
+
+Host donotdelete
+  HostName dont.delete.me
+  User please
+
+# --- START CODER VSCODE dev.coder.com ---
+Host coder-vscode.dev.coder.com--*
+  ConnectTimeout 0
+  LogLevel ERROR
+  ProxyCommand some-command-here
+  StrictHostKeyChecking no
+  UserKnownHostsFile /dev/null
+# --- END CODER VSCODE dev.coder.com ---
+
+Host afterconfig
+  HostName after.config.tld
+  User after`
+
+  const sshConfig = new SSHConfig(sshFilePath, mockFileSystem)
+  mockFileSystem.readFile.mockResolvedValueOnce(existentSSHConfig)
+  await sshConfig.load()
+
+  // When we try to update the config, it should throw an error.
+  await expect(
+    sshConfig.update("dev.coder.com", {
+      Host: "coder-vscode.dev.coder.com--*",
+      ProxyCommand: "some-command-here",
+      ConnectTimeout: "0",
+      StrictHostKeyChecking: "no",
+      UserKnownHostsFile: "/dev/null",
+      LogLevel: "ERROR",
+    }),
+  ).rejects.toThrow(
+    `Malformed config: ssh config has 2 dev.coder.com START comments but 1 dev.coder.com END comments. Each START must have a matching END.`,
+  )
+})
+
+it("throws an error if there are more than one sections with the same label", async () => {
+  const existentSSHConfig = `Host beforeconfig
+  HostName before.config.tld
+  User before
+
+# --- START CODER VSCODE dev.coder.com ---
+Host coder-vscode.dev.coder.com--*
+  ConnectTimeout 0
+  LogLevel ERROR
+  ProxyCommand some-command-here
+  StrictHostKeyChecking no
+  UserKnownHostsFile /dev/null
+# --- END CODER VSCODE dev.coder.com ---
+
+Host donotdelete
+  HostName dont.delete.me
+  User please
+
+# --- START CODER VSCODE dev.coder.com ---
+Host coder-vscode.dev.coder.com--*
+  ConnectTimeout 0
+  LogLevel ERROR
+  ProxyCommand some-command-here
+  StrictHostKeyChecking no
+  UserKnownHostsFile /dev/null
+# --- END CODER VSCODE dev.coder.com ---
+
+Host afterconfig
+  HostName after.config.tld
+  User after`
+
+  const sshConfig = new SSHConfig(sshFilePath, mockFileSystem)
+  mockFileSystem.readFile.mockResolvedValueOnce(existentSSHConfig)
+  await sshConfig.load()
+
+  // When we try to update the config, it should throw an error.
+  await expect(
+    sshConfig.update("dev.coder.com", {
+      Host: "coder-vscode.dev.coder.com--*",
+      ProxyCommand: "some-command-here",
+      ConnectTimeout: "0",
+      StrictHostKeyChecking: "no",
+      UserKnownHostsFile: "/dev/null",
+      LogLevel: "ERROR",
+    }),
+  ).rejects.toThrow(`Malformed config: ssh config has 2 dev.coder.com sections, please remove all but one.`)
+})
+
+it("correctly handles interspersed blocks with and without label", async () => {
+  const existentSSHConfig = `Host beforeconfig
+  HostName before.config.tld
+  User before
+
+# --- START CODER VSCODE ---
+Host coder-vscode.dev.coder.com--*
+  ConnectTimeout 0
+  LogLevel ERROR
+  ProxyCommand some-command-here
+  StrictHostKeyChecking no
+  UserKnownHostsFile /dev/null
+# --- END CODER VSCODE ---
+
+Host donotdelete
+  HostName dont.delete.me
+  User please
+
+# --- START CODER VSCODE dev.coder.com ---
+Host coder-vscode.dev.coder.com--*
+  ConnectTimeout 0
+  LogLevel ERROR
+  ProxyCommand some-command-here
+  StrictHostKeyChecking no
+  UserKnownHostsFile /dev/null
+# --- END CODER VSCODE dev.coder.com ---
+
+Host afterconfig
+  HostName after.config.tld
+  User after`
+
+  const sshConfig = new SSHConfig(sshFilePath, mockFileSystem)
+  mockFileSystem.readFile.mockResolvedValueOnce(existentSSHConfig)
+  await sshConfig.load()
+
+  const expectedOutput = `Host beforeconfig
+  HostName before.config.tld
+  User before
+
+# --- START CODER VSCODE ---
+Host coder-vscode.dev.coder.com--*
+  ConnectTimeout 0
+  LogLevel ERROR
+  ProxyCommand some-command-here
+  StrictHostKeyChecking no
+  UserKnownHostsFile /dev/null
+# --- END CODER VSCODE ---
+
+Host donotdelete
+  HostName dont.delete.me
+  User please
+
+# --- START CODER VSCODE dev.coder.com ---
+Host coder-vscode.dev.coder.com--*
+  ConnectTimeout 0
+  LogLevel ERROR
+  ProxyCommand some-command-here
+  StrictHostKeyChecking no
+  UserKnownHostsFile /dev/null
+# --- END CODER VSCODE dev.coder.com ---
+
+Host afterconfig
+  HostName after.config.tld
+  User after`
+
+  await sshConfig.update("dev.coder.com", {
+    Host: "coder-vscode.dev.coder.com--*",
+    ProxyCommand: "some-command-here",
+    ConnectTimeout: "0",
+    StrictHostKeyChecking: "no",
+    UserKnownHostsFile: "/dev/null",
+    LogLevel: "ERROR",
+  })
+
+  expect(mockFileSystem.writeFile).toBeCalledWith(sshFilePath, expectedOutput, {
+    encoding: "utf-8",
+    mode: 384,
+  })
+})
+
 it("override values", async () => {
   mockFileSystem.readFile.mockRejectedValueOnce("No file found")
   const sshConfig = new SSHConfig(sshFilePath, mockFileSystem)

src/sshConfig.ts

@@ -1,5 +1,6 @@
 import { mkdir, readFile, writeFile } from "fs/promises"
 import path from "path"
+import { countSubstring } from "./util"
 
 class SSHConfigBadFormat extends Error {}
 
@@ -123,14 +124,29 @@ export class SSHConfig {
    */
   private getBlock(label: string): Block | undefined {
     const raw = this.getRaw()
-    const startBlockIndex = raw.indexOf(this.startBlockComment(label))
-    const endBlockIndex = raw.indexOf(this.endBlockComment(label))
+    const startBlock = this.startBlockComment(label)
+    const endBlock = this.endBlockComment(label)
+    const startBlockCount = countSubstring(startBlock, raw)
+    const endBlockCount = countSubstring(endBlock, raw)
+    const startBlockIndex = raw.indexOf(startBlock)
+    const endBlockIndex = raw.indexOf(endBlock)
     const hasBlock = startBlockIndex > -1 && endBlockIndex > -1
 
     if (!hasBlock) {
       return
     }
 
+    if (startBlockCount !== endBlockCount) {
+      throw new SSHConfigBadFormat(
+        `Malformed config: ssh config has ${startBlockCount} ${label} START comments but ${endBlockCount} ${label} END comments. Each START must have a matching END.`,
+      )
+    }
+    if (startBlockCount > 1 || endBlockCount > 1) {
+      throw new SSHConfigBadFormat(
+        `Malformed config: ssh config has ${startBlockCount} ${label} sections, please remove all but one.`,
+      )
+    }
+
     if (startBlockIndex === -1) {
       throw new SSHConfigBadFormat("Start block not found")
     }
@@ -144,7 +160,7 @@ export class SSHConfig {
     }
 
     return {
-      raw: raw.substring(startBlockIndex, endBlockIndex + this.endBlockComment(label).length),
+      raw: raw.substring(startBlockIndex, endBlockIndex + endBlock.length),
     }
   }
 

src/util.ts

@@ -120,3 +120,22 @@ export function expandPath(input: string): string {
   const userHome = os.homedir()
   return input.replace(/\${userHome}/g, userHome)
 }
+
+/**
+ * Return the number of times a substring appears in a string.
+ * @param needle string
+ * @param haystack string
+ * @returns number
+ */
+export function countSubstring(needle: string, haystack: string): number {
+  if (needle.length < 1 || haystack.length < 1) {
+    return 0
+  }
+  let count = 0
+  let pos = haystack.indexOf(needle)
+  while (pos !== -1) {
+    count++
+    pos = haystack.indexOf(needle, pos + needle.length)
+  }
+  return count
+}