add worker threads and finally fix promise resolves

index.js

@@ -1,90 +1,60 @@
-const { VM } = require('@n8n/vm2')
 const { Server } = require('socket.io')
-const util = require('util')
+const { Worker } = require('worker_threads')
-const { stylize } = require('./colors')
+const path = require('path')
-const randomstring = require('randomstring')
-const ChatMessage = require('prismarine-chat')('1.20.1')
-const mc = require('minecraft-protocol')
-const mineflayer = require('mineflayer')
-const moment = require('moment-timezone')
-const crypto = require('crypto')
-const nbt = require('prismarine-nbt')
-const net = require('net')
-const axios = require('axios')
-
-const BRIDGE_PREFIX = 'function:'
 
 const io = new Server(3069)
 
-io.on('connection', async (socket) => {
+io.on('connection', (socket) => {
-    let functions
+    let worker
-    let proxy
-    let vm
 
-    const handler = {
+    let isFirst = true
-        get (target, prop) {
-            if (!target[prop]) throw new Error(`Function "${prop}" not available`)
 
-            return (...args) => target[prop](...args)
+    let jsonArray
-        }
-    }
 
-    socket.on('setFunctions', (jsonArray) => {
+    function reset () {
-        const parsed = JSON.parse(jsonArray)
+        worker = new Worker(path.join(__dirname, 'vm.js'))
 
-        functions = {}
+        worker.on('message', (msg) => {
+            switch (msg.type) {
+                case 'socketEmit':
+                    socket.emit(...msg.data)
 
-        for (const eachFuntion of parsed) {
+                    break
-            functions[eachFuntion] = (...args) => {
+                case 'socketOnce':
-                socket.emit(BRIDGE_PREFIX + eachFuntion, ...args)
+                    socket.once(msg.name, (message, parseJSON) => {
-
+                        if (parseJSON) worker.postMessage({ type: 'resolvePromise', uuid: msg.uuid, data: JSON.parse(message) })
-                return new Promise((resolve) => {
+                        else worker.postMessage({ type: 'resolvePromise', uuid: msg.uuid, data: message })
-                    socket.once(`functionOutput:${eachFuntion}`, (message, parseJSON) => {
-                        if (parseJSON) resolve(JSON.parse(message))
-                        else resolve(message)
                     })
-                })
-            }
-        }
 
-        proxy = new Proxy(functions, handler)
+                    break
-
-        resetVM()
-    })
-
-    async function resetVM () {
-        vm = new VM({
-            timeout: 1000,
-            sandbox: {
-                get bridge () { return proxy },
-                randomstring,
-                ChatMessage,
-                mc,
-                mineflayer,
-                moment,
-                crypto,
-                nbt,
-                net,
-                axios
             }
         })
+
+        // ohio
+        if (!isFirst) {
+            worker.postMessage({ type: 'setFunctions', jsonArray })
+        } else isFirst = false
     }
 
-    resetVM()
+    reset()
 
-    socket.on('runCode', async (transactionId, code) => {
+    socket.on('setFunctions', (_jsonArray) => {
-        new Promise(() => {
+        jsonArray = _jsonArray
-            try {
-                const output = vm.run(code)
 
-                socket.emit('codeOutput', transactionId, false, util.inspect(output, { stylize }))
+        worker.postMessage({ type: 'setFunctions', jsonArray })
-            } catch (e) {
+    })
-                socket.emit('codeOutput', transactionId, true, e.toString())
+
-            }
+    socket.on('runCode', (transactionId, code) => {
+        worker.postMessage({ type: 'runCode', code })
+
+        worker.on('message', ({ type, error, output }) => {
+            if (type !== 'codeOutput') return
+
+            socket.emit('codeOutput', transactionId, error, output)
         })
     })
 
-    socket.on('reset', resetVM)
+    socket.on('reset', reset)
 })
 
 process.on('uncaughtException', (e) => {

vm.js

@@ -0,0 +1,92 @@
+const { parentPort } = require('worker_threads')
+const { VM } = require('@n8n/vm2')
+const util = require('util')
+const { stylize } = require('./colors')
+const randomstring = require('randomstring')
+const ChatMessage = require('prismarine-chat')('1.20.6')
+const mc = require('minecraft-protocol')
+const mineflayer = require('mineflayer')
+const moment = require('moment-timezone')
+const crypto = require('crypto')
+const nbt = require('prismarine-nbt')
+const net = require('net')
+const axios = require('axios')
+
+const BRIDGE_PREFIX = 'function:'
+
+let promiseResolves = {}
+
+let proxy
+
+let vm
+
+const handler = {
+    get (target, prop) {
+        if (!target[prop]) throw new Error(`Function "${prop}" not available`)
+
+        return (...args) => target[prop](...args)
+    }
+}
+
+let functions
+
+parentPort.on('message', (msg) => {
+    switch (msg.type) {
+        case 'setFunctions':
+            const parsed = JSON.parse(msg.jsonArray)
+
+            functions = {}
+
+            for (const eachFuntion of parsed) {
+                functions[eachFuntion] = (...args) => {
+                    parentPort.postMessage({ type: 'socketEmit', data: [BRIDGE_PREFIX + eachFuntion, ...args] })
+
+                    const uuid = crypto.randomUUID()
+
+                    return new Promise((resolve) => {
+                        promiseResolves[uuid] = resolve
+
+                        parentPort.postMessage({ type: 'socketOnce', uuid, name: `functionOutput:${eachFuntion}`})
+                    })
+                }
+            }
+
+            proxy = new Proxy(functions, handler)
+
+            vm = new VM({
+                timeout: 1000,
+                sandbox: {
+                    get bridge () { return proxy },
+                    randomstring,
+                    ChatMessage,
+                    mc,
+                    mineflayer,
+                    moment,
+                    crypto,
+                    nbt,
+                    net,
+                    axios
+                }
+            })
+
+            break
+        case 'resolvePromise':
+            promiseResolves[msg.uuid](msg.data)
+
+            break
+        case 'runCode':
+            try {
+                const output = vm.run(msg.code)
+
+                parentPort.postMessage({ type: 'codeOutput', output: util.inspect(output, { stylize }), error: false })
+            } catch (e) {
+                parentPort.postMessage({ type: 'codeOutput', output: e.toString(), error: true })
+            }
+
+            break
+    }
+})
+
+process.on('uncaughtException', (e) => {
+    console.log(`Caught an uncaught exception in the worker!\n${e.stack}`)
+})