Use the SSH client directly

coder · kylecarbs · Jul 27, 2023 · Jul 27, 2023 · Jul 27, 2023 · Jul 27, 2023
commit f2db72a7a957d52b8152f97a4a50acf25bc8da2d
diff --git a/site/e2e/helpers.ts b/site/e2e/helpers.ts
@@ -1,4 +1,4 @@
-import { expect, Page, TestInfo } from "@playwright/test"
+import { expect, Page } from "@playwright/test"
 import { spawn } from "child_process"
 import { randomUUID } from "crypto"
 import path from "path"
@@ -15,6 +15,8 @@ import {
   Resource,
 } from "./provisionerGenerated"
 import { port } from "./playwright.config"
+import * as ssh from "ssh2"
+import { Duplex } from "stream"
 
 // createWorkspace creates a workspace for a template.
 // It does not wait for it to be running, but it does navigate to the page.
@@ -59,22 +61,53 @@ export const createTemplate = async (
   return name
 }
 
+// sshIntoWorkspace spawns a Coder SSH process and a client connected to it.
+export const sshIntoWorkspace = async (page: Page, workspace: string): Promise<ssh.Client> => {
+  const sessionToken = await findSessionToken(page)
+  return new Promise<ssh.Client>((resolve, reject) => {
+    const cp = spawn("go", ["run", coderMainPath(), "ssh", "--stdio", workspace], {
+      env: {
+        ...process.env,
+        CODER_SESSION_TOKEN: sessionToken,
+        CODER_URL: "http://localhost:3000",
+      },
+    })
+    cp.on("error", (err) => reject(err))
+    const proxyStream = new Duplex({
+      read: (size) => {
+        return cp.stdout.read(Math.min(size, cp.stdout.readableLength))
+      },
+      write: cp.stdin.write.bind(cp.stdin),
+    })
+    // eslint-disable-next-line no-console -- Helpful for debugging
+    cp.stderr.on("data", (data) => console.log(data.toString()))
+    cp.stdout.on("readable", (...args) => {
+      proxyStream.emit('readable', ...args);
+      if (cp.stdout.readableLength > 0) {
+        proxyStream.emit("data", cp.stdout.read());
+      }
+    });
+    const client = new ssh.Client()
+    client.connect({
+      sock: proxyStream,
+      username: "coder",
+    })
+    client.on("error", (err) => reject(err))
+    client.on("ready", () => {
+      resolve(client)
+    })
+  })
+}
+
 // startAgent runs the coder agent with the provided token.
 // It awaits the agent to be ready before returning.
 export const startAgent = async (page: Page, token: string): Promise<void> => {
-  const coderMain = path.join(
-    __dirname,
-    "..",
-    "..",
-    "enterprise",
-    "cmd",
-    "coder",
-    "main.go",
-  )
-  return startAgentWithCommand(page, token, "go", "run", coderMain)
+  return startAgentWithCommand(page, token, "go", "run", coderMainPath())
 }
 
