From 5b30b876e53e4fbb84a96f0e8cdc4a0754e53471 Mon Sep 17 00:00:00 2001 From: Boris D Date: Thu, 9 Oct 2025 17:20:33 +0300 Subject: [PATCH] feat: enhance DatabaseManagerService and QueryExecuterService with timeout settings, add AxiosPlugin for HTTP requests, and update DatabasePlugin to use query method --- .../database/database.manager.service.ts | 2 + src/query/executer/query.executer.service.ts | 4 +- src/vm/plugins/axios.plugin.ts | 54 +++++++++++++++++++ src/vm/plugins/database.plugin.ts | 17 +++--- src/vm/vm.class.ts | 18 +++++++ src/vm/vm.constants.ts | 4 ++ tests/base/case1-payload.js | 18 +++++-- 7 files changed, 104 insertions(+), 13 deletions(-) create mode 100644 src/vm/plugins/axios.plugin.ts diff --git a/src/databaseManager/database/database.manager.service.ts b/src/databaseManager/database/database.manager.service.ts index 5f7239b..d5509d7 100644 --- a/src/databaseManager/database/database.manager.service.ts +++ b/src/databaseManager/database/database.manager.service.ts @@ -79,6 +79,8 @@ export class DatabaseManagerService extends DatabaseEncryptionService { user: queryUser ? database.q_username : database.c_username, password: this.decryptPassword(database.password), database: database.database, + idleTimeout: 150e3, + connectTimeout: 2e3, }; } diff --git a/src/query/executer/query.executer.service.ts b/src/query/executer/query.executer.service.ts index 36d2bfb..c708e8d 100644 --- a/src/query/executer/query.executer.service.ts +++ b/src/query/executer/query.executer.service.ts @@ -61,7 +61,7 @@ export class QueryExecuterService { { removeOnComplete: true, removeOnFail: true, - attempts: 3, + attempts: 0, } ); @@ -130,6 +130,8 @@ export class QueryExecuterService { const vm = new Vm({ memoryLimit: 128, + timeLimit: BigInt(100e9), + cpuTimeLimit: BigInt(5e9), modules: modules, plugins: plugins, }); diff --git a/src/vm/plugins/axios.plugin.ts b/src/vm/plugins/axios.plugin.ts new file mode 100644 index 0000000..78ec54a --- /dev/null +++ b/src/vm/plugins/axios.plugin.ts @@ -0,0 +1,54 @@ +import axios from "axios"; +import { Plugin } from "../plugin.class"; + +export class AxiosPlugin extends Plugin { + constructor(name: string) { + super(name, ["get", "post", "put", "delete", "head", "options", "patch"]); + } + + static init(name: string): AxiosPlugin { + return new AxiosPlugin(name); + } + + private configure(config) { + if (config.timeout > 10e3) { + config.timeout = 10e3; + } else { + config.timeout = 5e3; + } + + return config; + } + + async get(url: string, config?: any): Promise { + return await axios.get(url, this.configure(config)); + } + + async post(url: string, data?: any, config?: any): Promise { + return await axios.post(url, data, this.configure(config)); + } + + async put(url: string, data?: any, config?: any): Promise { + return await axios.put(url, data, this.configure(config)); + } + + async delete(url: string, config?: any): Promise { + return await axios.delete(url, this.configure(config)); + } + + async head(url: string, config?: any): Promise { + return await axios.head(url, this.configure(config)); + } + + async options(url: string, config?: any): Promise { + return await axios.options(url, this.configure(config)); + } + + async patch(url: string, data?: any, config?: any): Promise { + return await axios.patch(url, data, this.configure(config)); + } + + onFinish() { + // No resources to clean up + } +} diff --git a/src/vm/plugins/database.plugin.ts b/src/vm/plugins/database.plugin.ts index 8b7317f..1589d51 100644 --- a/src/vm/plugins/database.plugin.ts +++ b/src/vm/plugins/database.plugin.ts @@ -9,7 +9,7 @@ export class DatabasePlugin extends Plugin { ); } - super(name, ["execute"]); + super(name, ["query"]); } static async init( @@ -20,6 +20,8 @@ export class DatabasePlugin extends Plugin { user: string; password: string; database: string; + idleTimeout: number; + connectTimeout: number; } ): Promise { const dbConnection = await mysql.createConnection({ @@ -28,19 +30,18 @@ export class DatabasePlugin extends Plugin { port: config.port, password: config.password, database: config.database, + idleTimeout: config.idleTimeout, + connectTimeout: config.connectTimeout, enableKeepAlive: true, }); + await dbConnection.query("SET SESSION MAX_EXECUTION_TIME=2000;"); + return new DatabasePlugin(name, dbConnection); } - async execute(query): Promise { - const [rows, fields] = await this.dbConnection.query(query); - - return { - rows: rows, - fields: fields ?? [], - }; + async query(query): Promise { + return await this.dbConnection.query(query); } onFinish() { diff --git a/src/vm/vm.class.ts b/src/vm/vm.class.ts index e5843f2..f8a7453 100644 --- a/src/vm/vm.class.ts +++ b/src/vm/vm.class.ts @@ -10,15 +10,21 @@ export class Vm { private jail: any; private plugins: Plugin[]; private isolate: ivm.Isolate; + private timeLimit?: bigint; + private cpuTimeLimit?: bigint; constructor(configs: { memoryLimit: number; + timeLimit?: bigint; + cpuTimeLimit?: bigint; modules: VModule[]; plugins: Plugin[]; }) { this.memoryLimit = configs.memoryLimit; this.modules = configs.modules; this.plugins = configs.plugins; + this.timeLimit = configs.timeLimit; + this.cpuTimeLimit = configs.cpuTimeLimit; } async init(): Promise { @@ -111,11 +117,23 @@ export class Vm { const compiledScript = await this.isolate.compileScript(scriptWithResult); + const interval = setInterval(() => { + if ( + this.isolate.cpuTime > this.cpuTimeLimit || + this.isolate.wallTime > this.timeLimit + ) { + this.isolate.dispose(); + + rejectPromise(new Error("Script execution timed out")); + } + }, 500); + compiledScript.run(this.context); try { return await resultPromise; } finally { + clearInterval(interval); this.onFinish(); } } diff --git a/src/vm/vm.constants.ts b/src/vm/vm.constants.ts index 40ed2e2..f09a715 100644 --- a/src/vm/vm.constants.ts +++ b/src/vm/vm.constants.ts @@ -2,6 +2,7 @@ import { QueryExecuterService } from "src/query/executer/query.executer.service" import { DatabasePlugin } from "./plugins/database.plugin"; import { Query } from "src/query/entities/query.entity"; import { QueryPlugin } from "./plugins/query.plugin"; +import { AxiosPlugin } from "./plugins/axios.plugin"; export const registeredPlugins = { db: async (service: QueryExecuterService, query: Query) => { @@ -17,6 +18,9 @@ export const registeredPlugins = { return DatabasePlugin.init("db", databaseConnection); }, + axios: async () => { + return AxiosPlugin.init("axios"); + }, query: async (service: QueryExecuterService, query: Query) => { return QueryPlugin.init(query, service); }, diff --git a/tests/base/case1-payload.js b/tests/base/case1-payload.js index 3eb2860..02514ed 100644 --- a/tests/base/case1-payload.js +++ b/tests/base/case1-payload.js @@ -3,6 +3,7 @@ import "module/squel"; import "plugin/db"; +import "plugin/axios"; function createSQL(id) { return squel.select().from("test").where("id = ?", id).toString(); @@ -11,13 +12,22 @@ function createSQL(id) { async function main(input, headers) { const sql = createSQL(input.id); - await db.execute("START TRANSACTION"); + await db.query("START TRANSACTION"); - // log(await db.execute('insert into test (name) values ("Test")')); + // log(await db.query('insert into test (name) values ("Test")')); - const res = await db.execute(sql); + log(new Date().toISOString()); + // const a = await axios.get("https://httpbin.dev/delay/10", { + // timeout: 50000000, + // }); - log(res); + // log(a); + + const res = await db.query(` + SELECT SLEEP(10000); +`); + + log(new Date().toISOString()); return { response: {