mirror of
https://github.com/HeyPuter/puter.git
synced 2026-10-11 22:32:26 +00:00
perf(ai-ocr): read Puter files only when the provider needs the bytes
resolveFileInput resolves and access-checks a Puter file and defers the read to load(); loadFileInput is that plus the read. OCR resolves first: Textract reads the object straight from S3 and downloads it only if Textract can't reach the bucket, and only Mistral loads the bytes. A file over the size cap is refused from its recorded size before any read. countPdfPages inflates object streams on the libuv pool instead of the request thread.
This commit is contained in:
1 parent
af4c03e5b6
commit
bbf0ccc1f4
6 files changed
+300
-96
No files matched your search
@@ -43,7 +43,10 @@ import {
|
||||
vi,
|
||||
type MockInstance,
|
||||
} from 'vitest';
|
||||
import { UnsupportedDocumentException } from '@aws-sdk/client-textract';
|
||||
import {
|
||||
InvalidS3ObjectException,
|
||||
UnsupportedDocumentException,
|
||||
} from '@aws-sdk/client-textract';
|
||||
import { v4 as uuidv4 } from 'uuid';
|
||||
|
||||
import type { Actor } from '../../core/actor.js';
|
||||
@@ -349,13 +352,63 @@ describe('OCRDriver.recognize (aws-textract)', () => {
|
||||
expect(textractSendMock).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
// Note: the S3Object-source branch (driver picks `Document.S3Object`
|
||||
// over inline Bytes when fsEntry has a bucket, and constructs a
|
||||
// TextractClient for that bucket's region) is best exercised by an
|
||||
// *.integration.test.ts against real S3 + Textract — the in-memory
|
||||
// S3 store doesn't deterministically produce a bucket-bearing
|
||||
// fsEntry, and the driver's per-region TextractClient cache leaks
|
||||
// across tests.
|
||||
const writePuterImage = async (): Promise<{
|
||||
actor: Actor;
|
||||
path: string;
|
||||
bytes: Buffer;
|
||||
}> => {
|
||||
const { actor, userId } = await makeUser();
|
||||
const path = `/${actor.user.username}/scan.png`;
|
||||
const bytes = Buffer.from('png-bytes');
|
||||
await server.services.fs.write(userId, {
|
||||
fileMetadata: {
|
||||
path,
|
||||
size: bytes.byteLength,
|
||||
contentType: 'image/png',
|
||||
},
|
||||
fileContent: bytes,
|
||||
});
|
||||
return { actor, path, bytes };
|
||||
};
|
||||
|
||||
it('reads a Puter file straight from S3, never downloading it', async () => {
|
||||
const { actor, path } = await writePuterImage();
|
||||
const s3Read = vi.spyOn(server.stores.s3Object, 'getObjectStream');
|
||||
textractSendMock.mockResolvedValueOnce(sampleTextractResponse);
|
||||
|
||||
await withActor(actor, () =>
|
||||
driver.recognize({ source: { path }, provider: 'aws-textract' }),
|
||||
);
|
||||
|
||||
const sentCmd = textractSendMock.mock.calls[0]![0];
|
||||
expect(sentCmd.input.Document.S3Object).toMatchObject({
|
||||
Name: expect.any(String),
|
||||
});
|
||||
expect(sentCmd.input.Document.Bytes).toBeUndefined();
|
||||
expect(s3Read).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('downloads the bytes only when Textract cannot read the S3 object', async () => {
|
||||
const { actor, path, bytes } = await writePuterImage();
|
||||
const s3Read = vi.spyOn(server.stores.s3Object, 'getObjectStream');
|
||||
textractSendMock
|
||||
.mockRejectedValueOnce(
|
||||
new InvalidS3ObjectException({
|
||||
message: 'unreachable bucket',
|
||||
$metadata: {},
|
||||
}),
|
||||
)
|
||||
.mockResolvedValueOnce(sampleTextractResponse);
|
||||
|
||||
await withActor(actor, () =>
|
||||
driver.recognize({ source: { path }, provider: 'aws-textract' }),
|
||||
);
|
||||
|
||||
expect(s3Read).toHaveBeenCalledTimes(1);
|
||||
expect(
|
||||
textractSendMock.mock.calls[1]![0].input.Document.Bytes,
|
||||
).toEqual(bytes);
|
||||
});
|
||||
|
||||
it('meters one usage line per detected page at the per-page rate from costs.ts', async () => {
|
||||
const { actor } = await makeUser();
|
||||
|
||||
@@ -39,7 +39,11 @@ import {
|
||||
withAiCostFactor,
|
||||
} from '../util/aiCostFactor.js';
|
||||
import { AI_CONCURRENT, AI_RATE_LIMIT } from '../util/aiLimits.js';
|
||||
import { loadFileInput, type LoadedFile } from '../util/fileInput.js';
|
||||
import {
|
||||
type LoadedFile,
|
||||
type ResolvedFileInput,
|
||||
resolveFileInput,
|
||||
} from '../util/fileInput.js';
|
||||
import {
|
||||
type ProviderCatalog,
|
||||
ProviderRegistry,
|
||||
@@ -132,7 +136,7 @@ const OCR_CATALOG: ProviderCatalog = {
|
||||
|
||||
/** Recognizes one input with one provider's model. */
|
||||
type OcrRecognizer = (
|
||||
loaded: LoadedFile,
|
||||
input: ResolvedFileInput,
|
||||
args: RecognizeArgs,
|
||||
model: OcrModel,
|
||||
actor: Actor,
|
||||
@@ -211,20 +215,23 @@ const isDocumentInput = (loaded: LoadedFile): boolean => {
|
||||
/**
|
||||
* The most pages a call can bill, known before the provider runs: a PDF's page
|
||||
* count, or a size-based guess for other documents, capped by a `pages`
|
||||
* selection and by what the provider reads in one call.
|
||||
* selection and by what the provider reads in one call. Reads the input only
|
||||
* for a provider that takes more than one page.
|
||||
*/
|
||||
const estimateOcrPages = (
|
||||
loaded: LoadedFile,
|
||||
const estimateOcrPages = async (
|
||||
input: ResolvedFileInput,
|
||||
provider: OcrProviderId,
|
||||
selection?: unknown,
|
||||
): number => {
|
||||
): Promise<number> => {
|
||||
const maxPages = OCR_MAX_PAGES[provider];
|
||||
if (maxPages <= 1 || !isDocumentInput(loaded)) return 1;
|
||||
if (maxPages <= 1) return 1;
|
||||
const loaded = await input.load();
|
||||
if (!isDocumentInput(loaded)) return 1;
|
||||
// A document we can't count is billed in proportion to its size, reaching
|
||||
// the provider's page limit at its size limit (20 pages per MB for
|
||||
// Mistral). Small files stay cheap; a large one can't pass as one page.
|
||||
const documentPages =
|
||||
countPdfPages(loaded.buffer, maxPages) ??
|
||||
(await countPdfPages(loaded.buffer, maxPages)) ??
|
||||
Math.ceil(
|
||||
(loaded.buffer.length * maxPages) / OCR_MAX_INPUT_BYTES[provider],
|
||||
);
|
||||
@@ -305,8 +312,8 @@ export class OCRDriver extends PuterDriver {
|
||||
};
|
||||
this.#providers.register(
|
||||
'aws-textract',
|
||||
(loaded, _args, model, actor) =>
|
||||
this.#textractRecognize(loaded, model, actor),
|
||||
(input, _args, model, actor) =>
|
||||
this.#textractRecognize(input, model, actor),
|
||||
);
|
||||
}
|
||||
|
||||
@@ -318,8 +325,8 @@ export class OCRDriver extends PuterDriver {
|
||||
}) as unknown as MistralOcrClient;
|
||||
this.#providers.register(
|
||||
'mistral',
|
||||
(loaded, args, model, actor) =>
|
||||
this.#mistralRecognize(loaded, args, model, actor),
|
||||
(input, args, model, actor) =>
|
||||
this.#mistralRecognize(input, args, model, actor),
|
||||
);
|
||||
} catch (e) {
|
||||
console.warn(
|
||||
@@ -355,7 +362,9 @@ export class OCRDriver extends PuterDriver {
|
||||
// reject anyway once it runs.
|
||||
await assertActorHasCredits(this.services.metering, actor, this.config);
|
||||
|
||||
const loaded = await loadFileInput(
|
||||
// A Puter file is resolved and access-checked here but read only if
|
||||
// the provider needs its bytes; Textract reads it straight from S3.
|
||||
const resolved = await resolveFileInput(
|
||||
this.stores,
|
||||
this.services.fs,
|
||||
actor,
|
||||
@@ -366,7 +375,7 @@ export class OCRDriver extends PuterDriver {
|
||||
},
|
||||
);
|
||||
|
||||
return recognize(loaded, args, model, actor);
|
||||
return recognize(resolved, args, model, actor);
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -444,19 +453,18 @@ export class OCRDriver extends PuterDriver {
|
||||
}
|
||||
|
||||
async #textractRecognize(
|
||||
loaded: LoadedFile,
|
||||
input: ResolvedFileInput,
|
||||
model: OcrModel,
|
||||
actor: Actor,
|
||||
) {
|
||||
// Prefer S3 direct source if the file is FS-backed; fall back to raw bytes.
|
||||
const { fsEntry } = input;
|
||||
const s3Info =
|
||||
loaded.fsEntry &&
|
||||
loaded.fsEntry.bucket &&
|
||||
loaded.fsEntry.bucketRegion
|
||||
fsEntry && fsEntry.bucket && fsEntry.bucketRegion
|
||||
? {
|
||||
bucket: loaded.fsEntry.bucket,
|
||||
bucketRegion: loaded.fsEntry.bucketRegion,
|
||||
key: loaded.fsEntry.uuid,
|
||||
bucket: fsEntry.bucket,
|
||||
bucketRegion: fsEntry.bucketRegion,
|
||||
key: fsEntry.uuid,
|
||||
}
|
||||
: null;
|
||||
|
||||
@@ -469,7 +477,7 @@ export class OCRDriver extends PuterDriver {
|
||||
const document =
|
||||
s3Info && useS3
|
||||
? { S3Object: { Bucket: s3Info.bucket, Name: s3Info.key } }
|
||||
: { Bytes: loaded.buffer };
|
||||
: { Bytes: (await input.load()).buffer };
|
||||
return client.send(
|
||||
new DetectDocumentTextCommand({ Document: document }),
|
||||
);
|
||||
@@ -479,7 +487,7 @@ export class OCRDriver extends PuterDriver {
|
||||
const hold = await this.#holdCredits(
|
||||
actor,
|
||||
model.pageUsageType,
|
||||
costPerPage * estimateOcrPages(loaded, model.provider),
|
||||
costPerPage * (await estimateOcrPages(input, model.provider)),
|
||||
);
|
||||
try {
|
||||
let response;
|
||||
@@ -535,11 +543,12 @@ export class OCRDriver extends PuterDriver {
|
||||
// -- Mistral OCR --------------------------------------------------
|
||||
|
||||
async #mistralRecognize(
|
||||
loaded: LoadedFile,
|
||||
input: ResolvedFileInput,
|
||||
args: RecognizeArgs,
|
||||
model: OcrModel,
|
||||
actor: Actor,
|
||||
) {
|
||||
const loaded = await input.load();
|
||||
const payload: Record<string, unknown> = {
|
||||
model: model.id,
|
||||
document: this.#mistralBuildChunk(loaded),
|
||||
@@ -584,7 +593,8 @@ export class OCRDriver extends PuterDriver {
|
||||
const hold = await this.#holdCredits(
|
||||
actor,
|
||||
model.pageUsageType,
|
||||
costPerPage * estimateOcrPages(loaded, model.provider, args.pages),
|
||||
costPerPage *
|
||||
(await estimateOcrPages(input, model.provider, args.pages)),
|
||||
);
|
||||
|
||||
try {
|
||||
|
||||
@@ -24,21 +24,21 @@ import { countPdfPages } from './pdfPages.js';
|
||||
const pdf = (body: string) => Buffer.from(`%PDF-1.7\n${body}`, 'latin1');
|
||||
|
||||
describe('countPdfPages', () => {
|
||||
it('counts the pages of a multi-page PDF', () => {
|
||||
expect(countPdfPages(buildPdf(1))).toBe(1);
|
||||
expect(countPdfPages(buildPdf(130))).toBe(130);
|
||||
it('counts the pages of a multi-page PDF', async () => {
|
||||
expect(await countPdfPages(buildPdf(1))).toBe(1);
|
||||
expect(await countPdfPages(buildPdf(130))).toBe(130);
|
||||
});
|
||||
|
||||
it('reads page objects packed in a compressed object stream', () => {
|
||||
it('reads page objects packed in a compressed object stream', async () => {
|
||||
const packed = buildPdf(42, { objectStream: true });
|
||||
// Nothing about the pages is visible without inflating the stream.
|
||||
expect(packed.includes('/Type /Page')).toBe(false);
|
||||
expect(countPdfPages(packed)).toBe(42);
|
||||
expect(await countPdfPages(packed)).toBe(42);
|
||||
});
|
||||
|
||||
it('reads the total from a nested page tree', () => {
|
||||
it('reads the total from a nested page tree', async () => {
|
||||
expect(
|
||||
countPdfPages(
|
||||
await countPdfPages(
|
||||
pdf(
|
||||
'1 0 obj\n<< /Type /Pages /Kids [2 0 R 3 0 R] /Count 7 >>\nendobj\n' +
|
||||
'2 0 obj\n<< /Type /Pages /Parent 1 0 R /Kids [4 0 R] /Count 4 >>\nendobj\n' +
|
||||
@@ -48,21 +48,21 @@ describe('countPdfPages', () => {
|
||||
).toBe(7);
|
||||
});
|
||||
|
||||
it('counts page objects when the tree understates them', () => {
|
||||
it('counts page objects when the tree understates them', async () => {
|
||||
const pages = Array.from(
|
||||
{ length: 5 },
|
||||
(_, i) => `${i + 2} 0 obj\n<</Type/Page/Parent 1 0 R>>\nendobj\n`,
|
||||
).join('');
|
||||
expect(
|
||||
countPdfPages(
|
||||
await countPdfPages(
|
||||
pdf(`1 0 obj\n<</Type/Pages/Count 1>>\nendobj\n${pages}`),
|
||||
),
|
||||
).toBe(5);
|
||||
});
|
||||
|
||||
it('reads names written with # escapes', () => {
|
||||
it('reads names written with # escapes', async () => {
|
||||
expect(
|
||||
countPdfPages(
|
||||
await countPdfPages(
|
||||
pdf(
|
||||
'1 0 obj\n<< /Type /P#61ges /C#6funt 9 >>\nendobj\n' +
|
||||
'2 0 obj\n<< /Type /P#61ge >>\nendobj\n',
|
||||
@@ -71,11 +71,11 @@ describe('countPdfPages', () => {
|
||||
).toBe(9);
|
||||
});
|
||||
|
||||
it("can't be shrunk by redefining an object number lower", () => {
|
||||
it("can't be shrunk by redefining an object number lower", async () => {
|
||||
// An unreferenced redefinition trying to pass off a smaller tree;
|
||||
// a genuine object number reuse can only raise the count, never lower it.
|
||||
expect(
|
||||
countPdfPages(
|
||||
await countPdfPages(
|
||||
pdf(
|
||||
'1 0 obj\n<< /Type /Pages /Count 3 >>\nendobj\n' +
|
||||
'2 0 obj\n<< /Type /Page >>\nendobj\n' +
|
||||
@@ -87,9 +87,9 @@ describe('countPdfPages', () => {
|
||||
).toBe(3);
|
||||
});
|
||||
|
||||
it('treats an indirect /Count as unknown rather than reading its reference as the total', () => {
|
||||
it('treats an indirect /Count as unknown rather than reading its reference as the total', async () => {
|
||||
expect(
|
||||
countPdfPages(
|
||||
await countPdfPages(
|
||||
pdf(
|
||||
'1 0 obj\n<< /Type /Pages /Kids [2 0 R 3 0 R] /Count 3 0 R >>\nendobj\n',
|
||||
),
|
||||
@@ -97,9 +97,9 @@ describe('countPdfPages', () => {
|
||||
).toBeNull();
|
||||
});
|
||||
|
||||
it('counts a kid with no /Type as long as it has no /Kids of its own', () => {
|
||||
it('counts a kid with no /Type as long as it has no /Kids of its own', async () => {
|
||||
expect(
|
||||
countPdfPages(
|
||||
await countPdfPages(
|
||||
pdf(
|
||||
'1 0 obj\n<< /Type /Pages /Kids [2 0 R 3 0 R] /Count 1 >>\nendobj\n' +
|
||||
'2 0 obj\n<< /Type /Page /Parent 1 0 R >>\nendobj\n' +
|
||||
@@ -109,11 +109,11 @@ describe('countPdfPages', () => {
|
||||
).toBe(2);
|
||||
});
|
||||
|
||||
it("doesn't count untyped objects that aren't page tree kids", () => {
|
||||
it("doesn't count untyped objects that aren't page tree kids", async () => {
|
||||
// Outline items, merged field widgets and name-tree leaves have no
|
||||
// /Type and no /Kids either, but no reader takes them for pages.
|
||||
expect(
|
||||
countPdfPages(
|
||||
await countPdfPages(
|
||||
pdf(
|
||||
'1 0 obj\n<< /Type /Catalog /Pages 2 0 R /Outlines 4 0 R /Names << /Dests 9 0 R >> >>\nendobj\n' +
|
||||
'2 0 obj\n<< /Type /Pages /Kids [3 0 R] /Count 1 >>\nendobj\n' +
|
||||
@@ -130,59 +130,59 @@ describe('countPdfPages', () => {
|
||||
).toBe(1);
|
||||
});
|
||||
|
||||
it('decodes a # escaped /ObjStm type before deciding whether to unpack it', () => {
|
||||
it('decodes a # escaped /ObjStm type before deciding whether to unpack it', async () => {
|
||||
const header = '2 0\n';
|
||||
const packedObject = '<< /Type /Page >>';
|
||||
const objStm =
|
||||
`5 0 obj\n<< /Type /ObjS#74m /N 1 /First ${header.length} >>\n` +
|
||||
`stream\n${header}${packedObject}\nendstream\nendobj\n`;
|
||||
expect(countPdfPages(pdf(objStm))).toBe(1);
|
||||
expect(await countPdfPages(pdf(objStm))).toBe(1);
|
||||
});
|
||||
|
||||
it("refuses an object stream that claims more objects than it's worth counting", () => {
|
||||
it("refuses an object stream that claims more objects than it's worth counting", async () => {
|
||||
const objStm =
|
||||
'1 0 obj\n<< /Type /ObjStm /N 100001 /First 4 >>\nstream\n1 0\nendstream\nendobj\n';
|
||||
expect(countPdfPages(pdf(objStm))).toBeNull();
|
||||
expect(await countPdfPages(pdf(objStm))).toBeNull();
|
||||
});
|
||||
|
||||
it('refuses an oversized /First without scanning the data behind it', () => {
|
||||
it('refuses an oversized /First without scanning the data behind it', async () => {
|
||||
// Without a cap this would run a regex match over ~2MB of digits.
|
||||
const data = '1 '.repeat(1_000_000);
|
||||
const objStm =
|
||||
`1 0 obj\n<< /Type /ObjStm /N 1 /First ${data.length} >>\n` +
|
||||
`stream\n${data}\nendstream\nendobj\n`;
|
||||
const start = Date.now();
|
||||
expect(countPdfPages(pdf(objStm))).toBeNull();
|
||||
expect(await countPdfPages(pdf(objStm))).toBeNull();
|
||||
expect(Date.now() - start).toBeLessThan(500);
|
||||
});
|
||||
|
||||
it('returns null for bytes that are not a PDF', () => {
|
||||
expect(countPdfPages(Buffer.from('PK\x03\x04 not a pdf'))).toBeNull();
|
||||
expect(countPdfPages(Buffer.alloc(0))).toBeNull();
|
||||
it('returns null for bytes that are not a PDF', async () => {
|
||||
expect(await countPdfPages(Buffer.from('PK\x03\x04 not a pdf'))).toBeNull();
|
||||
expect(await countPdfPages(Buffer.alloc(0))).toBeNull();
|
||||
});
|
||||
|
||||
it('returns null when no pages are visible', () => {
|
||||
expect(countPdfPages(pdf('garbage'))).toBeNull();
|
||||
it('returns null when no pages are visible', async () => {
|
||||
expect(await countPdfPages(pdf('garbage'))).toBeNull();
|
||||
});
|
||||
|
||||
it('returns null when an object stream cannot be decoded', () => {
|
||||
it('returns null when an object stream cannot be decoded', async () => {
|
||||
// An encrypted or corrupt stream could hide the page tree.
|
||||
const unreadable = pdf(
|
||||
'1 0 obj\n<< /Type /Page >>\nendobj\n' +
|
||||
'2 0 obj\n<< /Type /ObjStm /N 1 /First 4 /Filter /FlateDecode >>\nstream\n' +
|
||||
'not deflate data\nendstream\nendobj\n',
|
||||
);
|
||||
expect(countPdfPages(unreadable)).toBeNull();
|
||||
expect(await countPdfPages(unreadable)).toBeNull();
|
||||
|
||||
const otherFilter = pdf(
|
||||
'2 0 obj\n<< /Type /ObjStm /N 1 /First 4 /Filter /LZWDecode >>\nstream\n' +
|
||||
'data\nendstream\nendobj\n',
|
||||
);
|
||||
expect(countPdfPages(otherFilter)).toBeNull();
|
||||
expect(await countPdfPages(otherFilter)).toBeNull();
|
||||
});
|
||||
|
||||
it('stops at the limit', () => {
|
||||
expect(countPdfPages(buildPdf(50), 10)).toBe(10);
|
||||
expect(countPdfPages(buildPdf(5), 10)).toBe(5);
|
||||
it('stops at the limit', async () => {
|
||||
expect(await countPdfPages(buildPdf(50), 10)).toBe(10);
|
||||
expect(await countPdfPages(buildPdf(5), 10)).toBe(5);
|
||||
});
|
||||
});
|
||||
@@ -17,7 +17,11 @@
|
||||
* along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
*/
|
||||
|
||||
import { inflateSync } from 'node:zlib';
|
||||
import { promisify } from 'node:util';
|
||||
import { inflate } from 'node:zlib';
|
||||
|
||||
// On the libuv pool, so a large object stream doesn't stall the event loop.
|
||||
const inflateAsync = promisify(inflate);
|
||||
|
||||
/** Hard ceiling on inflated object-stream bytes, regardless of input size. */
|
||||
const MAX_INFLATED_BYTES = 64 * 1024 * 1024;
|
||||
@@ -55,11 +59,11 @@ const decodeNameEscapes = (dict: string): string =>
|
||||
: dict;
|
||||
|
||||
/** The objects packed in an object stream, or null when it can't be read. */
|
||||
const readObjectStream = (
|
||||
const readObjectStream = async (
|
||||
dict: string,
|
||||
data: Buffer,
|
||||
budget: number,
|
||||
): { objects: Array<[number, string]>; inflated: number } | null => {
|
||||
): Promise<{ objects: Array<[number, string]>; inflated: number } | null> => {
|
||||
let content = data;
|
||||
let inflated = 0;
|
||||
if (/\/Filter/.test(dict)) {
|
||||
@@ -67,7 +71,7 @@ const readObjectStream = (
|
||||
if (!FLATE_FILTER.test(dict) || predictor > 1 || budget < 1)
|
||||
return null;
|
||||
try {
|
||||
content = inflateSync(data, { maxOutputLength: budget });
|
||||
content = await inflateAsync(data, { maxOutputLength: budget });
|
||||
} catch {
|
||||
return null;
|
||||
}
|
||||
@@ -111,7 +115,10 @@ const readObjectStream = (
|
||||
* understate either. Null when the bytes show neither, including an encrypted
|
||||
* or undecodable object stream that could hide them. Stops at `limit`.
|
||||
*/
|
||||
export function countPdfPages(pdf: Buffer, limit = Infinity): number | null {
|
||||
export async function countPdfPages(
|
||||
pdf: Buffer,
|
||||
limit = Infinity,
|
||||
): Promise<number | null> {
|
||||
if (!pdf.subarray(0, 1024).includes('%PDF-')) return null;
|
||||
const text = pdf.toString('latin1');
|
||||
// Zlib bombs are a fixed multiple of their compressed size; scale the
|
||||
@@ -180,7 +187,11 @@ export function countPdfPages(pdf: Buffer, limit = Infinity): number | null {
|
||||
start + stream.index + stream[0].length,
|
||||
start + dataEnd,
|
||||
);
|
||||
const packed = readObjectStream(dict, data, inflateBudget - inflated);
|
||||
const packed = await readObjectStream(
|
||||
dict,
|
||||
data,
|
||||
inflateBudget - inflated,
|
||||
);
|
||||
if (!packed) return null;
|
||||
inflated += packed.inflated;
|
||||
for (const [id, objectDict] of packed.objects) {
|
||||
|
||||
@@ -24,7 +24,11 @@ import { runWithContext } from '../../core/context.js';
|
||||
import { PuterServer } from '../../server.js';
|
||||
import { setupTestServer } from '../../testUtil.js';
|
||||
import { generateDefaultFsentries } from '../../util/userProvisioning.js';
|
||||
import { inferFilenameFromUrlOrPath, loadFileInput } from './fileInput.js';
|
||||
import {
|
||||
inferFilenameFromUrlOrPath,
|
||||
loadFileInput,
|
||||
resolveFileInput,
|
||||
} from './fileInput.js';
|
||||
|
||||
// Web inputs leave the process through secureFetch; stub that boundary so
|
||||
// the tests control the remote response without real network access.
|
||||
@@ -418,6 +422,71 @@ describe('loadFileInput FS path', () => {
|
||||
});
|
||||
});
|
||||
|
||||
// ── resolveFileInput ───────────────────────────────────────────────
|
||||
|
||||
describe('resolveFileInput', () => {
|
||||
it('checks access to a Puter file without reading it, then reads it once on load', async () => {
|
||||
const { actor, userId } = await makeUser();
|
||||
const path = `/${actor.user!.username!}/Documents/lazy.txt`;
|
||||
const body = Buffer.from('read me later');
|
||||
const entry = await withActor(actor, () =>
|
||||
writeFile(userId, path, body, 'text/plain'),
|
||||
);
|
||||
const s3Read = vi.spyOn(server.stores.s3Object, 'getObjectStream');
|
||||
try {
|
||||
const resolved = await withActor(actor, () =>
|
||||
resolveFileInput(
|
||||
server.stores,
|
||||
server.services.fs,
|
||||
actor,
|
||||
path,
|
||||
),
|
||||
);
|
||||
expect(resolved.filename).toBe('lazy.txt');
|
||||
expect(resolved.fsEntry?.uuid).toBe(entry.uuid);
|
||||
expect(s3Read).not.toHaveBeenCalled();
|
||||
|
||||
const [first, second] = await Promise.all([
|
||||
resolved.load(),
|
||||
resolved.load(),
|
||||
]);
|
||||
expect(first.buffer.equals(body)).toBe(true);
|
||||
expect(second).toBe(first);
|
||||
expect(s3Read).toHaveBeenCalledTimes(1);
|
||||
} finally {
|
||||
s3Read.mockRestore();
|
||||
}
|
||||
});
|
||||
|
||||
it('refuses a Puter file over maxBytes from its recorded size, before any read', async () => {
|
||||
const { actor, userId } = await makeUser();
|
||||
const path = `/${actor.user!.username!}/Documents/big.bin`;
|
||||
await withActor(actor, () =>
|
||||
writeFile(userId, path, Buffer.alloc(2048, 0x42)),
|
||||
);
|
||||
const s3Read = vi.spyOn(server.stores.s3Object, 'getObjectStream');
|
||||
try {
|
||||
await expect(
|
||||
withActor(actor, () =>
|
||||
resolveFileInput(
|
||||
server.stores,
|
||||
server.services.fs,
|
||||
actor,
|
||||
path,
|
||||
{ maxBytes: 64 },
|
||||
),
|
||||
),
|
||||
).rejects.toMatchObject({
|
||||
statusCode: 413,
|
||||
legacyCode: 'storage_limit_reached',
|
||||
});
|
||||
expect(s3Read).not.toHaveBeenCalled();
|
||||
} finally {
|
||||
s3Read.mockRestore();
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
// ── loadFileInput web URL ──────────────────────────────────────────
|
||||
|
||||
/** A web stream that yields `chunks` one pull at a time, counting pulls. */
|
||||
|
||||
@@ -62,6 +62,14 @@ export interface LoadedFile {
|
||||
} | null;
|
||||
}
|
||||
|
||||
/** A file input resolved and access-checked, its bytes read on demand. */
|
||||
export interface ResolvedFileInput {
|
||||
filename: string;
|
||||
fsEntry: LoadedFile['fsEntry'];
|
||||
/** Reads the bytes; later calls share the first read. */
|
||||
load(): Promise<LoadedFile>;
|
||||
}
|
||||
|
||||
/**
|
||||
* A Puter FS reference as drivers receive it: a path string or `{ path?, uid?,
|
||||
* uuid? }`.
|
||||
@@ -87,6 +95,28 @@ export async function loadFileInput(
|
||||
input: unknown,
|
||||
options: { maxBytes?: number; acceptWebInput?: true } = {},
|
||||
): Promise<LoadedFile> {
|
||||
const resolved = await resolveFileInput(
|
||||
stores,
|
||||
fsService,
|
||||
actor,
|
||||
input,
|
||||
options,
|
||||
);
|
||||
return resolved.load();
|
||||
}
|
||||
|
||||
/**
|
||||
* {@link loadFileInput} without the read: a Puter FS reference is resolved and
|
||||
* access-checked, and its bytes are read only when `load()` is called. Data and
|
||||
* web URLs carry their bytes, so those are read up front.
|
||||
*/
|
||||
export async function resolveFileInput(
|
||||
stores: FileInputStores,
|
||||
fsService: FSService,
|
||||
actor: Actor,
|
||||
input: unknown,
|
||||
options: { maxBytes?: number; acceptWebInput?: true } = {},
|
||||
): Promise<ResolvedFileInput> {
|
||||
if (!input) {
|
||||
throw new HttpError(400, 'Missing file input', {
|
||||
legacyCode: 'bad_request',
|
||||
@@ -103,12 +133,12 @@ export async function loadFileInput(
|
||||
});
|
||||
const buffer = dataUriBytes(uri);
|
||||
assertMax(buffer, options.maxBytes);
|
||||
return {
|
||||
return alreadyLoaded({
|
||||
buffer,
|
||||
filename: filenameFromMime(uri.mimeType),
|
||||
mimeType: uri.mimeType,
|
||||
fsEntry: null,
|
||||
};
|
||||
});
|
||||
}
|
||||
|
||||
// Web URL — fetch via SSRF-guarded secureFetch.
|
||||
@@ -131,31 +161,51 @@ export async function loadFileInput(
|
||||
contentType?.split(';')[0]?.trim() ||
|
||||
mimeFromName(input) ||
|
||||
'application/octet-stream';
|
||||
return {
|
||||
return alreadyLoaded({
|
||||
buffer,
|
||||
filename: inferFilenameFromUrlOrPath(input),
|
||||
mimeType: mime,
|
||||
fsEntry: null,
|
||||
};
|
||||
});
|
||||
}
|
||||
|
||||
// Path string or object reference → resolve into FSEntry, then S3 read.
|
||||
const opened = await openFileInputStream(
|
||||
// Path string or object reference → resolve into FSEntry; S3 read on load.
|
||||
const entry = await resolveFileInputEntry(
|
||||
stores,
|
||||
fsService,
|
||||
actor,
|
||||
input as FileInputRef,
|
||||
options,
|
||||
);
|
||||
const buffer = await collectStream(opened.body, options.maxBytes);
|
||||
if (options.maxBytes && (entry.size ?? 0) > options.maxBytes) {
|
||||
throw tooLarge(options.maxBytes);
|
||||
}
|
||||
let loading: Promise<LoadedFile> | undefined;
|
||||
return {
|
||||
buffer,
|
||||
filename: opened.filename,
|
||||
mimeType: opened.mimeType,
|
||||
fsEntry: opened.fsEntry,
|
||||
filename: entry.name,
|
||||
fsEntry: fsEntryInfo(entry),
|
||||
load: () =>
|
||||
(loading ??= (async () => {
|
||||
const opened = await openEntryStream(stores, entry, options);
|
||||
const buffer = await collectStream(
|
||||
opened.body,
|
||||
options.maxBytes,
|
||||
);
|
||||
return {
|
||||
buffer,
|
||||
filename: opened.filename,
|
||||
mimeType: opened.mimeType,
|
||||
fsEntry: opened.fsEntry,
|
||||
};
|
||||
})()),
|
||||
};
|
||||
}
|
||||
|
||||
const alreadyLoaded = (loaded: LoadedFile): ResolvedFileInput => ({
|
||||
filename: loaded.filename,
|
||||
fsEntry: loaded.fsEntry,
|
||||
load: () => Promise.resolve(loaded),
|
||||
});
|
||||
|
||||
/**
|
||||
* Resolve a Puter FS reference to its entry, with the checks every driver read
|
||||
* needs: it must be a file (not a directory, symlink or shortcut) and `actor`
|
||||
@@ -218,14 +268,25 @@ export async function openFileInputStream(
|
||||
options: { maxBytes?: number } = {},
|
||||
): Promise<OpenedFileInput> {
|
||||
const entry = await resolveFileInputEntry(stores, fsService, actor, input);
|
||||
const fsEntry = {
|
||||
uuid: entry.uuid,
|
||||
path: entry.path,
|
||||
bucket: entry.bucket,
|
||||
bucketRegion: entry.bucketRegion,
|
||||
size: entry.size,
|
||||
sqlId: entry.id,
|
||||
};
|
||||
return openEntryStream(stores, entry, options);
|
||||
}
|
||||
|
||||
const fsEntryInfo = (entry: FSEntry): NonNullable<LoadedFile['fsEntry']> => ({
|
||||
uuid: entry.uuid,
|
||||
path: entry.path,
|
||||
bucket: entry.bucket,
|
||||
bucketRegion: entry.bucketRegion,
|
||||
size: entry.size,
|
||||
sqlId: entry.id,
|
||||
});
|
||||
|
||||
/** {@link openFileInputStream} for an entry that is already resolved. */
|
||||
async function openEntryStream(
|
||||
stores: FileInputStores,
|
||||
entry: FSEntry,
|
||||
options: { maxBytes?: number },
|
||||
): Promise<OpenedFileInput> {
|
||||
const fsEntry = fsEntryInfo(entry);
|
||||
// Empty files (created via `touch`) have no backing S3 object —
|
||||
// getObjectStream would throw NoSuchKey, so return empty content.
|
||||
if (hasNoBackingS3Object(entry)) {
|
||||
|
||||
Reference in new issue
Block a user