diff --git a/src/api/entities/token.entity.ts b/src/api/entities/token.entity.ts index 0a018c4..77b8c40 100644 --- a/src/api/entities/token.entity.ts +++ b/src/api/entities/token.entity.ts @@ -9,8 +9,8 @@ export class Token { @Column({ type: "tinyint", default: 1 }) isActive: boolean; - @Column({ type: "tinyint", default: 0 }) - isAdmin: boolean; + // @Column({ type: "tinyint", default: 0 }) + // isAdmin: boolean; @ManyToOne(() => Project, (project) => project.apiTokens) project: Project; diff --git a/src/api/guards/admin.guard.ts b/src/api/guards/admin.guard.ts index 6d419f3..1291e01 100644 --- a/src/api/guards/admin.guard.ts +++ b/src/api/guards/admin.guard.ts @@ -15,13 +15,15 @@ export class AdminGuard implements CanActivate { ) {} async canActivate(context: ExecutionContext): Promise { - const request = context.switchToHttp().getRequest(); - const apiToken = request.apiToken; - - if (!apiToken || !apiToken.isAdmin) { - throw new UnauthorizedException("Admin privileges are required"); - } - return true; + + // const request = context.switchToHttp().getRequest(); + // const apiToken = request.apiToken; + + // if (!apiToken || !apiToken.isAdmin) { + // throw new UnauthorizedException("Admin privileges are required"); + // } + + // return true; } } diff --git a/src/query/base/base-query.controller.ts b/src/query/base/base-query.controller.ts index b9eec0c..dbcfcf9 100644 --- a/src/query/base/base-query.controller.ts +++ b/src/query/base/base-query.controller.ts @@ -5,10 +5,11 @@ import { Inject, Param, Post, + Req, Res, UseGuards, } from "@nestjs/common"; -import { Response } from "express"; +import { Response, Request } from "express"; import { QueryHandlerService } from "../handler/query.handler.service"; import { ApiTokenGuard } from "src/api/guards/api-token.guard"; import { QueryExecuterService } from "../executer/query.executer.service"; @@ -52,11 +53,22 @@ export abstract class BaseQueryController { const queryResult = await this.queryExecuterService.runQueryQueued( id, query, - headers + headers, + headers.cookie.split("; ").reduce((acc, cookie) => { + const [key, value] = cookie.split("="); + acc[key] = value; + return acc; + }, {}) ); res.status(queryResult?.statusCode || 200); + if (queryResult?.cookies) { + for (const [key, value] of Object.entries(queryResult?.cookies || {})) { + res.cookie(key, value); + } + } + if ( queryResult?.redirect || (queryResult?.statusCode === 302 && diff --git a/src/query/executer/query.executer.service.ts b/src/query/executer/query.executer.service.ts index 349914e..1c6b43e 100644 --- a/src/query/executer/query.executer.service.ts +++ b/src/query/executer/query.executer.service.ts @@ -15,6 +15,7 @@ import { InjectQueue } from "@nestjs/bullmq"; import { QUEUE_NAMES } from "src/queue/constants"; import { Queue, QueueEvents } from "bullmq"; import { FunctionService } from "src/query/function/function.service"; +import { SessionService } from "../session/session.service"; @Injectable() export class QueryExecuterService { @@ -27,6 +28,7 @@ export class QueryExecuterService { private readonly functionService: FunctionService, readonly databaseManagerService: DatabaseManagerService, readonly redisNodeService: RedisNodeService, + readonly sessionService: SessionService, @InjectQueue(QUEUE_NAMES.QUERY) private queryQueue: Queue ) { this.queueEvents = new QueueEvents(this.queryQueue.name); @@ -54,7 +56,8 @@ export class QueryExecuterService { async runQueryQueued( token: string, queryData: any, - headers: Record = {} + headers: Record = {}, + cookies: Record = {} ): Promise { const job = await this.queryQueue.add( `${new Date().getTime()}_${token}`, @@ -62,6 +65,7 @@ export class QueryExecuterService { token, queryData, headers, + cookies, }, { removeOnComplete: true, @@ -77,7 +81,8 @@ export class QueryExecuterService { async runQuery( token: string, queryData: any, - headers: Record = {} + headers: Record = {}, + cookies: Record = {} ): Promise { const query = await this.queryRepository.findOne({ where: { id: token }, @@ -88,7 +93,16 @@ export class QueryExecuterService { throw new Error("Query not found"); } - const vm = await this.createVm(query); + const sessionId = cookies["x-session-id"] || null; + + console.log("Session ID:", sessionId); + + if (!sessionId) { + const session = await this.sessionService.create(query.project.id); + cookies["x-session-id"] = session.sessionId; + } + + const vm = await this.createVm(query, cookies["x-session-id"]); const result = await vm.runScript( this.clearImports(query.source), queryData, @@ -99,10 +113,20 @@ export class QueryExecuterService { throw new Error(`Error initializing VM: ${JSON.stringify(result)}`); } + if (!result?.cookies || !result?.cookies["x-session-id"]) { + if (result.cookies === undefined) { + result.cookies = {}; + } + + result.cookies["x-session-id"] = ( + await this.sessionService.get(cookies["x-session-id"], query.project.id) + ).sessionId; + } + return result; } - private async createVm(query: Query) { + private async createVm(query: Query, sessionId: string = null) { const imports = this.parseImports(query.source); const importsParsed = imports.map((imp) => { const item = imp.split("/"); @@ -140,7 +164,9 @@ export class QueryExecuterService { throw new Error(`Plugin ${plugin.name} not found`); } - plugins.push(await registeredPlugins[plugin.name](this, query)); + plugins.push( + await registeredPlugins[plugin.name](this, query, sessionId) + ); } const vm = new Vm({ diff --git a/src/query/query.module.ts b/src/query/query.module.ts index 4917a4f..b186533 100644 --- a/src/query/query.module.ts +++ b/src/query/query.module.ts @@ -14,6 +14,7 @@ import { FunctionService } from "src/query/function/function.service"; import { FunctionController } from "src/query/function/function.controller"; import { RedisManagerModule } from "src/redisManager/redisManager.module"; import { RedisModule } from "src/redis/redis.module"; +import { SessionService } from "./session/session.service"; @Module({ imports: [ @@ -26,7 +27,12 @@ import { RedisModule } from "src/redis/redis.module"; TypeOrmModule.forFeature([Query, FunctionEntity]), ], controllers: [QueryController, CommandController, FunctionController], - providers: [QueryExecuterService, QueryHandlerService, FunctionService], + providers: [ + QueryExecuterService, + SessionService, + QueryHandlerService, + FunctionService, + ], exports: [QueryExecuterService, TypeOrmModule, QueryHandlerService], }) export class QueryModule {} diff --git a/src/query/session/session.service.ts b/src/query/session/session.service.ts new file mode 100644 index 0000000..f680182 --- /dev/null +++ b/src/query/session/session.service.ts @@ -0,0 +1,50 @@ +import { Inject, Injectable } from "@nestjs/common"; +import { RedisClient } from "src/redis/redis.service"; + +@Injectable() +export class SessionService { + constructor(@Inject(RedisClient) private readonly redisClient: RedisClient) {} + + private generateSessionId(length: number = 16): string { + const characters = + "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789"; + let result = ""; + const charactersLength = characters.length; + for (let i = 0; i < length; i++) { + result += characters.charAt(Math.floor(Math.random() * charactersLength)); + } + return `${result}_${new Date().getTime()}`; + } + + async create(prefix: string): Promise<{ sessionId: string }> { + const sessionId = this.generateSessionId(); + await this.set(sessionId, prefix, { + sessionId, + }); + + return { sessionId }; + } + + async get(sessionId: string | null, prefix: string): Promise { + if (!sessionId) { + return await this.create(prefix); + } + + const data = await this.redisClient.get(`${prefix}:${sessionId}`); + + if (data) { + await this.redisClient.set(`${prefix}:${sessionId}`, data, 3600); + return data; + } + + return await this.create(prefix); + } + + async set(sessionId: string, prefix: string, data: any): Promise { + await this.redisClient.set(`${prefix}:${sessionId}`, data, 3600); + } + + async delete(sessionId: string, prefix: string): Promise { + await this.redisClient.del(`${prefix}:${sessionId}`); + } +} diff --git a/src/queue/processors/query.processor.ts b/src/queue/processors/query.processor.ts index b4534e0..4e734fa 100644 --- a/src/queue/processors/query.processor.ts +++ b/src/queue/processors/query.processor.ts @@ -7,6 +7,7 @@ export interface QueryJob { token: string; queryData: any; headers: Record; + cookies: Record; } @Processor(QUEUE_NAMES.QUERY, { concurrency: 5 }) @@ -16,8 +17,13 @@ export class QueryProcessor extends WorkerHost { } async process(job: Job) { - const { token, queryData, headers } = job.data; + const { token, queryData, headers, cookies } = job.data; - return await this.queryExecuterService.runQuery(token, queryData, headers); + return await this.queryExecuterService.runQuery( + token, + queryData, + headers, + cookies + ); } } diff --git a/src/vm/plugins/session.plugin.ts b/src/vm/plugins/session.plugin.ts new file mode 100644 index 0000000..c12e09d --- /dev/null +++ b/src/vm/plugins/session.plugin.ts @@ -0,0 +1,57 @@ +import { Query } from "src/query/entities/query.entity"; +import { Plugin } from "../plugin.class"; +import { QueryExecuterService } from "src/query/executer/query.executer.service"; + +export class SessionPlugin extends Plugin { + constructor( + name: string, + private query: Query, + private queryExecuterService: QueryExecuterService, + private session: Record + ) { + super(name, ["get", "set", "delete"]); + } + + static async init( + query: Query, + queryExecuterService: QueryExecuterService, + sessionId: string | null + ) { + const session = await queryExecuterService.sessionService.get( + sessionId, + query.project.id + ); + + return new SessionPlugin("session", query, queryExecuterService, session); + } + + async get(property: string): Promise { + this.session = await this.queryExecuterService.sessionService.get( + this.session?.sessionId, + this.query.project.id + ); + + return this.session[property]; + } + + async set(property: string, data: any): Promise { + this.session[property] = data; + + await this.queryExecuterService.sessionService.set( + this.session?.sessionId, + this.query.project.id, + this.session + ); + } + + async delete(): Promise { + return await this.queryExecuterService.sessionService.delete( + this.session?.sessionId, + this.query.project.id + ); + } + + onFinish() { + // No resources to clean up + } +} diff --git a/src/vm/vm.class.ts b/src/vm/vm.class.ts index 9a0b361..11ace9c 100644 --- a/src/vm/vm.class.ts +++ b/src/vm/vm.class.ts @@ -13,6 +13,7 @@ export class Vm { private isolate: ivm.Isolate; private timeLimit?: bigint; private cpuTimeLimit?: bigint; + private sessionId: string | null; constructor(configs: { memoryLimit: number; diff --git a/src/vm/vm.constants.ts b/src/vm/vm.constants.ts index 028f567..86a5086 100644 --- a/src/vm/vm.constants.ts +++ b/src/vm/vm.constants.ts @@ -4,6 +4,7 @@ import { Query } from "src/query/entities/query.entity"; import { QueryPlugin } from "./plugins/query.plugin"; import { AxiosPlugin } from "./plugins/axios.plugin"; import { RedisPlugin } from "./plugins/redis.plugin"; +import { SessionPlugin } from "./plugins/session.plugin"; export const registeredPlugins = { db: async (service: QueryExecuterService, query: Query) => { @@ -30,6 +31,13 @@ export const registeredPlugins = { return RedisPlugin.init("redis", redisConnection, query.project.id); }, + session: async ( + service: QueryExecuterService, + query: Query, + sessionId: string | null + ) => { + return SessionPlugin.init(query, service, sessionId); + }, axios: async () => { return AxiosPlugin.init("axios"); }, @@ -50,4 +58,5 @@ export type QueryResponse = { response: any; headers?: Record; redirect?: string; + cookies?: Record; }; diff --git a/tests/base/case1-payload.js b/tests/base/case1-payload.js index 02514ed..ce448a3 100644 --- a/tests/base/case1-payload.js +++ b/tests/base/case1-payload.js @@ -4,6 +4,7 @@ import "module/squel"; import "plugin/db"; import "plugin/axios"; +import "plugin/session"; function createSQL(id) { return squel.select().from("test").where("id = ?", id).toString(); @@ -23,6 +24,10 @@ async function main(input, headers) { // log(a); + log(await session.get("test")); + + await session.set("test", 1); + const res = await db.query(` SELECT SLEEP(10000); `); diff --git a/tests/base/case1.ts b/tests/base/case1.ts index d637724..ff1332c 100644 --- a/tests/base/case1.ts +++ b/tests/base/case1.ts @@ -1,9 +1,9 @@ -import createMigration from "../functions/createMigration"; -import createDatabase from "../functions/createDatabase"; -import createDatabaseNode from "../functions/createDatabaseNode"; -import createProject from "../functions/createProject"; +// import createMigration from "../functions/createMigration"; +// import createDatabase from "../functions/createDatabase"; +// import createDatabaseNode from "../functions/createDatabaseNode"; +// import createProject from "../functions/createProject"; import createQuery from "../functions/createQuery"; -import databaseMigrationUp from "../functions/databaseMigrationUp"; +// import databaseMigrationUp from "../functions/databaseMigrationUp"; import runQuery from "../functions/runQuery"; import * as fs from "fs"; import * as path from "path"; @@ -45,6 +45,7 @@ import * as path from "path"; const result = await runQuery(query.id, { id: 1 }); console.log("Query Result:", result.data); + console.log("headers:", result.headers); } catch (error) { console.error("Error during test execution"); } diff --git a/tests/functions/runQuery.ts b/tests/functions/runQuery.ts index 8798ca8..a81856f 100644 --- a/tests/functions/runQuery.ts +++ b/tests/functions/runQuery.ts @@ -6,7 +6,12 @@ export default async (token: string, queryData: Record) => { const response = await axios.post( `${config.url}/query/run/${token}`, queryData, - { headers: { "x-api-token": "efbeccd6-dde1-47dc-b3aa-4fbd773d5429" } } + { + headers: { + "x-api-token": "efbeccd6-dde1-47dc-b3aa-4fbd773d5429", + Cookie: `x-session-id=psaCHcYFMrt6RnUz_1760094020243`, + }, + } ); return response;