From 32bcf82b07150a6b4d131f1c8c08341b7b6eefcc Mon Sep 17 00:00:00 2001 From: ChomeNS <95471003+ChomeNS@users.noreply.github.com> Date: Tue, 22 Oct 2024 08:11:25 +0700 Subject: [PATCH] add worker threads and finally fix promise resolves --- index.js | 104 ++++++++++++++++++++----------------------------------- vm.js | 92 ++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 129 insertions(+), 67 deletions(-) create mode 100644 vm.js diff --git a/index.js b/index.js index 4f21e84..c0a0cfc 100644 --- a/index.js +++ b/index.js @@ -1,90 +1,60 @@ -const { VM } = require('@n8n/vm2') const { Server } = require('socket.io') -const util = require('util') -const { stylize } = require('./colors') -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 { Worker } = require('worker_threads') +const path = require('path') const io = new Server(3069) -io.on('connection', async (socket) => { - let functions - let proxy - let vm +io.on('connection', (socket) => { + let worker - const handler = { - get (target, prop) { - if (!target[prop]) throw new Error(`Function "${prop}" not available`) + let isFirst = true - return (...args) => target[prop](...args) - } - } + let jsonArray - socket.on('setFunctions', (jsonArray) => { - const parsed = JSON.parse(jsonArray) + function reset () { + worker = new Worker(path.join(__dirname, 'vm.js')) - functions = {} - - for (const eachFuntion of parsed) { - functions[eachFuntion] = (...args) => { - socket.emit(BRIDGE_PREFIX + eachFuntion, ...args) + worker.on('message', (msg) => { + switch (msg.type) { + case 'socketEmit': + socket.emit(...msg.data) - return new Promise((resolve) => { - socket.once(`functionOutput:${eachFuntion}`, (message, parseJSON) => { - if (parseJSON) resolve(JSON.parse(message)) - else resolve(message) + break + case 'socketOnce': + socket.once(msg.name, (message, parseJSON) => { + if (parseJSON) worker.postMessage({ type: 'resolvePromise', uuid: msg.uuid, data: JSON.parse(message) }) + else worker.postMessage({ type: 'resolvePromise', uuid: msg.uuid, data: message }) }) - }) - } - } - proxy = new Proxy(functions, handler) - - resetVM() - }) - - async function resetVM () { - vm = new VM({ - timeout: 1000, - sandbox: { - get bridge () { return proxy }, - randomstring, - ChatMessage, - mc, - mineflayer, - moment, - crypto, - nbt, - net, - axios + break } }) + + // ohio + if (!isFirst) { + worker.postMessage({ type: 'setFunctions', jsonArray }) + } else isFirst = false } - resetVM() + reset() - socket.on('runCode', async (transactionId, code) => { - new Promise(() => { - try { - const output = vm.run(code) + socket.on('setFunctions', (_jsonArray) => { + jsonArray = _jsonArray - socket.emit('codeOutput', transactionId, false, util.inspect(output, { stylize })) - } catch (e) { - socket.emit('codeOutput', transactionId, true, e.toString()) - } + worker.postMessage({ type: 'setFunctions', jsonArray }) + }) + + 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) => { diff --git a/vm.js b/vm.js new file mode 100644 index 0000000..0ba713f --- /dev/null +++ b/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}`) +})