Skip to content

Commit 0228620

Browse files
committed
fix(files): bound CSV previews and preserve render provenance
1 parent d60fdd7 commit 0228620

10 files changed

Lines changed: 131 additions & 12 deletions

File tree

‎apps/sim/lib/file-parsers/csv-preview-slice.test.ts‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -81,3 +81,37 @@ describe('getCsvPreviewSlice', () => {
8181
expect(destroySpy).toHaveBeenCalled()
8282
})
8383
})
84+
85+
it('rejects an oversized streamed record before consuming the entire file', async () => {
86+
let chunksRead = 0
87+
const source = Readable.from(
88+
(async function* () {
89+
yield Buffer.from('header\n"')
90+
for (let i = 0; i < 1000; i++) {
91+
chunksRead++
92+
yield Buffer.alloc(16 * 1024, 'x')
93+
}
94+
})()
95+
)
96+
mockDownloadFileStream.mockResolvedValue(source)
97+
await expect(getCsvPreviewSlice(args)).rejects.toMatchObject({ code: 'CSV_MAX_RECORD_SIZE' })
98+
expect(chunksRead).toBeLessThan(30)
99+
expect(source.destroyed).toBe(true)
100+
})
101+
102+
it('bounds bytes consumed even when blank records produce no preview rows', async () => {
103+
let chunksRead = 0
104+
const source = Readable.from(
105+
(async function* () {
106+
yield Buffer.from('header\n')
107+
for (let i = 0; i < 1000; i++) {
108+
chunksRead++
109+
yield Buffer.alloc(64 * 1024, '\n')
110+
}
111+
})()
112+
)
113+
mockDownloadFileStream.mockResolvedValue(source)
114+
await expect(getCsvPreviewSlice(args)).rejects.toThrow('CSV preview')
115+
expect(chunksRead).toBeLessThan(140)
116+
expect(source.destroyed).toBe(true)
117+
})

‎apps/sim/lib/file-parsers/csv-preview-slice.ts‎

Lines changed: 26 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2,11 +2,14 @@ import { Readable } from 'node:stream'
22
import { truncate } from '@sim/utils/string'
33
import { parse as parseCsvStream } from 'csv-parse'
44
import { CSV_PREVIEW_MAX_ROWS } from '@/lib/api/contracts/workspace-file-table'
5+
import { assertKnownSizeWithinLimit } from '@/lib/core/utils/stream-limits'
56
import type { StorageContext } from '@/lib/uploads/config'
67
import { downloadFileStream } from '@/lib/uploads/core/storage-service'
78

89
/** Cap a single cell so one pathological field can't bloat the preview payload. */
910
const MAX_CELL_LENGTH = 2_000
11+
const MAX_RECORD_SIZE = 256 * 1024
12+
const MAX_PREVIEW_BYTES = 8 * 1024 * 1024
1013

