diff --git a/src/lib/features/filesystem/cache.svelte.ts b/src/lib/features/filesystem/cache.svelte.ts index bb2585d..38d3e32 100644 --- a/src/lib/features/filesystem/cache.svelte.ts +++ b/src/lib/features/filesystem/cache.svelte.ts @@ -11,6 +11,7 @@ interface FilesystemCacheOptions { interface Entry { value?: V; promise?: Promise; + isInvalidated: boolean; } // Caches filesystem entries as reactive objects which are updated in place, so that every view @@ -24,7 +25,7 @@ export class FilesystemCache { private getOrCreateEntry(key: K) { let entry = this.map.get(key); if (!entry) { - const newEntry: Entry = $state({}); + const newEntry: Entry = $state({ isInvalidated: false }); entry = newEntry; this.map.set(key, entry); } @@ -40,6 +41,15 @@ export class FilesystemCache { return entry.value!; } + private finish(key: K, entry: Entry) { + entry.promise = undefined; + // Mutations may invalidate the entry while it is being revalidated, in which case it is queried + // again unless it was evicted or replaced in the meantime + if (entry.isInvalidated && this.map.get(key) === entry) { + void Promise.resolve(this.get(key)).catch(() => {}); + } + } + // Returns the cached value synchronously if available, and revalidates it in the background get(key: K, options?: { fetchFromServer?: (cachedValue: V | undefined) => Promise }) { return untrack(() => { @@ -49,6 +59,7 @@ export class FilesystemCache { const { promise, resolve, reject } = Promise.withResolvers(); promise.catch(() => {}); // Rejections are handled by callers waiting for the promise entry.promise = promise; + entry.isInvalidated = false; (async () => { let cachedValue = entry.value && ($state.snapshot(entry.value) as V); @@ -67,9 +78,7 @@ export class FilesystemCache { console.error(e); reject(e); }) - .finally(() => { - entry.promise = undefined; - }); + .finally(() => this.finish(key, entry)); return entry.value ?? promise; }); @@ -77,20 +86,23 @@ export class FilesystemCache { bulkGet(keys: Set) { return untrack(() => { + const entries = new Map>(); const pending = new Map>(); for (const key of keys) { const entry = this.getOrCreateEntry(key); + entries.set(key, entry); if (entry.promise) continue; const resolvers = Promise.withResolvers(); resolvers.promise.catch(() => {}); entry.promise = resolvers.promise; + entry.isInvalidated = false; pending.set(key, resolvers); } const settleAll = (values: Map) => { values.forEach((value, key) => { - pending.get(key)?.resolve(this.settle(this.map.get(key)!, value)); + pending.get(key)?.resolve(this.settle(entries.get(key)!, value)); }); }; @@ -98,12 +110,12 @@ export class FilesystemCache { const pendingKeys = [...pending.keys()]; settleAll( await this.options.bulkFetchFromIndexedDB!( - new Set(pendingKeys.filter((key) => !this.map.get(key)!.value)), + new Set(pendingKeys.filter((key) => !entries.get(key)!.value)), ), ); settleAll( await this.options.bulkFetchFromServer!( - new Map(pendingKeys.map((key) => [key, { cachedValue: this.map.get(key)!.value }])), + new Map(pendingKeys.map((key) => [key, { cachedValue: entries.get(key)!.value }])), ), ); })() @@ -112,14 +124,11 @@ export class FilesystemCache { pending.forEach(({ reject }) => reject(e)); }) .finally(() => { - pending.forEach((_, key) => { - this.map.get(key)!.promise = undefined; - }); + pending.forEach((_, key) => this.finish(key, entries.get(key)!)); }); - const makeResult = () => new Map([...keys].map((key) => [key, this.map.get(key)!.value!])); - const waiting = [...keys] - .map((key) => this.map.get(key)!) + const makeResult = () => new Map([...entries].map(([key, entry]) => [key, entry.value!])); + const waiting = [...entries.values()] .filter((entry) => !entry.value) .map((entry) => entry.promise!); return waiting.length > 0 ? Promise.all(waiting).then(makeResult) : makeResult(); @@ -128,7 +137,9 @@ export class FilesystemCache { // Revalidates the entry if it has been loaded invalidate(key: K) { - if (this.map.has(key)) { + const entry = this.map.get(key); + if (entry) { + entry.isInvalidated = true; void Promise.resolve(this.get(key)).catch(() => {}); } } diff --git a/src/lib/server/db/repositories/file.ts b/src/lib/server/db/repositories/file.ts index 52b6c2d..1ccaa38 100644 --- a/src/lib/server/db/repositories/file.ts +++ b/src/lib/server/db/repositories/file.ts @@ -99,17 +99,21 @@ export const lockFile = async (trx: Tx, userId: number, fileId: number) => { return !!file; }; -// Locks the file in share mode so that its DEK isn't changed until the transaction ends. Returns -// its DEK version if it exists. -export const lockFileDekVersion = async (trx: Tx, userId: number, fileId: number) => { - const file = await trx +// Locks the file in the given mode so that its DEK isn't changed until the transaction ends. +// Returns its DEK version if it exists. +export const lockFileDekVersion = async ( + trx: Tx, + userId: number, + fileId: number, + mode: "share" | "update" = "share", +) => { + const query = trx .selectFrom("file") .select("data_encryption_key_version") .where("id", "=", fileId) .where("user_id", "=", userId) - .limit(1) - .forShare() - .executeTakeFirst(); + .limit(1); + const file = await (mode === "update" ? query.forUpdate() : query.forShare()).executeTakeFirst(); return file?.data_encryption_key_version ?? null; }; diff --git a/src/lib/server/services/common.ts b/src/lib/server/services/common.ts index 254f270..ce83b3c 100644 --- a/src/lib/server/services/common.ts +++ b/src/lib/server/services/common.ts @@ -26,8 +26,9 @@ export const ensureFileDekVersion = async ( userId: number, fileId: number, dekVersion: Date, + mode: "share" | "update" = "share", ) => { - const currentDekVersion = await FileRepo.lockFileDekVersion(trx, userId, fileId); + const currentDekVersion = await FileRepo.lockFileDekVersion(trx, userId, fileId, mode); if (!currentDekVersion) { throw new AppError("NOT_FOUND", "File not found"); } else if (currentDekVersion.getTime() !== dekVersion.getTime()) { diff --git a/src/lib/server/services/upload.ts b/src/lib/server/services/upload.ts index 7b0360f..f392396 100644 --- a/src/lib/server/services/upload.ts +++ b/src/lib/server/services/upload.ts @@ -15,6 +15,7 @@ import { deleteFile, hashFile, reserveFile, + withTemporaryChunk, writeChunkAtOffset, } from "#lib/server/storage/index.ts"; import { @@ -126,11 +127,22 @@ export const uploadChunk = async ( throw new AppError("CONFLICT", "Chunk already uploaded"); } - // Concurrent uploads of the same chunk may corrupt the file, which will be detected by the - // hash verification on completion const isLastChunk = chunkIndex === session.totalChunks; - await writeChunkAtOffset(session.path, encChunkStream, chunkIndex, isLastChunk); - await UploadRepo.markChunkAsUploaded(db, sessionId, chunkIndex); + await withTemporaryChunk(encChunkStream, isLastChunk, (readChunk) => + transaction(async (trx) => { + // The chunk is received before locking the session, which is checked again so that late + // requests can't overwrite completed chunks or files even if they passed the checks above + const lockedSession = await UploadRepo.lockUploadSession(trx, sessionId, userId); + if (!lockedSession) { + throw new AppError("NOT_FOUND", "Invalid upload id"); + } else if (isChunkUploaded(lockedSession.bitmap, chunkIndex)) { + throw new AppError("CONFLICT", "Chunk already uploaded"); + } + + await writeChunkAtOffset(lockedSession.path, readChunk(), chunkIndex, isLastChunk); + await UploadRepo.markChunkAsUploaded(trx, sessionId, chunkIndex); + }), + ); }; const getCompletedUploadSession = async ( @@ -212,7 +224,8 @@ export const completeFileThumbnailUpload = async (userId: number, uploadId: stri throw new AppError("NOT_FOUND", "Invalid upload id"); } - await ensureFileDekVersion(trx, userId, session.fileId, session.dekVersion); + // Locks the file exclusively, as there may be no thumbnail row to lock yet + await ensureFileDekVersion(trx, userId, session.fileId, session.dekVersion, "update"); const oldPath = await MediaRepo.lockFileThumbnailPath(trx, session.fileId); await MediaRepo.upsertFileThumbnail(trx, session.fileId, session.path, null); await UploadRepo.deleteUploadSession(trx, uploadId); diff --git a/src/lib/server/storage/chunks.ts b/src/lib/server/storage/chunks.ts index 8e25739..fab7231 100644 --- a/src/lib/server/storage/chunks.ts +++ b/src/lib/server/storage/chunks.ts @@ -1,10 +1,37 @@ +import { randomUUID } from "crypto"; +import { createReadStream } from "fs"; import { open } from "fs/promises"; +import { tmpdir } from "os"; +import { join } from "path"; import type { Readable } from "stream"; import { AppError } from "#lib/server/errors.ts"; import { ENCRYPTION_OVERHEAD, ENCRYPTED_CHUNK_SIZE } from "#lib/shared/format/index.ts"; +import { deleteFile } from "./files.ts"; const invalidChunkSize = () => new AppError("BAD_REQUEST", "Invalid chunk size"); +// Receives and validates a chunk before it is copied into the upload's file. Temporary chunks +// aren't placed in the persistent storage, and are removed even if receiving or applying fails. +export const withTemporaryChunk = async ( + encChunkStream: Readable, + isLastChunk: boolean, + callback: (readChunk: () => Readable) => Promise, +) => { + const path = join(tmpdir(), `arkvault-chunk-${randomUUID()}`); + // Created exclusively with restricted permissions, as `writeChunkAtOffset` writes into existing + // files only + const file = await open(path, "wx", 0o600); + try { + await file.close(); + await writeChunkAtOffset(path, encChunkStream, 1, isLastChunk); + // A new stream is created for each call, as the chunk may be read several times (e.g. when + // retrying a transaction) + return await callback(() => createReadStream(path)); + } finally { + await deleteFile(path).catch(console.error); + } +}; + export const writeChunkAtOffset = async ( path: string, encChunkStream: Readable, diff --git a/src/lib/shared/format/stream.ts b/src/lib/shared/format/stream.ts index 8d1c9dc..d8a8325 100644 --- a/src/lib/shared/format/stream.ts +++ b/src/lib/shared/format/stream.ts @@ -40,7 +40,7 @@ const createByteQueue = () => { // Decrypts the encrypted chunks covering the plaintext `range` (as returned by // `getEncryptedRange`) into the plaintext of the range. `onChunk` receives each decrypted chunk -// as a whole, along with its index in the file. +// as a whole, along with its index in the file. Async callbacks are awaited to apply backpressure. export const createDecryptionStream = ({ range, firstChunkIndex = 0, @@ -54,7 +54,7 @@ export const createDecryptionStream = ({ chunkCount: number; dataKey: CryptoKey; isLegacy: boolean; - onChunk?: (chunkIndex: number, chunk: Uint8Array) => void; + onChunk?: (chunkIndex: number, chunk: Uint8Array) => void | Promise; }) => { const queue = createByteQueue(); let chunkIndex = 0; @@ -64,7 +64,7 @@ export const createDecryptionStream = ({ encryptedSize: number, ) => { const decrypted = new Uint8Array(await decryptChunk(queue.take(encryptedSize).buffer, dataKey)); - onChunk?.(firstChunkIndex + chunkIndex, decrypted); + await onChunk?.(firstChunkIndex + chunkIndex, decrypted); const isFirst = chunkIndex === 0; const isLast = chunkIndex === chunkCount - 1; diff --git a/src/service-worker/handlers/decrypt-file.ts b/src/service-worker/handlers/decrypt-file.ts index b014164..b3d8c4b 100644 --- a/src/service-worker/handlers/decrypt-file.ts +++ b/src/service-worker/handlers/decrypt-file.ts @@ -168,10 +168,10 @@ async function* fetchSegment( chunkCount: segment.lastChunkIndex - segment.firstChunkIndex + 1, dataKey: metadata.dataKey, isLegacy: metadata.isLegacy, - // Chunks are written atomically, so the worker being terminated doesn't corrupt the cache - onChunk: (chunkIndex, chunk) => { - storeCachedChunk(request.fileId, chunkIndex, chunk, layout).catch(console.error); - }, + // The atomic write is awaited so that slow storage doesn't queue up the whole plaintext, + // and its failure doesn't fail the download + onChunk: (chunkIndex, chunk) => + storeCachedChunk(request.fileId, chunkIndex, chunk, layout).catch(console.error), }), ); } diff --git a/tests/server/upload.test.ts b/tests/server/upload.test.ts index 4a6e057..0c519bf 100644 --- a/tests/server/upload.test.ts +++ b/tests/server/upload.test.ts @@ -4,6 +4,7 @@ import { beforeEach, describe, expect, it } from "vitest"; import { UploadRepo, db } from "#lib/server/db/index.ts"; import { FileService, UploadService } from "#lib/server/services/index.ts"; import { isChunkUploaded } from "#lib/server/services/upload.ts"; +import { hashFile } from "#lib/server/storage/index.ts"; import { ciphertext, createUser, newEntryParams, resetDatabase } from "../helpers/database.ts"; beforeEach(resetDatabase); @@ -38,6 +39,61 @@ describe("upload session bitmap", () => { }); describe("upload completion", () => { + it.each([false, true])( + "rejects a late duplicate without changing the file (completed: %s)", + async (completeBeforeDuplicate) => { + const { userId } = await createUser(); + const { uploadId } = await UploadService.startFileUpload({ + ...newEntryParams(userId, "root"), + totalChunks: 1, + hskVersion: null, + contentType: "image/png", + encCreatedAt: null, + encLastModifiedAt: ciphertext(), + }); + const session = (await UploadRepo.getUploadSession(db, uploadId, userId))!; + const started = Promise.withResolvers(); + const release = Promise.withResolvers(); + const duplicate = UploadService.uploadChunk( + userId, + uploadId, + 1, + Readable.from( + (async function* () { + started.resolve(); + await release.promise; + yield Buffer.alloc(100, 2); + })(), + ), + ); + await started.promise; + + try { + const chunk = Buffer.alloc(100, 1); + const hash = createHash("sha256").update(chunk).digest("base64"); + await UploadService.uploadChunk(userId, uploadId, 1, Readable.from([chunk])); + if (completeBeforeDuplicate) { + await UploadService.completeFileUpload(userId, uploadId, null, hash); + } + + const rejected = expect(duplicate).rejects.toMatchObject({ + code: completeBeforeDuplicate ? "NOT_FOUND" : "CONFLICT", + }); + release.resolve(); + await rejected; + expect(await hashFile(session.path)).toBe(hash); + + if (!completeBeforeDuplicate) { + await UploadService.completeFileUpload(userId, uploadId, null, hash); + } + expect(await FileService.getAllFileIds(userId)).toHaveLength(1); + } finally { + release.resolve(); + await duplicate.catch(() => {}); + } + }, + ); + it("registers the file only once when completed concurrently", async () => { const { userId } = await createUser(); const { uploadId } = await UploadService.startFileUpload({ diff --git a/tests/unit/filesystem-cache.test.ts b/tests/unit/filesystem-cache.test.ts index ecbe1b4..ea20fdf 100644 --- a/tests/unit/filesystem-cache.test.ts +++ b/tests/unit/filesystem-cache.test.ts @@ -30,6 +30,92 @@ const createCache = (options?: { const flush = () => new Promise((resolve) => setTimeout(resolve, 0)); describe("FilesystemCache", () => { + it("revalidates again after invalidation during a pending request, coalescing repeated calls", async () => { + const pending = Promise.withResolvers(); + const fetchFromServer = vi + .fn<(id: number) => Promise>() + .mockImplementationOnce(() => pending.promise) + .mockResolvedValue({ id: 1, name: "new" }); + const cache = new FilesystemCache({ + fetchFromIndexedDB: async () => undefined, + fetchFromServer, + }); + const loading = cache.get(1); + await vi.waitFor(() => expect(fetchFromServer).toHaveBeenCalledTimes(1)); + + cache.invalidate(1); + cache.invalidate(1); + cache.invalidate(1); + pending.resolve({ id: 1, name: "old" }); + const entry = await loading; + await vi.waitFor(() => expect(entry.name).toBe("new")); + expect(fetchFromServer).toHaveBeenCalledTimes(2); + }); + + it("revalidates only invalidated entries after a pending bulk request", async () => { + const pending = Promise.withResolvers>(); + const fetchFromServer = vi.fn(async (id: number) => ({ id, name: "new" })); + const bulkFetchFromServer = vi.fn(() => pending.promise); + const cache = new FilesystemCache({ + fetchFromIndexedDB: async () => undefined, + fetchFromServer, + bulkFetchFromIndexedDB: async () => new Map(), + bulkFetchFromServer, + }); + const loading = cache.bulkGet(new Set([1, 2])); + await vi.waitFor(() => expect(bulkFetchFromServer).toHaveBeenCalledTimes(1)); + cache.invalidate(1); + pending.resolve(new Map([1, 2].map((id) => [id, { id, name: "old" }]))); + + const entries = await loading; + await vi.waitFor(() => expect(entries.get(1)?.name).toBe("new")); + expect(entries.get(2)?.name).toBe("old"); + expect(fetchFromServer).toHaveBeenCalledTimes(1); + }); + + it("preserves invalidation when a pending request fails", async () => { + vi.spyOn(console, "error").mockImplementation(() => {}); + const pending = Promise.withResolvers(); + const fetchFromServer = vi + .fn<() => Promise>() + .mockImplementationOnce(() => pending.promise) + .mockResolvedValue({ id: 1, name: "new" }); + const cache = new FilesystemCache({ + fetchFromIndexedDB: async () => ({ id: 1, name: "local" }), + fetchFromServer, + }); + const entry = await cache.get(1); + await vi.waitFor(() => expect(fetchFromServer).toHaveBeenCalledTimes(1)); + cache.invalidate(1); + pending.reject(new Error("Failed")); + + await vi.waitFor(() => expect(entry.name).toBe("new")); + expect(fetchFromServer).toHaveBeenCalledTimes(2); + }); + + it("doesn't revalidate evicted entries after a pending bulk request", async () => { + const pending = Promise.withResolvers>(); + const fetchFromServer = vi.fn(async (id: number) => ({ id, name: "new" })); + const bulkFetchFromServer = vi.fn(() => pending.promise); + const cache = new FilesystemCache({ + fetchFromIndexedDB: async () => undefined, + fetchFromServer, + bulkFetchFromIndexedDB: async () => new Map(), + bulkFetchFromServer, + }); + const loading = cache.bulkGet(new Set([1])); + await vi.waitFor(() => expect(bulkFetchFromServer).toHaveBeenCalledTimes(1)); + cache.invalidate(1); + cache.evict(1); + const replacement = await cache.get(1); + pending.resolve(new Map([[1, { id: 1, name: "old" }]])); + await loading; + await flush(); + + expect(replacement.name).toBe("new"); + expect(fetchFromServer).toHaveBeenCalledTimes(1); + }); + it("resolves with the local value first and updates it in place", async () => { const { cache } = createCache({ local: { 1: "local" }, server: { 1: "server" } });