implement stateless vm and update ssh package

howardchung · Jan 13, 2024 · a6fd46e · a6fd46e
1 parent cef6bc4
commit a6fd46e
Show file tree

Hide file tree

Showing 9 changed files with 369 additions and 359 deletions.
diff --git a/README.md b/README.md
@@ -66,7 +66,7 @@ This project supports creating virtual browsers (using https://github.com/m1k1o/
 - Install Docker: `curl -fsSL https://get.docker.com | sh`
 - Make sure you have an SSH key pair set up on the server (`id_rsa` in `~/.ssh` directory)
 - Add `DOCKER_VM_HOST=localhost` to your .env file (can substitute localhost for a public hostname)
-- Configure Postgres by adding `DATABASE_URL` to your .env file (Postgres is required for virtual browser management)
+- Add `ENABLE_STATELESS_VM=1` if you don't want to configure Postgres for full VM management via vmWorker and just create/delete VMs on-demand
 
 ### Room Persistence
 

diff --git a/package-lock.json b/package-lock.json
diff --git a/package.json b/package.json
@@ -23,6 +23,7 @@
     "hls.js": "^1.3.4",
     "ioredis": "^5.3.1",
     "mediasoup-client": "^3.6.82",
+    "node-ssh": "^13.1.0",
     "pg": "^8.10.0",
     "pm2": "^5.2.2",
     "react": "^18.2.0",
@@ -42,7 +43,6 @@
     "socket.io": "^4.6.1",
     "socket.io-client": "^4.6.1",
     "srt-webvtt": "^2.0.0",
-    "ssh-exec": "^2.0.0",
     "stripe": "^8.222.0",
     "twitch-m3u8": "^1.1.5",
     "uuid": "^9.0.0",

diff --git a/server/config.ts b/server/config.ts
@@ -29,6 +29,7 @@ const defaults = {
   DOCKER_VM_HOST: '', // Optional, for Docker VMs
   DOCKER_VM_HOST_SSH_USER: '', // Optional, username for Docker host
   DOCKER_VM_HOST_SSH_KEY_BASE64: '', // Optional, private SSH key for Docker host
+  ENABLE_STATELESS_VM: '', // Optional, create and delete Docker VMs without management/Postgres
   RECAPTCHA_SECRET_KEY: '', // Optional, Recaptcha for VBrowser creation
   SSL_KEY_FILE: '', // Optional, Filename of SSL key (to use https)
   SSL_CRT_FILE: '', // Optional, Filename of SSL cert (to use https)

diff --git a/server/room.ts b/server/room.ts
@@ -14,6 +14,19 @@ import { v4 as uuidv4 } from 'uuid';
 //@ts-ignore
 import twitch from 'twitch-m3u8';
 import { QueryResult } from 'pg';
+import { Docker } from './vm/docker';
+
+// Stateless pool instance to use for VMs if full management isn't needed
+let stateless: Docker | undefined = undefined;
+if (config.ENABLE_STATELESS_VM && config.DOCKER_VM_HOST) {
+  stateless = new Docker({
+    provider: 'Docker',
+    isLarge: false,
+    region: 'US',
+    limitSize: 0,
+    minSize: 0,
+  });
+}
 
 export class Room {
   // Serialized state
@@ -261,7 +274,7 @@ export class Room {
     if (postgres) {
       try {
         const roomString = this.serialize();
-        await postgres?.query(
+        await postgres.query(
           `UPDATE room SET "lastUpdateTime" = $1, data = $2 WHERE "roomId" = $3`,
           [this.lastUpdateTime, roomString, this.roomId]
         );
@@ -341,19 +354,26 @@ export class Room {
       await redis.lpush('vBrowserSessionMS', Number(new Date()) - assignTime);
       await redis.ltrim('vBrowserSessionMS', 0, 24);
     }
-    try {
-      await axios.post(
-        'http://localhost:' + config.VMWORKER_PORT + '/releaseVM',
-        {
-          provider,
-          isLarge,
-          region,
-          id,
-          uid,
+
+    if (id) {
+      try {
+        if (stateless) {
+          await stateless.terminateVM(id);
+        } else {
+          await axios.post(
+            'http://localhost:' + config.VMWORKER_PORT + '/releaseVM',
+            {
+              provider,
+              isLarge,
+              region,
+              id,
+              uid,
+            }
+          );
         }
-      );
-    } catch (e) {
-      console.warn(e);
+      } catch (e) {
+        console.warn(e);
+      }
     }
   };
 
@@ -943,16 +963,29 @@ export class Room {
     while (this.vBrowserQueue) {
       const { queueTime, isLarge, region, uid, roomId, clientId } =
         this.vBrowserQueue;
-      const assignmentResp = await axios.post<AssignedVM>(
-        'http://localhost:' + config.VMWORKER_PORT + '/assignVM',
-        {
-          isLarge,
-          region,
-          uid,
-          roomId,
+      let assignment: AssignedVM | undefined = undefined;
+      try {
+        if (stateless) {
+          const id = await stateless.startVM(uuidv4());
+          assignment = {
+            ...(await stateless.getVM(id)),
+            assignTime: Date.now(),
+          };
+        } else {
+          const { data } = await axios.post<AssignedVM>(
+            'http://localhost:' + config.VMWORKER_PORT + '/assignVM',
+            {
+              isLarge,
+              region,
+              uid,
+              roomId,
+            }
+          );
+          assignment = data;
         }
-      );
-      const assignment = assignmentResp.data;
+      } catch (e) {
+        console.warn(e);
+      }
       if (assignment) {
         this.vBrowser = assignment;
         this.vBrowser.controllerClient = clientId;

diff --git a/server/server.ts b/server/server.ts
@@ -340,8 +340,8 @@ app.get('/metadata', async (req, res) => {
         'http://localhost:' + config.VMWORKER_PORT + '/isFreePoolFull'
       )
     ).data.isFull;
-  } catch (e) {
-    console.warn(e);
+  } catch (e: any) {
+    console.warn('[WARNING]: free pool check failed: %s', e.code);
   }
   const beta =
     decoded?.email != null &&

diff --git a/server/vm/base.ts b/server/vm/base.ts
@@ -4,11 +4,14 @@ import { v4 as uuidv4 } from 'uuid';
 import { redis, redisCount } from '../utils/redis';
 import { newPostgres } from '../utils/postgres';
 import { PoolConfig, PoolRegion } from './utils';
+import type { Client } from 'pg';
 const incrInterval = 5 * 1000;
 const decrInterval = 30 * 1000;
 const cleanupInterval = 5 * 60 * 1000;
 
-const postgres = newPostgres();
+const postgres = config.ENABLE_STATELESS_VM
+  ? (null as unknown as Client)
+  : newPostgres();
 
 export abstract class VMManager {
   protected isLarge = false;
@@ -231,11 +234,11 @@ export abstract class VMManager {
     console.log('[TERMINATE]', vmid);
     // Update the DB before calling terminate
     // If we don't actually complete the termination, cleanup will eventually remove it
-    const { rowCount } = await postgres.query(
+    const { rows } = await postgres.query(
       `DELETE FROM vbrowser WHERE pool = $1 AND vmid = $2 RETURNING id`,
       [this.getPoolName(), vmid]
     );
-    console.log('DELETE', rowCount);
+    console.log('DELETE', rows.length);
     // We can log the VM lifetime by returning the creationTime and diffing
     await this.terminateVM(vmid);
   };

diff --git a/server/vm/docker.ts b/server/vm/docker.ts
@@ -5,63 +5,62 @@ import { VMManager, VM } from './base';
 import { imageName } from './utils';
 import fs from 'fs';
 import { homedir } from 'os';
-//@ts-ignore
-import sshExec from 'ssh-exec';
+import { NodeSSH } from 'node-ssh';
 
 const gatewayHost = config.DOCKER_VM_HOST;
-const sshConfig = {
-  user: config.DOCKER_VM_HOST_SSH_USER || 'root',
-  host: gatewayHost,
-  // The private key the Docker host is configured to accept
-  key: config.DOCKER_VM_HOST_SSH_KEY_BASE64
-    ? Buffer.from(config.DOCKER_VM_HOST_SSH_KEY_BASE64, 'base64')
-    : // Defaults to ~/.ssh/id_rsa on the local server
-      fs.readFileSync(homedir() + '/.ssh/id_rsa'),
-};
+let ssh: NodeSSH | undefined = undefined;
+
+async function getSSH() {
+  if (ssh && ssh.isConnected()) {
+    return ssh;
+  }
+  if (!gatewayHost) {
+    throw new Error('DOCKER_VM_HOST not defined');
+  }
+  const sshConfig = {
+    username: config.DOCKER_VM_HOST_SSH_USER || 'root',
+    host: gatewayHost,
+    // The private key the Docker host is configured to accept
+    privateKey: config.DOCKER_VM_HOST_SSH_KEY_BASE64
+      ? Buffer.from(config.DOCKER_VM_HOST_SSH_KEY_BASE64, 'base64').toString()
+      : // Defaults to ~/.ssh/id_rsa on the local server
+        fs.readFileSync(homedir() + '/.ssh/id_rsa').toString(),
+  };
+  ssh = new NodeSSH();
+  await ssh.connect(sshConfig);
+  return ssh;
+}
 
 export class Docker extends VMManager {
   size = '';
   largeSize = '';
   minRetries = 0;
   reuseVMs = false;
   id = 'Docker';
+
   startVM = async (name: string) => {
-    return new Promise<string>(async (resolve, reject) => {
-      sshExec(
-        `
-        #!/bin/bash
-        set -e
-        PORT=$(comm -23 <(seq 5000 5063 | sort) <(ss -Htan | awk '{print $4}' | cut -d':' -f2 | sort -u) | shuf | head -n 1)
-        INDEX=$(($PORT - 5000))
-        UDP_START=$((59000+$INDEX*100))
-        UDP_END=$((59099+$INDEX*100))
-        docker run -d --rm --name=${name} --memory="2g" --cpus="2" -p $PORT:$PORT -p $UDP_START-$UDP_END:$UDP_START-$UDP_END/udp -v /etc/letsencrypt:/etc/letsencrypt -l ${this.getTag()} -l index=$INDEX --log-opt max-size=1g --shm-size=1g --cap-add="SYS_ADMIN" -e NEKO_KEY="/etc/letsencrypt/live/${gatewayHost}/privkey.pem" -e NEKO_CERT="/etc/letsencrypt/live/${gatewayHost}/fullchain.pem" -e DISPLAY=":99.0" -e NEKO_PASSWORD=${name} -e NEKO_PASSWORD_ADMIN=${name} -e NEKO_BIND=":$PORT" -e NEKO_EPR=":$UDP_START-$UDP_END" -e NEKO_H264="1" ${imageName}
-        `,
-        sshConfig,
-        (err: string, stdout: string) => {
-          if (err) {
-            return reject(err);
-          }
-          console.log(stdout);
-          resolve(stdout.trim());
-        }
-      );
-    });
+    const tag = this.getTag();
+    const conn = await getSSH();
+    const { stdout, stderr } = await conn.execCommand(
+      `
+      #!/bin/bash
+      set -e
+      PORT=$(comm -23 <(seq 5000 5063 | sort) <(ss -Htan | awk '{print $4}' | cut -d':' -f2 | sort -u) | shuf | head -n 1)
+      INDEX=$(($PORT - 5000))
+      UDP_START=$((59000+$INDEX*100))
+      UDP_END=$((59099+$INDEX*100))
+      docker run -d --rm --name=${name} --memory="2g" --cpus="2" -p $PORT:$PORT -p $UDP_START-$UDP_END:$UDP_START-$UDP_END/udp -v /etc/letsencrypt:/etc/letsencrypt -l ${tag} -l index=$INDEX --log-opt max-size=1g --shm-size=1g --cap-add="SYS_ADMIN" -e NEKO_KEY="/etc/letsencrypt/live/${gatewayHost}/privkey.pem" -e NEKO_CERT="/etc/letsencrypt/live/${gatewayHost}/fullchain.pem" -e DISPLAY=":99.0" -e NEKO_PASSWORD=${name} -e NEKO_PASSWORD_ADMIN=${name} -e NEKO_BIND=":$PORT" -e NEKO_EPR=":$UDP_START-$UDP_END" -e NEKO_H264="1" ${imageName}
+      `
+    );
+    console.log(stdout, stderr);
+    return stdout.trim();
   };
 
   terminateVM = async (id: string) => {
-    return new Promise<void>((resolve, reject) => {
-      sshExec(
-        `docker rm -f ${id}`,
-        sshConfig,
-        (err: string, stdout: string) => {
-          if (err) {
-            return reject(err);
-          }
-          resolve();
-        }
-      );
-    });
+    const conn = await getSSH();
+    const { stdout, stderr } = await conn.execCommand(`docker rm -fv ${id}`);
+    console.log(stdout, stderr);
+    return;
   };
 
   rebootVM = async (id: string) => {
@@ -70,55 +69,37 @@ export class Docker extends VMManager {
   };
 
   getVM = async (id: string) => {
-    return new Promise<VM>((resolve, reject) => {
-      sshExec(
-        `docker inspect ${id}`,
-        sshConfig,
-        (err: string, stdout: string) => {
-          if (err) {
-            return reject(err);
-          }
-          let data = null;
-          try {
-            data = JSON.parse(stdout)[0];
-            if (!data) {
-              return reject(new Error('no container with this ID found'));
-            }
-          } catch {
-            console.warn(stdout);
-            return reject('failed to parse json');
-          }
-          let server = this.mapServerObject(data);
-          return resolve(server);
-        }
-      );
-    });
+    const conn = await getSSH();
+    const { stdout } = await conn.execCommand(`docker inspect ${id}`);
+    let data = null;
+    try {
+      data = JSON.parse(stdout)[0];
+      if (!data) {
+        throw new Error('no container with this ID found');
+      }
+    } catch {
+      console.warn(stdout);
+      throw new Error('failed to parse json');
+    }
+    let server = this.mapServerObject(data);
+    return server;
   };
 
   listVMs = async (filter?: string) => {
-    return new Promise<VM[]>((resolve, reject) => {
-      sshExec(
-        `docker inspect $(docker ps --filter label=${filter} --quiet --no-trunc) || true`,
-        sshConfig,
-        (err: string, stdout: string) => {
-          // Swallow exceptions and return empty array
-          if (err) {
-            return resolve([]);
-          }
-          if (!stdout) {
-            return resolve([]);
-          }
-          let data = [];
-          try {
-            data = JSON.parse(stdout);
-          } catch (e) {
-            console.warn(stdout);
-            return reject('failed to parse json');
-          }
-          return resolve(data.map(this.mapServerObject));
-        }
-      );
-    });
+    const conn = await getSSH();
+    const listCmd = `docker inspect $(docker ps --filter label=${filter} --quiet --no-trunc)`;
+    const { stdout } = await conn.execCommand(listCmd);
+    if (!stdout) {
+      return [];
+    }
+    let data = [];
+    try {
+      data = JSON.parse(stdout);
+    } catch (e) {
+      console.warn(stdout);
+      throw new Error('failed to parse json');
+    }
+    return data.map(this.mapServerObject);
   };
 
   powerOn = async (id: string) => {};

diff --git a/src/components/App/App.tsx b/src/components/App/App.tsx
@@ -512,9 +512,9 @@ export default class App extends React.Component<AppProps, AppState> {
                   const filtered = files.filter(
                     (f: WebTorrent.TorrentFile) => f.length >= 10 * 1024 * 1024
                   );
-                  const fileIndex = new URLSearchParams(src).get(
-                    'fileIndex'
-                  ) as unknown as number;
+                  const fileIndex = Number(
+                    new URLSearchParams(src).get('fileIndex')
+                  );
                   // Try to find a single large file to play
                   const target =
                     files[fileIndex] ??