Skip to content

Commit ff0010a

Browse files
thinhnx128thinhnx128
committed
fix(security): harden PostgresClient SQL interpolation (CWE-89) and sync-rpc execution (CWE-78, CWE-668)
1 parent 0f1a443 commit ff0010a

4 files changed

Lines changed: 120 additions & 20 deletions

File tree

‎JS/edgechains/arakoodev/src/db/src/lib/postgres-client/PostgresClient.ts‎

Lines changed: 56 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,28 @@
11
import { createConnection, getManager } from "typeorm";
2+
3+
export function escapeSqlLiteral(val: unknown): string {
4+
if (val === null || val === undefined) {
5+
return "";
6+
}
7+
return String(val).replace(/\0/g, "").replace(/'/g, "''");
8+
}
9+
10+
const SQL_IDENTIFIER_RE = /^[A-Za-z_][A-Za-z0-9_]{0,62}$/;
11+
export function validateSqlIdentifier(name: string): string {
12+
if (typeof name !== "string" || !SQL_IDENTIFIER_RE.test(name)) {
13+
throw new Error(`Invalid SQL identifier: ${JSON.stringify(name)}`);
14+
}
15+
return name;
16+
}
17+
18+
export function sanitizeNumber(num: unknown): number {
19+
const parsed = typeof num === "number" ? num : Number(num);
20+
if (!Number.isFinite(parsed)) {
21+
throw new Error(`Invalid numeric value: ${String(num)}`);
22+
}
23+
return parsed;
24+
}
25+
226
export class PostgresClient {
327
wordEmbeddings: number[][];
428
metric: PostgresDistanceMetric;
@@ -30,22 +54,40 @@ export class PostgresClient {
3054
}
3155

3256
async dbQuery() {
57+
const safeTableName = validateSqlIdentifier(this.tableName);
58+
const safeMetadataTable = validateSqlIdentifier(this.arkRequest?.metadataTable);
59+
const safeJoinTable = validateSqlIdentifier(`${safeTableName}_join_${safeMetadataTable}`);
60+
const safeMetaTargetTable = validateSqlIdentifier(`${safeTableName}_${safeMetadataTable}`);
61+
const safeNamespace = escapeSqlLiteral(this.namespace);
62+
const safeSearchQuery = escapeSqlLiteral(this.arkRequest?.query);
63+
const safeProbes = Math.max(1, Math.trunc(sanitizeNumber(this.probes)));
64+
const safeTopK = Math.max(1, Math.trunc(sanitizeNumber(this.topK)));
65+
const safeUpperLimit = Math.max(1, Math.trunc(sanitizeNumber(this.upperLimit)));
66+
67+
const textBaseWeight = sanitizeNumber(this.arkRequest?.textWeight?.baseWeight);
68+
const textFineTuneWeight = sanitizeNumber(this.arkRequest?.textWeight?.fineTuneWeight);
69+
const simBaseWeight = sanitizeNumber(this.arkRequest?.similarityWeight?.baseWeight);
70+
const simFineTuneWeight = sanitizeNumber(this.arkRequest?.similarityWeight?.fineTuneWeight);
71+
const dateBaseWeight = sanitizeNumber(this.arkRequest?.dateWeight?.baseWeight);
72+
const dateFineTuneWeight = sanitizeNumber(this.arkRequest?.dateWeight?.fineTuneWeight);
73+
3374
const connection = await createConnection();
3475
const entityManager = await getManager();
3576
try {
36-
const query1 = `SET LOCAL ivfflat.probes = ${this.probes};`;
77+
const query1 = `SET LOCAL ivfflat.probes = ${safeProbes};`;
3778
await entityManager.query(query1);
3879

3980
let query: string = "";
4081

4182
for (let i = 0; i < this.wordEmbeddings.length; i++) {
42-
const embedding: string = JSON.stringify(this.wordEmbeddings[i]);
83+
const sanitizedVector = this.wordEmbeddings[i].map((v) => sanitizeNumber(v));
84+
const embedding: string = JSON.stringify(sanitizedVector);
4385

4486
query += `( SELECT id, raw_text, document_date, metadata, namespace, filename, timestamp,
45-
${this.arkRequest.textWeight.baseWeight} / (ROW_NUMBER() OVER (ORDER BY text_rank DESC) + ${this.arkRequest.textWeight.fineTuneWeight}) +
46-
${this.arkRequest.similarityWeight.baseWeight} / (ROW_NUMBER() OVER (ORDER BY similarity DESC) + ${this.arkRequest.similarityWeight.fineTuneWeight}) +
47-
${this.arkRequest.dateWeight.baseWeight} / (ROW_NUMBER() OVER (ORDER BY date_rank DESC) + ${this.arkRequest.dateWeight.fineTuneWeight}) AS rrf_score
48-
FROM ( SELECT sv.id, sv.raw_text, sv.namespace, sv.filename, sv.timestamp, svtm.document_date, svtm.metadata, ts_rank_cd(sv.tsv, plainto_tsquery('${"english"}', '${this.arkRequest.query}')) AS text_rank, `;
87+
${textBaseWeight} / (ROW_NUMBER() OVER (ORDER BY text_rank DESC) + ${textFineTuneWeight}) +
88+
${simBaseWeight} / (ROW_NUMBER() OVER (ORDER BY similarity DESC) + ${simFineTuneWeight}) +
89+
${dateBaseWeight} / (ROW_NUMBER() OVER (ORDER BY date_rank DESC) + ${dateFineTuneWeight}) AS rrf_score
90+
FROM ( SELECT sv.id, sv.raw_text, sv.namespace, sv.filename, sv.timestamp, svtm.document_date, svtm.metadata, ts_rank_cd(sv.tsv, plainto_tsquery('${"english"}', '${safeSearchQuery}')) AS text_rank, `;
4991

5092
if (this.metric === PostgresDistanceMetric.COSINE)
5193
query += `1 - (sv.embedding <=> '${embedding}') AS similarity, `;
@@ -54,16 +96,16 @@ export class PostgresClient {
5496
if (this.metric === PostgresDistanceMetric.L2)
5597
query += `sv.embedding <-> '${embedding}' AS similarity, `;
5698

57-
query += `CASE WHEN svtm.document_date IS NULL THEN 0 ELSE EXTRACT(YEAR FROM svtm.document_date) * 365 + EXTRACT(DOY FROM svtm.document_date) END AS date_rank FROM (SELECT id, raw_text, embedding, tsv, namespace, filename, timestamp from ${this.tableName} WHERE namespace = '${this.namespace}'`;
99+
query += `CASE WHEN svtm.document_date IS NULL THEN 0 ELSE EXTRACT(YEAR FROM svtm.document_date) * 365 + EXTRACT(DOY FROM svtm.document_date) END AS date_rank FROM (SELECT id, raw_text, embedding, tsv, namespace, filename, timestamp from ${safeTableName} WHERE namespace = '${safeNamespace}'`;
58100

59101
if (this.metric === PostgresDistanceMetric.COSINE)
60-
query += ` ORDER BY embedding <=> '${embedding}' LIMIT ${this.topK}`;
102+
query += ` ORDER BY embedding <=> '${embedding}' LIMIT ${safeTopK}`;
61103
if (this.metric === PostgresDistanceMetric.IP)
62-
query += ` ORDER BY embedding <#> '${embedding}' LIMIT ${this.topK}`;
104+
query += ` ORDER BY embedding <#> '${embedding}' LIMIT ${safeTopK}`;
63105
if (this.metric === PostgresDistanceMetric.L2)
64-
query += ` ORDER BY embedding <-> '${embedding}' LIMIT ${this.topK}`;
106+
query += ` ORDER BY embedding <-> '${embedding}' LIMIT ${safeTopK}`;
65107

66-
query += `) sv JOIN ${this.tableName}_join_${this.arkRequest.metadataTable} jtm ON sv.id = jtm.id JOIN ${this.tableName}_${this.arkRequest.metadataTable} svtm ON jtm.metadata_id = svtm.metadata_id) subquery `;
108+
query += `) sv JOIN ${safeJoinTable} jtm ON sv.id = jtm.id JOIN ${safeMetaTargetTable} svtm ON jtm.metadata_id = svtm.metadata_id) subquery `;
67109

68110
switch (this.arkRequest.orderRRF) {
69111
case "text_rank":
@@ -80,16 +122,16 @@ export class PostgresClient {
80122
break;
81123
}
82124

83-
query += ` LIMIT ${this.topK})`;
125+
query += ` LIMIT ${safeTopK})`;
84126
if (i < this.wordEmbeddings.length - 1) {
85127
query += " UNION ALL \n";
86128
}
87129
}
88130

89131
if (this.wordEmbeddings.length > 1) {
90-
query = `SELECT * FROM (SELECT DISTINCT ON (result.id) * FROM ( ${query} ) result) subquery ORDER BY rrf_score DESC LIMIT ${this.upperLimit};`;
132+
query = `SELECT * FROM (SELECT DISTINCT ON (result.id) * FROM ( ${query} ) result) subquery ORDER BY rrf_score DESC LIMIT ${safeUpperLimit};`;
91133
} else {
92-
query += ` ORDER BY rrf_score DESC LIMIT ${this.topK};`;
134+
query += ` ORDER BY rrf_score DESC LIMIT ${safeTopK};`;
93135
}
94136
const results = await entityManager.query(query);
95137
return results;

‎JS/edgechains/arakoodev/src/db/src/tests/postgres-client/postgresClient.test.ts‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,4 +52,57 @@ describe("PostgresClient", () => {
5252
// Check that the correct SQL query was generated
5353
expect(mockQuery).toHaveBeenCalledWith(expect.any(String));
5454
});
55+
56+
test("dbQuery should escape single quotes in search query and namespace to prevent SQL injection (CWE-89)", async () => {
57+
const mockQuery = jest.fn().mockResolvedValue([]);
58+
(getManager as jest.Mock).mockResolvedValueOnce({ query: mockQuery });
59+
(createConnection as jest.Mock).mockResolvedValueOnce({ close: jest.fn() });
60+
61+
const clientWithQuotes = new PostgresClient(
62+
[[0.1, 0.2]],
63+
PostgresDistanceMetric.COSINE,
64+
10,
65+
5,
66+
"test_table",
67+
"ns'); DROP TABLE users; --",
68+
{
69+
textWeight: { baseWeight: 1, fineTuneWeight: 0.5 },
70+
similarityWeight: { baseWeight: 1, fineTuneWeight: 0.5 },
71+
dateWeight: { baseWeight: 1, fineTuneWeight: 0.5 },
72+
orderRRF: "text_rank",
73+
metadataTable: "test_metadata",
74+
query: "search'); DROP TABLE secrets; --",
75+
},
76+
100
77+
);
78+
79+
await clientWithQuotes.dbQuery();
80+
81+
expect(mockQuery).toHaveBeenCalledTimes(2);
82+
const generatedSql = mockQuery.mock.calls[1][0] as string;
83+
expect(generatedSql).toContain("search''); DROP TABLE secrets; --");
84+
expect(generatedSql).toContain("ns''); DROP TABLE users; --");
85+
});
86+
87+
test("dbQuery should reject invalid SQL table identifiers (CWE-89)", async () => {
88+
const badClient = new PostgresClient(
89+
[[0.1, 0.2]],
90+
PostgresDistanceMetric.COSINE,
91+
10,
92+
5,
93+
"test_table; DROP TABLE users; --",
94+
"test_namespace",
95+
{
96+
textWeight: { baseWeight: 1, fineTuneWeight: 0.5 },
97+
similarityWeight: { baseWeight: 1, fineTuneWeight: 0.5 },
98+
dateWeight: { baseWeight: 1, fineTuneWeight: 0.5 },
99+
orderRRF: "text_rank",
100+
metadataTable: "test_metadata",
101+
query: "test_query",
102+
},
103+
100
104+
);
105+
106+
await expect(badClient.dbQuery()).rejects.toThrow(/Invalid SQL identifier/);
107+
});
55108
});

‎JS/edgechains/arakoodev/src/sync-rpc/lib/sync-rpc.ts‎

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { execSync } from "node:child_process";
1+
import { execFileSync } from "node:child_process";
22
import path from "node:path";
33
import fs from "node:fs";
44
import crypto from "node:crypto";
@@ -49,13 +49,11 @@ function createSyncRPC(filename: string) {
4949
fs.writeFileSync(scriptPath, scriptContent);
5050

5151
try {
52-
const output = execSync(`node ${scriptPath}`, {
52+
const output = execFileSync(process.execPath, [scriptPath], {
5353
encoding: "utf8",
5454
stdio: ["pipe", "pipe", "inherit"],
5555
});
5656

57-
fs.unlinkSync(scriptPath);
58-
5957
const trimmedOutput = output.trim();
6058

6159
try {
@@ -71,7 +69,6 @@ function createSyncRPC(filename: string) {
7169
);
7270
}
7371
} catch (error: any) {
74-
fs.unlinkSync(scriptPath);
7572
if (error) {
7673
console.error("Execution error:", error.message);
7774
if (error.stderr) {
@@ -80,6 +77,14 @@ function createSyncRPC(filename: string) {
8077
throw new Error(`Execution error: ${error.message}`);
8178
}
8279
throw error;
80+
} finally {
81+
try {
82+
if (fs.existsSync(scriptPath)) {
83+
fs.unlinkSync(scriptPath);
84+
}
85+
} catch {
86+
// Ignore cleanup failure so original execution error is preserved
87+
}
8388
}
8489
};
8590
}

‎JS/wasm/sync-rpc/worker.js‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -73,4 +73,4 @@ function init(filename) {
7373
}
7474

7575
// server.listen(6553);
76-
server.listen(process.argv[2]);
76+
server.listen(process.argv[2], "127.0.0.1");

0 commit comments

Comments
 (0)