From 31301419ad344a2e81b398eb9567616fed008a72 Mon Sep 17 00:00:00 2001
From: Justin Lin <wancat@wancat.cc>
Date: Mon, 14 Oct 2024 17:31:35 +0800
Subject: [PATCH] feat: add ssh connection retry

---
 .env.sample                     |  2 +-
 server/src/config.ts            |  4 +--
 server/src/containers/docker.ts | 54 ++++++++-------------------------
 server/src/containers/ssh.ts    | 48 +++++++++++++++++++++++++++++
 4 files changed, 64 insertions(+), 44 deletions(-)

diff --git a/.env.sample b/.env.sample
index a389a42..06be919 100644
--- a/.env.sample
+++ b/.env.sample
@@ -1,11 +1,11 @@
 HOST=localhost
 PORT=3000
 BASE_URL=http://localhost:5173
-BACKEND_URL=http://localhost:3000
 MONGO_URL=mongodb://localhost:27017/odyssey-test
 SECRET_KEY=sample-secret
 
 # Optional
+# BACKEND_URL=http://host.docker.internal:3000
 # CONTAINER_EXPIRY=60     # 60 minutes
 # SESSION_MAX_AGE=168      # 7 days
 # TRUSTED_PROXIES=1.2.3.4    
diff --git a/server/src/config.ts b/server/src/config.ts
index a16399d..ea258af 100644
--- a/server/src/config.ts
+++ b/server/src/config.ts
@@ -35,6 +35,7 @@ function createConfig() {
   const host = get('HOST', '0.0.0.0')
   const port = Number(get('PORT', 3000))
   const baseUrl = get('BASE_URL', `http://${host}:${port}`)
+  const backendUrl = get('BACKEND_URL', `http://host.docker.internal:${port}`)
   const url = new URL(baseUrl)
   const projectRoot = getProjectRoot()
   const keypairPath = get(
@@ -49,7 +50,7 @@ function createConfig() {
     baseUrl,
     domain: url.hostname,
     protocol: url.protocol,
-    backendUrl: get('BACKEND_URL', baseUrl),
+    backendUrl,
     db: get('MONGO_URL', 'mongodb://localhost:27017/odyssey-test'),
     secret: get('SECRET_KEY', ''),
     isProduction: process.env.NODE_ENV === 'production',
@@ -97,7 +98,6 @@ function getProjectRoot(): string {
 }
 
 const config = createConfig()
-console.log(config.docker.keypair)
 
 export function getQuestImage(id: string): string {
   return `${config.docker.imagePrefix}${id}`
diff --git a/server/src/containers/docker.ts b/server/src/containers/docker.ts
index 848a3fd..92acd07 100644
--- a/server/src/containers/docker.ts
+++ b/server/src/containers/docker.ts
@@ -1,9 +1,9 @@
 import fs from 'fs/promises'
 import { Duplex } from 'stream'
-import { Client } from 'ssh2'
 import Docker from 'dockerode'
 import config, { getQuestImage } from '../config.js'
 import logger from '../utils/logger.js'
+import { connectToSSH } from './ssh.js'
 
 const engine = new Docker()
 
@@ -70,9 +70,6 @@ export async function getAndStartContainer(
   if (!(await container.inspect()).State.Running) {
     await container.start()
   }
-  await new Promise((resolve) => {
-    setTimeout(resolve, 2000)
-  })
   return container
 }
 
@@ -83,45 +80,20 @@ export async function attachContainer(
   const containerIp = (await container.inspect()).NetworkSettings.Networks[
     config.docker.network
   ].IPAddress
-  const conn = new Client()
-  console.log(config.docker.keypair.privateKey)
 
-  return new Promise((resolve, reject) => {
-    conn
-      .on('ready', () => {
-        console.log('SSH connection ready')
-        conn.shell(
-          {
-            env: {
-              TOKEN: token,
-              API_ENDPOINT: 'http://host.docker.internal:3000',
-              ZDOTDIR: '/etc/zsh',
-              NODE_PATH: '/usr/src/node_modules',
-            },
-          },
-          (err, stream) => {
-            if (err) {
-              reject(err)
-            }
-            stream.on('close', () => {
-              console.log('Stream :: close')
-              conn.end()
-            })
-            resolve(stream)
-          }
-        )
+  for (let i = 0; i < 10; i++) {
+    try {
+      // eslint-disable-next-line no-await-in-loop
+      const stream = await connectToSSH(containerIp, { token })
+      return stream
+    } catch {
+      // eslint-disable-next-line no-await-in-loop
+      await new Promise((resolve) => {
+        setTimeout(resolve, 1000)
       })
-      .on('error', (err) => {
-        reject(err)
-      })
-      .connect({
-        host: containerIp,
-        port: 22,
-        username: 'commander',
-        privateKey: config.docker.keypair.privateKey,
-        debug: console.log,
-      })
-  })
+    }
+  }
+  throw new Error('Failed to connect to SSH')
 }
 
 export async function deleteContainer(id: string) {
diff --git a/server/src/containers/ssh.ts b/server/src/containers/ssh.ts
index e69de29..ffb1338 100644
--- a/server/src/containers/ssh.ts
+++ b/server/src/containers/ssh.ts
@@ -0,0 +1,48 @@
+import { Client } from 'ssh2'
+import { Duplex } from 'stream'
+import config from '../config.js'
+
+export async function connectToSSH(
+  containerIp: string,
+  { token }: { token: string }
+): Promise<Duplex> {
+  const conn = new Client()
+
+  return new Promise((resolve, reject) => {
+    conn
+      .on('ready', () => {
+        console.log('SSH connection ready')
+        conn.shell(
+          {
+            env: {
+              TOKEN: token,
+              API_ENDPOINT: config.backendUrl,
+              ZDOTDIR: '/etc/zsh',
+              NODE_PATH: '/usr/src/node_modules',
+            },
+          },
+          (err, stream) => {
+            if (err) {
+              reject(err)
+            }
+            stream.on('close', () => {
+              console.log('Stream :: close')
+              conn.end()
+            })
+            resolve(stream)
+          }
+        )
+      })
+      .on('error', (err) => {
+        reject(err)
+      })
+      .connect({
+        host: containerIp,
+        port: 22,
+        username: 'commander',
+        privateKey: config.docker.keypair.privateKey,
+        debug: console.log,
+        timeout: 5000,
+      })
+  })
+}