1114
/** Read at most this many bytes while sniffing the first line for the delimiter. */
1215
const DELIMITER_SNIFF_MAX_BYTES = 256 * 1024
@@ -44,7 +47,7 @@ function cell(value: unknown): string {
4447
/**
4548
* Streams the first {@link CSV_PREVIEW_MAX_ROWS} rows of a CSV/TSV from storage without
4649
* ever buffering the whole file. The source stream is destroyed as soon as enough rows are
47-
* read (one past the cap, to detect truncation), so a multi-GB file costs O(rows) of memory.
50+
* read (one past the cap, to detect truncation), with record and total-byte limits bounding parser memory and work.
4851
*/
4952
export async function getCsvPreviewSlice({
5053
key,
@@ -60,6 +63,17 @@ export async function getCsvPreviewSlice({
6063
signal?.addEventListener('abort', onAbort, { once: true })
6164

6265
const reader = source[Symbol.asyncIterator]()
66+
let bytesRead = 0
67+
let piped: Readable | undefined
68+
let parser: ReturnType<typeof parseCsvStream> | undefined
69+
async function nextChunk() {
70+
const next = await reader.next()
71+
if (!next.done) {
72+
bytesRead += Buffer.byteLength(next.value)
73+
assertKnownSizeWithinLimit(bytesRead, MAX_PREVIEW_BYTES, 'CSV preview')
74+
}
75+
return next
76+
}
6377

6478
try {
6579
// Pull chunks until the first newline so the delimiter can be sniffed before parsing.
@@ -71,12 +85,14 @@ export async function getCsvPreviewSlice({
7185
let firstLine = ''
7286
let sniffedBytes = 0
7387
while (true) {
74-
const { value, done } = await reader.next()
88+
const { value, done } = await nextChunk()
7589
if (done) break
7690
const chunk = Buffer.isBuffer(value) ? value : Buffer.from(value)
7791
sniffed.push(chunk)
7892
sniffedBytes += chunk.length
79-
const text = chunk.toString('utf-8')
93+
const text = chunk
94+
.subarray(0, Math.max(0, DELIMITER_SNIFF_MAX_BYTES - sniffedBytes + chunk.length))
95+
.toString('utf-8')
8096
const nl = text.indexOf('\n')
8197
if (nl !== -1) {
8298
firstLine += text.slice(0, nl)
@@ -91,13 +107,13 @@ export async function getCsvPreviewSlice({
91107
}
92108

93109
const delimiter = detectDelimiter(firstLine)
94-
const parser = parseCsvStream({
110+
parser = parseCsvStream({
95111
columns: false,
96112
skip_empty_lines: true,
97113
trim: true,
98114
relax_column_count: true,
99115
relax_quotes: true,
100-
skip_records_with_error: true,
116+
max_record_size: MAX_RECORD_SIZE,
101117
cast: false,
102118
bom: true,
103119
delimiter,
@@ -107,13 +123,13 @@ export async function getCsvPreviewSlice({
107123
async function* rejoin() {
108124
for (const chunk of sniffed) yield chunk
109125
while (true) {
110-
const { value, done } = await reader.next()
126+
const { value, done } = await nextChunk()
111127
if (done) return
112128
yield value
113129
}
114130
}
115-
const piped = Readable.from(rejoin())
116-
piped.on('error', (err) => parser.destroy(err))
131+
piped = Readable.from(rejoin())
132+
piped.on('error', (err) => parser?.destroy(err))
117133
piped.pipe(parser)
118134

119135
let headers: string[] = []
@@ -141,6 +157,8 @@ export async function getCsvPreviewSlice({
141157
if (signal?.aborted) signal.throwIfAborted()
142158
throw error
143159
} finally {
160+
piped?.destroy()
161+
parser?.destroy()
144162
signal?.removeEventListener('abort', onAbort)
145163
source.destroy()
146164
}

‎apps/sim/lib/mothership/__integration__/content-delivery.integration.ts‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -132,6 +132,7 @@ import {
132132
} from '@/lib/mothership/__integration__/content-delivery-fixture'
133133
import { observePrivateFile } from '@/lib/mothership/agent-cli/engines/observe-private-file'
134134
import { createAgentCliTransport } from '@/lib/mothership/agent-cli/transport'
135+
import { createWorkbenchFileProvenance } from '@/lib/mothership/agent-cli/workbench-file-provenance'
135136
import { createCopilotChatPrincipal } from '@/lib/mothership/auth/application-delegation'
136137
import { createRegistryContentGate } from '@/lib/mothership/content-delivery/registry-gate'
137138
import { proxySandboxResourceRequest } from '@/lib/mothership/tools/sandbox-resource-transport'
@@ -251,6 +252,12 @@ describe('content reaching Chat is admitted per item', () => {
251252
chatId: generateId(),
252253
},
253254
registry,
255+
files: createWorkbenchFileProvenance({
256+
userId: fixture.ids.aliceId,
257+
workspaceId: fixture.ids.workspaceId,
258+
sessionKey: generateId(),
259+
resolvedSecretTraceRegistry: registry,
260+
}),
254261
resources: [],
255262
observeReads: true,
256263
})

‎apps/sim/lib/mothership/agent-cli/file-provenance.test.ts‎

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,7 @@ vi.mock('@/lib/mothership/agent-cli/scoped-transport', () => ({
5959

6060
import { createFileReadTransport } from '@/lib/mothership/agent-cli/file-read-transport'
6161
import { createAgentCliTransport } from '@/lib/mothership/agent-cli/transport'
62+
import { createWorkbenchFileProvenance } from '@/lib/mothership/agent-cli/workbench-file-provenance'
6263
import { createCopilotChatPrincipal } from '@/lib/mothership/auth/application-delegation'
6364
import { createRegistryContentGate } from '@/lib/mothership/content-delivery'
6465
import { inspectToolResultForCopilot } from '@/lib/mothership/request/tools/resolved-secret-result'
@@ -204,11 +205,19 @@ function runCli(invocation: CliInvocation, trace = registry()) {
204205
const invocation = { userId: 'reader', workspaceId: 'workspace', chatId: 'chat' }
205206

206207
/** Chat's CLI transport as every invocation composes it: the registry gate around the routes. */
207-
function readTransport(trace = registry()): typeof fetch {
208+
function readTransport(trace = registry(), tracked = true): typeof fetch {
208209
const transport = createAgentCliTransport({
209210
endpoint: 'https://sim.test',
210211
invocation,
211212
registry: trace,
213+
files: tracked
214+
? createWorkbenchFileProvenance({
215+
userId: 'reader',
216+
workspaceId: 'workspace',
217+
sessionKey: 'chat',
218+
resolvedSecretTraceRegistry: trace,
219+
})
220+
: undefined,
212221
resources: [],
213222
observeReads: true,
214223
})
@@ -429,6 +438,16 @@ describe('file provenance at the actual CLI and model-result boundary', () => {
429438
).not.toContain(content)
430439
})
431440

441+
it('refuses opaque protected downloads without workbench provenance tracking', async () => {
442+
classify('exact', revision, [
443+
{ name: 'FILE_SECRET', encryptedValue: 'fixture-ciphertext', sourceUserId: 'reader' },
444+
])
445+
const response = await readTransport(registry(), false)(fileRequest())
446+
expect(response.status).toBe(403)
447+
expect(await response.text()).not.toContain(content)
448+
expect(mocks.stream).not.toHaveBeenCalled()
449+
})
450+
432451
it('imports encrypted file provenance into a fresh registry before model projection', async () => {
433452
classify('exact', revision, [
434453
{ name: 'FILE_SECRET', encryptedValue: 'fixture-ciphertext', sourceUserId: 'reader' },

‎apps/sim/lib/mothership/agent-cli/file-read-transport.ts‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ export function createFileReadTransport(context: {
3838
transport: typeof fetch
3939
invocation: CopilotChatDelegationContext & { workspaceId: string }
4040
registry?: ResolvedSecretTraceRegistry
41+
tracksWorkbenchDownloads?: boolean
4142
}): typeof fetch {
4243
const base = new URL(context.endpoint)
4344
const basePath = base.pathname.replace(/\/$/, '')
@@ -105,7 +106,7 @@ export function createFileReadTransport(context: {
105106
workspaceId: parsed.data.query.workspaceId,
106107
actorUserId: context.invocation.userId,
107108
unknown: 'taint',
108-
destination: 'workbench',
109+
destination: context.tracksWorkbenchDownloads ? 'workbench' : 'model',
109110
}),
110111
() =>
111112
downloadWorkspaceFileStream.execute({

‎apps/sim/lib/mothership/agent-cli/transport.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@ export function createAgentCliTransport(context: {
4040
transport: createTracedCliTransport(endpoint, createScopedCliTransport(endpoint, invocation)),
4141
invocation,
4242
registry,
43+
tracksWorkbenchDownloads: Boolean(files),
4344
})
4445
const reads: typeof fetch = async (input, init) => {
4546
const { response, admitted } = await dispatchUnderContentGate(gate, () => routes(input, init))

‎apps/sim/lib/projects/files/application/artifacts.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -277,6 +277,7 @@ export const readProjectFileArtifact = defineAuthorizedProjectFileUseCase<
277277
const inputIdentity = fileDocumentInputIdentity(context.owner, manifest.dependencies)
278278
const compiled = await compileFileDocument({
279279
owner: context.owner,
280+
sourceMime: file.contentType,
280281
source: content.toString('utf-8'),
281282
fileName: file.originalName,
282283
inputs,

‎apps/sim/lib/uploads/documents/compile.test.ts‎

Lines changed: 31 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,11 @@ vi.mock('@/lib/uploads/documents/compiled-store', async (importOriginal) => ({
4040

4141
import { type ContentDeliveryGate, observeContentDelivery } from '@/lib/core/application'
4242
import { createCopilotChatFilePrincipal } from '@/lib/mothership/auth/file-delegation'
43-
import { collectReferencedFileIds, compileDoc } from '@/lib/uploads/documents/compile'
43+
import {
44+
collectReferencedFileIds,
45+
compileDoc,
46+
compileFileDocument,
47+
} from '@/lib/uploads/documents/compile'
4448
import { DocCompileUserError } from '@/lib/uploads/documents/compile-error'
4549

4650
const executeInSandboxMock = remoteSandboxMockFns.mockExecuteInSandbox
@@ -391,3 +395,29 @@ describe('isolated-VM compile for Chat', () => {
391395
expect(runSandboxTaskMock).toHaveBeenCalledTimes(2)
392396
})
393397
})
398+
399+
describe('compileFileDocument source engine', () => {
400+
it('uses the JavaScript PDF task even with the document sandbox enabled', async () => {
401+
setEnvFlags({ isDocSandboxEnabled: true })
402+
loadCompiledDocMock.mockResolvedValue(null)
403+
runSandboxTaskMock.mockResolvedValue(Buffer.from('%PDF-1.7'))
404+
runSandboxTaskMock.mockClear()
405+
executeInSandboxMock.mockClear()
406+
const source = 'const doc = await PDFDocument.create()'
407+
const result = await compileFileDocument({
408+
owner: { entityType: 'project', entityId: ID },
409+
source,
410+
sourceMime: 'text/x-pdflibjs',
411+
fileName: 'report.pdf',
412+
inputs: [],
413+
maxBytes: 1024,
414+
})
415+
expect(result.buffer.toString()).toBe('%PDF-1.7')
416+
expect(runSandboxTaskMock).toHaveBeenCalledWith(
417+
'pdf-generate',
418+
{ code: source },
419+
expect.any(Object)
420+
)
421+
expect(executeInSandboxMock).not.toHaveBeenCalled()
422+
})
423+
})

‎apps/sim/lib/uploads/documents/compile.ts‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -623,6 +623,7 @@ async function renderPreparedDocument(source: string, fmt: E2BDocFormat, inputs:
623623
/** Compiles a document against a fixed, authorized input manifest and canonical owner cache. */
624624
export async function compileFileDocument(args: {
625625
owner: EditableFileOwner
626+
sourceMime: string
626627
source: string
627628
fileName: string
628629
inputs: readonly { fileId: string; contentType: string; content: Buffer }[]
@@ -651,7 +652,7 @@ export async function compileFileDocument(args: {
651652
dependsOnReferencedFiles: args.inputs.length > 0,
652653
}
653654
let buffer: Buffer
654-
if (isDocSandboxEnabled) {
655+
if (isDocSandboxEnabled && args.sourceMime !== 'text/x-pdflibjs') {
655656
buffer = await renderPreparedDocument(
656657
args.source,
657658
fmt,

‎apps/sim/lib/workspace-files/__integration__/copilot-file-versions.integration.ts‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@ import {
4040
import { runCli } from '@/lib/mothership/agent-cli/run-cli'
4141
import { createScopedCliTransport } from '@/lib/mothership/agent-cli/scoped-transport'
4242
import { createAgentCliTransport } from '@/lib/mothership/agent-cli/transport'
43+
import { createWorkbenchFileProvenance } from '@/lib/mothership/agent-cli/workbench-file-provenance'
4344
import { createCopilotChatPrincipal } from '@/lib/mothership/auth/application-delegation'
4445
import {
4546
updateWorkspaceFileContent,
@@ -145,6 +146,12 @@ describe('chat-delegated file version history', () => {
145146
endpoint: ORIGIN,
146147
invocation,
147148
...(registry ? { registry } : {}),
149+
files: createWorkbenchFileProvenance({
150+
userId,
151+
workspaceId: fixture.workspaceId,
152+
sessionKey: invocation.chatId,
153+
resolvedSecretTraceRegistry: registry,
154+
}),
148155
resources: [],
149156
observeReads: true,
150157
})

0 commit comments

Comments
 (0)