Skip to content

Commit cea09af

Browse files
authored
Harden SQLite store logs, artifacts, and lifecycle
## Summary - append log chunks directly instead of read/modify/write rewriting the full log file - prevent artifact file path traversal outside each task artifact directory - add an explicit close() API for the SQLite connection - cover log append, artifact path validation, and close/reopen behavior in tests ## Tests - npm ci - npm run typecheck - npm test -- --run - npm run build Fixes #2 Fixes #3 Fixes #4
1 parent 9f0129d commit cea09af

2 files changed

Lines changed: 42 additions & 7 deletions

File tree

src/index.ts

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import Database from "better-sqlite3";
22
import { mkdirSync } from "node:fs";
3-
import { mkdir, readFile, stat, writeFile } from "node:fs/promises";
4-
import { dirname, join } from "node:path";
3+
import { appendFile, mkdir, readFile, stat, writeFile } from "node:fs/promises";
4+
import { dirname, join, resolve, sep } from "node:path";
55
import { createId, nowIso, type ArtifactRecord, type LogChunk, type RuntimeEvent, type RuntimeRecord, type SessionRecord, type TaskRecord, type TaskStore } from "@agent-dispatch/core";
66

77
export interface SqliteTaskStoreOptions {
@@ -75,9 +75,7 @@ export class SqliteTaskStore implements TaskStore {
7575

7676
async appendLog(taskId: string, chunk: string): Promise<void> {
7777
await this.ensureReady();
78-
const path = this.logPath(taskId);
79-
const existing = await readFile(path, "utf8").catch(() => "");
80-
await writeFile(path, `${existing}${chunk}`);
78+
await appendFile(this.logPath(taskId), chunk);
8179
}
8280

8381
async readLogs(taskId: string, cursor = 0, limit = 64_000): Promise<LogChunk> {
@@ -101,7 +99,7 @@ export class SqliteTaskStore implements TaskStore {
10199

102100
async saveArtifactFile(taskId: string, name: string, bytes: Uint8Array, contentType = "application/octet-stream"): Promise<ArtifactRecord> {
103101
await this.ensureReady();
104-
const uri = join(this.stateDir, "artifacts", taskId, name);
102+
const uri = this.artifactPath(taskId, name);
105103
await mkdir(dirname(uri), { recursive: true });
106104
await writeFile(uri, bytes);
107105
const sizeBytes = (await stat(uri)).size;
@@ -110,6 +108,10 @@ export class SqliteTaskStore implements TaskStore {
110108
return artifact;
111109
}
112110

111+
close(): void {
112+
this.db.close();
113+
}
114+
113115
private initialize(): void {
114116
this.db.exec("create table if not exists schema_migrations (id text primary key, applied_at text not null)");
115117
for (const migration of migrations) {
@@ -123,6 +125,15 @@ export class SqliteTaskStore implements TaskStore {
123125
}
124126
}
125127

128+
private artifactPath(taskId: string, name: string): string {
129+
const root = resolve(this.stateDir, "artifacts", taskId);
130+
const uri = resolve(root, name);
131+
if (uri !== root && !uri.startsWith(`${root}${sep}`)) {
132+
throw new Error("Artifact name must stay within the task artifact directory.");
133+
}
134+
return uri;
135+
}
136+
126137
private logPath(taskId: string): string {
127138
return join(this.stateDir, "logs", `${taskId}.log`);
128139
}

test/store.test.ts

Lines changed: 25 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { mkdtemp, rm } from "node:fs/promises";
1+
import { mkdtemp, readFile, rm, stat } from "node:fs/promises";
22
import { join } from "node:path";
33
import { tmpdir } from "node:os";
44
import { afterEach, beforeEach, describe, expect, it } from "vitest";
@@ -14,6 +14,7 @@ beforeEach(async () => {
1414
});
1515

1616
afterEach(async () => {
17+
store.close();
1718
await rm(stateDir, { recursive: true, force: true });
1819
});
1920

@@ -67,12 +68,35 @@ describe("SqliteTaskStore", () => {
6768
expect(second).toMatchObject({ data: "cde", nextCursor: 5 });
6869
});
6970

71+
it("appends log chunks without rewriting existing content", async () => {
72+
await Promise.all(Array.from({ length: 25 }, (_, index) => store.appendLog("task_1", `${index}\n`)));
73+
const logs = await store.readLogs("task_1", 0, 10_000);
74+
75+
for (let index = 0; index < 25; index += 1) {
76+
expect(logs.data).toContain(`${index}\n`);
77+
}
78+
});
79+
7080
it("persists artifact metadata and files", async () => {
7181
const artifact = await store.saveArtifactFile("task_1", "result.txt", Buffer.from("done"), "text/plain");
7282
const artifacts = await store.listArtifacts("task_1");
7383

7484
expect(artifacts).toHaveLength(1);
7585
expect(artifacts[0]).toMatchObject({ id: artifact.id, kind: "file", contentType: "text/plain", sizeBytes: 4 });
86+
await expect(readFile(artifact.uri, "utf8")).resolves.toBe("done");
87+
});
88+
89+
it("rejects artifact paths outside the task artifact directory", async () => {
90+
await expect(store.saveArtifactFile("task_1", "../escape.txt", Buffer.from("nope"))).rejects.toThrow("Artifact name");
91+
await expect(store.saveArtifactFile("task_1", "/tmp/escape.txt", Buffer.from("nope"))).rejects.toThrow("Artifact name");
92+
await expect(stat(join(stateDir, "artifacts", "escape.txt"))).rejects.toThrow();
93+
});
94+
95+
it("closes the database connection explicitly", async () => {
96+
store.close();
97+
expect(() => store.appliedMigrations()).toThrow();
98+
store = new SqliteTaskStore({ stateDir });
99+
expect(store.appliedMigrations()).toEqual(["001_initial"]);
76100
});
77101
});
78102

0 commit comments

Comments
 (0)