-export const downloadCoderVersion = async (testInfo: TestInfo, version: string): Promise<string> => {
+// downloadCoderVersion downloads the version provided into a temporary dir and
+// caches it so subsequent calls are fast.
+export const downloadCoderVersion = async (version: string): Promise<string> => {
   if (version.startsWith("v")) {
     version = version.slice(1)
   }
@@ -95,6 +128,8 @@ export const downloadCoderVersion = async (testInfo: TestInfo, version: string):
     return binaryPath
   }
 
+  // Runs our public install script using our options to
+  // install the binary!
   await new Promise<void>((resolve, reject) => {
     const cp = spawn("sh", ["-c", [
       "curl", "-L", "https://coder.com/install.sh",
@@ -105,8 +140,8 @@ export const downloadCoderVersion = async (testInfo: TestInfo, version: string):
       "--prefix", tempDir,
       "--binary-name", binaryName,
     ].join(" ")])
-    cp.stderr.on("data", (data) => testInfo.stderr.push(data))
-    cp.stdout.on("data", (data) => testInfo.stdout.push(data))
+    // eslint-disable-next-line no-console -- Needed for debugging
+    cp.stderr.on("data", (data) => console.log(data.toString()))
     cp.on("close", (code) => {
       if (code === 0) {
         resolve()
@@ -138,6 +173,18 @@ export const startAgentWithCommand = async (page: Page, token: string, command:
   }
 }
 
+const coderMainPath = (): string => {
+  return path.join(
+    __dirname,
+    "..",
+    "..",
+    "enterprise",
+    "cmd",
+    "coder",
+    "main.go",
+  )
+}
+
 // Allows users to more easily define properties they want for agents and resources!
 type RecursivePartial<T> = {
   [P in keyof T]?: T[P] extends (infer U)[]
@@ -308,7 +355,7 @@ export const createServer = async (
   return e
 }
 
-export const findSessionToken = async (page: Page): Promise<string> => {
+const findSessionToken = async (page: Page): Promise<string> => {
   const cookies = await page.context().cookies()
   const sessionCookie = cookies.find((c) => c.name === "coder_session_token")
   if (!sessionCookie) {

diff --git a/site/e2e/tests/outdatedAgent.spec.ts b/site/e2e/tests/outdatedAgent.spec.ts
@@ -1,10 +1,10 @@
 import { test } from "@playwright/test"
-import { createTemplate, createWorkspace, downloadCoderVersion, findSessionToken, startAgentWithCommand } from "../helpers"
 import { randomUUID } from "crypto"
-import { spawn } from "child_process"
-import path from "path"
+import { createTemplate, createWorkspace, downloadCoderVersion, sshIntoWorkspace, startAgentWithCommand } from "../helpers"
 
-test("create workspace with an outdated agent", async ({ page }, testInfo) => {
+const agentVersion = "v0.14.0"
+
+test("ssh with agent " + agentVersion, async ({ page }) => {
   const token = randomUUID()
   const template = await createTemplate(page, {
     apply: [
@@ -24,42 +24,23 @@ test("create workspace with an outdated agent", async ({ page }, testInfo) => {
     ],
   })
   const workspace = await createWorkspace(page, template)
-  const binaryPath = await downloadCoderVersion(testInfo, "v0.24.0")
+  const binaryPath = await downloadCoderVersion(agentVersion)
   await startAgentWithCommand(page, token, binaryPath)
-  const sessionToken = await findSessionToken(page)
-  const coderMain = path.join(
-    __dirname,
-    "..",
-    "..",
-    "..",
-    "enterprise",
-    "cmd",
-    "coder",
-    "main.go",
-  )
+
+  const client = await sshIntoWorkspace(page, workspace)
   await new Promise<void>((resolve, reject) => {
-    const cp = spawn("ssh", [
-      "-o", "StrictHostKeyChecking=no",
-      "-o", "UserKnownHostsFile=/dev/null",
-      "-o", "ProxyCommand=/usr/local/go/bin/go run "+coderMain+" ssh --stdio " + workspace,
-      "localhost",
-      "exit",
-      "0",
-    ], {
-      env: {
-        ...process.env,
-        CODER_SESSION_TOKEN: sessionToken,
-        CODER_URL: "http://localhost:3000",
-      },
-    })
-    cp.stderr.on("data", (data) => console.log(data.toString()))
-    cp.stdout.on("data", (data) => console.log(data.toString()))
-    cp.on("close", (code) => {
-      if (code === 0) {
-        resolve()
-      } else {
-        reject(new Error("ssh failed with code " + code))
+    // We just exec a command to be certain the agent is running!
+    client.exec("exit 0", (err, stream) => {
+      if (err) {
+        return reject(err)
       }
+      stream.on("exit", (code) => {
+        if (code !== 0) {
+          return reject(new Error(`Command exited with code ${code}`))
+        }
+        client.end();
+        resolve()
+      });
     })
   })
 })
diff --git a/site/package.json b/site/package.json
@@ -126,6 +126,7 @@
     "@types/react-syntax-highlighter": "15.5.5",
     "@types/react-virtualized-auto-sizer": "1.0.1",
     "@types/react-window": "1.8.5",
+    "@types/ssh2": "1.11.13",
     "@types/ua-parser-js": "0.7.36",
     "@types/uuid": "9.0.2",
     "@typescript-eslint/eslint-plugin": "5.62.0",
@@ -153,6 +154,7 @@
     "msw": "1.2.2",
     "prettier": "3.0.0",
     "resize-observer": "1.0.4",
+    "ssh2": "1.14.0",
     "storybook": "7.1.0",
     "storybook-addon-react-router-v6": "1.0.2",
     "storybook-react-context": "0.6.0",

diff --git a/site/yarn.lock b/site/yarn.lock
@@ -3751,6 +3751,11 @@
   resolved "https://registry.yarnpkg.com/@types/node/-/node-16.18.23.tgz#b6e934fe427eb7081d0015aad070acb3373c3c90"
   integrity sha512-XAMpaw1s1+6zM+jn2tmw8MyaRDIJfXxqmIQIS0HfoGYPuf7dUWeiUKopwq13KFX9lEp1+THGtlaaYx39Nxr58g==
 
+"@types/node@^18.11.18":
+  version "18.17.1"
+  resolved "https://registry.yarnpkg.com/@types/node/-/node-18.17.1.tgz#84c32903bf3a09f7878c391d31ff08f6fe7d8335"
+  integrity sha512-xlR1jahfizdplZYRU59JlUx9uzF1ARa8jbhM11ccpCJya8kvos5jwdm2ZAgxSCwOl0fq21svP18EVwPBXMQudw==
+
 "@types/normalize-package-data@^2.4.0":
   version "2.4.1"
   resolved "https://registry.yarnpkg.com/@types/normalize-package-data/-/normalize-package-data-2.4.1.tgz#d3357479a0fdfdd5907fe67e17e0a85c906e1301"
@@ -3905,6 +3910,13 @@
   dependencies:
     "@types/node" "*"
 
+"@types/ssh2@1.11.13":
+  version "1.11.13"
+  resolved "https://registry.yarnpkg.com/@types/ssh2/-/ssh2-1.11.13.tgz#e6224da936abec0541bf26aa826b1cc37ea70d69"
+  integrity sha512-08WbG68HvQ2YVi74n2iSUnYHYpUdFc/s2IsI0BHBdJwaqYJpWlVv9elL0tYShTv60yr0ObdxJR5NrCRiGJ/0CQ==
+  dependencies:
+    "@types/node" "^18.11.18"
+
 "@types/stack-utils@^2.0.0":
   version "2.0.1"
   resolved "https://registry.yarnpkg.com/@types/stack-utils/-/stack-utils-2.0.1.tgz#20f18294f797f2209b5f65c8e3b5c8e8261d127c"
@@ -4374,6 +4386,13 @@ array.prototype.flatmap@^1.3.0, array.prototype.flatmap@^1.3.1:
     es-abstract "^1.20.4"
     es-shim-unscopables "^1.0.0"
 
+asn1@^0.2.6:
+  version "0.2.6"
+  resolved "https://registry.yarnpkg.com/asn1/-/asn1-0.2.6.tgz#0d3a7bb6e64e02a90c0303b31f292868ea09a08d"
+  integrity sha512-ix/FxPn0MDjeyJ7i/yoHGFt/EX6LyNbxSEhPPXODPL+KB0VPk86UYfL0lMdy+KCnv+fmvIzySwaK5COwqVbWTQ==
+  dependencies:
+    safer-buffer "~2.1.0"
+
 assert@^2.0.0:
   version "2.0.0"
   resolved "https://registry.yarnpkg.com/assert/-/assert-2.0.0.tgz#95fc1c616d48713510680f2eaf2d10dd22e02d32"
@@ -4571,6 +4590,13 @@ base64-js@^1.3.1:
   resolved "https://registry.yarnpkg.com/base64-js/-/base64-js-1.5.1.tgz#1b1b440160a5bf7ad40b650f095963481903930a"
   integrity sha512-AKpaYlHn8t4SVbOHCy+b5+KKgvR4vrsD8vbvrbiQJps7fKDTkjkDry6ji0rUJjC0kzbNePLwzxq8iypo41qeWA==
 
+bcrypt-pbkdf@^1.0.2:
+  version "1.0.2"
+  resolved "https://registry.yarnpkg.com/bcrypt-pbkdf/-/bcrypt-pbkdf-1.0.2.tgz#a4301d389b6a43f9b67ff3ca11a3f6637e360e9e"
+  integrity sha512-qeFIXtP4MSoi6NLqO12WfqARWWuCKi2Rn/9hJLEmtB5yTNr9DqFWkJRCf2qShWzPeAMRnOgCrq0sg/KLv5ES9w==
+  dependencies:
+    tweetnacl "^0.14.3"
+
 better-opn@^3.0.2:
   version "3.0.2"
   resolved "https://registry.yarnpkg.com/better-opn/-/better-opn-3.0.2.tgz#f96f35deaaf8f34144a4102651babcf00d1d8817"
@@ -4701,6 +4727,11 @@ buffer@^5.5.0:
     base64-js "^1.3.1"
     ieee754 "^1.1.13"
 
+buildcheck@~0.0.6:
+  version "0.0.6"
+  resolved "https://registry.yarnpkg.com/buildcheck/-/buildcheck-0.0.6.tgz#89aa6e417cfd1e2196e3f8fe915eb709d2fe4238"
+  integrity sha512-8f9ZJCUXyT1M35Jx7MkBgmBMo3oHTTBIPLiY9xyL0pl3T5RwcPEY8cUHr5LBNfu/fk6c2T4DJZuVM/8ZZT2D2A==
+
 builtin-modules@^3.3.0:
   version "3.3.0"
   resolved "https://registry.yarnpkg.com/builtin-modules/-/builtin-modules-3.3.0.tgz#cae62812b89801e9656336e46223e030386be7b6"
@@ -5177,6 +5208,14 @@ cosmiconfig@^7.0.0, cosmiconfig@^7.0.1:
     path-type "^4.0.0"
     yaml "^1.10.0"
 
+cpu-features@~0.0.8:
+  version "0.0.8"
+  resolved "https://registry.yarnpkg.com/cpu-features/-/cpu-features-0.0.8.tgz#a2d464b023b8ad09004c8cdca23b33f192f63546"
+  integrity sha512-BbHBvtYhUhksqTjr6bhNOjGgMnhwhGTQmOoZGD+K7BCaQDCuZl/Ve1ZxUSMRwVC4D/rkCPQ2MAIeYzrWyK7eEg==
+  dependencies:
+    buildcheck "~0.0.6"
+    nan "^2.17.0"
+
 create-jest-runner@^0.11.2:
   version "0.11.2"
   resolved "https://registry.yarnpkg.com/create-jest-runner/-/create-jest-runner-0.11.2.tgz#4b4f62ccef1e4de12e80f81c2cf8211fa392a962"
@@ -10511,7 +10550,7 @@ safe-regex@^2.1.1:
   dependencies:
     regexp-tree "~0.1.1"
 
-"safer-buffer@>= 2.1.2 < 3", "safer-buffer@>= 2.1.2 < 3.0.0":
+"safer-buffer@>= 2.1.2 < 3", "safer-buffer@>= 2.1.2 < 3.0.0", safer-buffer@~2.1.0:
   version "2.1.2"
   resolved "https://registry.yarnpkg.com/safer-buffer/-/safer-buffer-2.1.2.tgz#44fa161b0187b9549dd84bb91802f9bd8385cd6a"
   integrity sha512-YZo3K82SD7Riyi0E1EQPojLz7kpepnSQI9IyPbHHg1XXXevb5dJI7tpyN2ADxGcQbHG7vcyRHk0cbwqcQriUtg==
@@ -10775,6 +10814,17 @@ sprintf-js@~1.0.2:
   resolved "https://registry.yarnpkg.com/sprintf-js/-/sprintf-js-1.0.3.tgz#04e6926f662895354f3dd015203633b857297e2c"
   integrity sha512-D9cPgkvLlV3t3IzL0D0YLvGA9Ahk4PcvVwUbN0dSGr1aP0Nrt4AEnTUbuGvquEC0mA64Gqt1fzirlRs5ibXx8g==
 
+ssh2@1.14.0:
+  version "1.14.0"
+  resolved "https://registry.yarnpkg.com/ssh2/-/ssh2-1.14.0.tgz#8f68440e1b768b66942c9e4e4620b2725b3555bb"
+  integrity sha512-AqzD1UCqit8tbOKoj6ztDDi1ffJZ2rV2SwlgrVVrHPkV5vWqGJOVp5pmtj18PunkPJAuKQsnInyKV+/Nb2bUnA==
+  dependencies:
+    asn1 "^0.2.6"
+    bcrypt-pbkdf "^1.0.2"
+  optionalDependencies:
+    cpu-features "~0.0.8"
+    nan "^2.17.0"
+
 stack-generator@^2.0.5:
   version "2.0.10"
   resolved "https://registry.yarnpkg.com/stack-generator/-/stack-generator-2.0.10.tgz#8ae171e985ed62287d4f1ed55a1633b3fb53bb4d"
@@ -11362,6 +11412,11 @@ tween-functions@^1.2.0:
   resolved "https://registry.yarnpkg.com/tween-functions/-/tween-functions-1.2.0.tgz#1ae3a50e7c60bb3def774eac707acbca73bbc3ff"
   integrity sha512-PZBtLYcCLtEcjL14Fzb1gSxPBeL7nWvGhO5ZFPGqziCcr8uvHp0NDmdjBchp6KHL+tExcg0m3NISmKxhU394dA==
 
+tweetnacl@^0.14.3:
+  version "0.14.5"
+  resolved "https://registry.yarnpkg.com/tweetnacl/-/tweetnacl-0.14.5.tgz#5ae68177f192d4456269d108afa93ff8743f4f64"
+  integrity sha512-KXXFFdAbFXY4geFIwoyNK+f5Z1b7swfXABfL7HXCmoIWMKU3dmS26672A4EeQtDzLKy7SXmfBu51JolvEKwtGA==
+
 type-check@^0.4.0, type-check@~0.4.0:
   version "0.4.0"
   resolved "https://registry.yarnpkg.com/type-check/-/type-check-0.4.0.tgz#07b8203bfa7056c0657050e3ccd2c37730bab8f1"