Skip to content

Commit 3aa8b9f

Browse files
committed
fix(account-deletion): close the storage cleanup gaps review found
- integration fixture rows get unique names: the active-name unique index on workspace_files rejected the second 'fixture.txt' row, failing seed() - the reference guard reads stored asset URLs unfiltered and decodes them, so a URL that percent-encodes the key's characters still keeps its object; the world-readable-key gate still decides whether any URL column is read - enqueueAccountStorageCleanupInTx re-reads the account's rows written since the pre-transaction collection (via the user_id and uploaded_by indexes) under the account lock, so a late Copilot or knowledge-base upload is queued - local listings walk keys in sorted order and page after the last key, so a workspace sweep advances past a full page of kept objects
1 parent 9dee1b4 commit 3aa8b9f

5 files changed

Lines changed: 275 additions & 172 deletions

File tree

‎apps/sim/lib/uploads/core/prefix-sweep.ts‎

Lines changed: 40 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
1-
import { opendir, rmdir } from 'node:fs/promises'
1+
import { readdir, rmdir } from 'node:fs/promises'
22
import { dirname, join } from 'node:path'
33
import { describeError } from '@sim/utils/errors'
44
import { isRecordLike } from '@sim/utils/object'
5+
import { compareStrings } from '@sim/utils/string'
56
import type { OutboxEventContext } from '@/lib/core/outbox/service'
67
import {
78
getStorageConfig,
@@ -40,22 +41,50 @@ async function removeEmptyDirectory(path: string): Promise<void> {
4041
}
4142
}
4243

43-
async function localKeys(prefix: string, signal: AbortSignal): Promise<string[]> {
44+
/** Orders keys path segment by segment, the order a sorted depth-first walk visits them in. */
45+
function compareKeyPaths(left: readonly string[], right: readonly string[]): number {
46+
for (let index = 0; index < Math.min(left.length, right.length); index++) {
47+
const order = compareStrings(left[index], right[index])
48+
if (order !== 0) return order
49+
}
50+
return left.length - right.length
51+
}
52+
53+
/**
54+
* The first local keys under `prefix` in {@link compareKeyPaths} order, after `startAfter` when
55+
* given, so a caller can page past objects it keeps. Prunes directories it walks to the end empty.
56+
*/
57+
async function localKeys(
58+
prefix: string,
59+
signal: AbortSignal,
60+
startAfter?: string
61+
): Promise<string[]> {
62+
const cursor = startAfter?.split('/')
4463
const keys: string[] = []
4564
async function visit(key: string): Promise<void> {
4665
signal.throwIfAborted()
47-
let directory
66+
let entries
4867
try {
49-
directory = await opendir(join(UPLOAD_DIR_SERVER, key))
68+
entries = await readdir(join(UPLOAD_DIR_SERVER, key), { withFileTypes: true })
5069
} catch (error) {
5170
if (describeError(error).code === 'ENOENT') return
5271
throw error
5372
}
54-
for await (const entry of directory) {
73+
entries.sort((left, right) => compareStrings(left.name, right.name))
74+
for (const entry of entries) {
5575
signal.throwIfAborted()
5676
const child = `${key}${entry.name}`
57-
if (entry.isDirectory()) await visit(`${child}/`)
58-
else keys.push(child)
77+
if (entry.isDirectory()) {
78+
const segments = child.split('/')
79+
if (
80+
!cursor ||
81+
segments.every((segment, index) => cursor[index] === segment) ||
82+
compareKeyPaths(segments, cursor) > 0
83+
)
84+
await visit(`${child}/`)
85+
} else if (!cursor || compareKeyPaths(child.split('/'), cursor) > 0) {
86+
keys.push(child)
87+
}
5988
if (keys.length >= STORAGE_KEY_PAGE_SIZE) return
6089
}
6190
await removeEmptyDirectory(join(UPLOAD_DIR_SERVER, key))
@@ -77,8 +106,8 @@ function keyPage(keys: string[], nextPageToken: string | undefined): StorageKeyP
77106
}
78107

79108
/**
80-
* Lists one page of object keys under `prefix` in the bucket serving `context`. A local listing
81-
* has no cursor: it returns the first keys still on disk, so callers advance by deleting them.
109+
* Lists one page of object keys under `prefix` in the bucket serving `context`. A local listing's
110+
* page token is the page's last key, so a caller keeping objects still advances past them.
82111
*/
83112
export async function listStorageKeyPage(
84113
prefix: string,
@@ -141,8 +170,8 @@ export async function listStorageKeyPage(
141170
: undefined
142171
)
143172
}
144-
const keys = await localKeys(prefix, signal)
145-
return { keys, done: keys.length < STORAGE_KEY_PAGE_SIZE }
173+
const keys = await localKeys(prefix, signal, pageToken)
174+
return keyPage(keys, keys.length < STORAGE_KEY_PAGE_SIZE ? undefined : keys[keys.length - 1])
146175
}
147176

148177
/**

‎apps/sim/lib/users/storage-reference-guard.ts‎

Lines changed: 26 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -12,14 +12,10 @@ import {
1212
workspaceFileVersion,
1313
workspace as workspaceTable,
1414
} from '@sim/db/schema'
15-
import { chunkArray } from '@sim/utils/helpers'
1615
import { isRecordLike } from '@sim/utils/object'
17-
import { inArray, isNotNull, like, or, type SQL, type SQLWrapper, sql } from 'drizzle-orm'
16+
import { inArray, isNotNull, like, or } from 'drizzle-orm'
1817
import { extractStorageKey, isPublicStorageKey } from '@/lib/uploads/utils/file'
1918

20-
const MIN_NEEDLE_LENGTH = 4
21-
const NEEDLES_PER_QUERY = 50
22-
2319
/** Every string inside a URL column or a JSON settings blob, including JSON held as a string. */
2420
function collectStrings(value: unknown, into: string[]): void {
2521
if (typeof value === 'string') {
@@ -81,54 +77,40 @@ function keysNamedByUrl(value: string): string[] {
8177
return [...keys]
8278
}
8379

84-
/**
85-
* The longest run of URL-safe characters after a key's root. Percent-encoding never splits such a
86-
* run, so any stored URL naming the key contains it; `null` when the key has none long enough.
87-
*/
88-
function urlNeedle(key: string): string | null {
89-
let best = ''
90-
for (const run of key.slice(key.indexOf('/') + 1).match(/[A-Za-z0-9_-]+/g) ?? [])
91-
if (run.length > best.length) best = run
92-
return best.length >= MIN_NEEDLE_LENGTH ? best : null
93-
}
94-
9580
/**
9681
* Keys named by the columns that hold a stored asset as a URL rather than a row:
9782
* `user.image`, `workspace.logo_url`, `organization.logo`, `organization.whitelabel_settings`,
9883
* `organization.metadata`, `custom_block.icon_url`, `chat.customizations` (deployed-chat logo
99-
* and avatar), and `oauth_client.icon`. With `needles`, reads only rows whose value contains one;
100-
* without, every row that holds a URL.
84+
* and avatar), and `oauth_client.icon`. Every non-empty value is read and decoded here rather than
85+
* matched in SQL, since a URL may percent-encode any character of the key it names. Only
86+
* `user.image`, mostly external OAuth avatars, is narrowed, to values that point at this app's
87+
* file-serve route or name a world-readable root.
10188
*/
102-
async function loadUrlReferencedKeys(needles: readonly string[] | null): Promise<Set<string>> {
103-
const matching = (column: SQLWrapper, fallback: SQL | undefined): SQL | undefined =>
104-
needles ? or(...needles.map((needle) => sql`${column}::text LIKE ${`%${needle}%`}`)) : fallback
89+
async function loadUrlReferencedKeys(): Promise<Set<string>> {
10590
const values: string[] = []
10691
const sources = [
10792
db
10893
.select({ value: user.image })
10994
.from(user)
11095
.where(
111-
matching(
112-
user.image,
113-
or(like(user.image, '%/api/files/serve/%'), like(user.image, '%profile-pictures/%'))
96+
or(
97+
like(user.image, '%/api/files/serve/%'),
98+
like(user.image, '%profile-pictures%'),
99+
like(user.image, '%og-images%'),
100+
like(user.image, '%workspace-logos%'),
101+
like(user.image, '%organization-logos%')
114102
)
115103
),
116104
db
117105
.select({ value: workspaceTable.logoUrl })
118106
.from(workspaceTable)
119-
.where(matching(workspaceTable.logoUrl, isNotNull(workspaceTable.logoUrl))),
107+
.where(isNotNull(workspaceTable.logoUrl)),
120108
db
121109
.select({ value: customBlock.iconUrl })
122110
.from(customBlock)
123-
.where(matching(customBlock.iconUrl, isNotNull(customBlock.iconUrl))),
124-
db
125-
.select({ value: chat.customizations })
126-
.from(chat)
127-
.where(matching(chat.customizations, sql`${chat.customizations}::text LIKE '%Url%'`)),
128-
db
129-
.select({ value: oauthClient.icon })
130-
.from(oauthClient)
131-
.where(matching(oauthClient.icon, isNotNull(oauthClient.icon))),
111+
.where(isNotNull(customBlock.iconUrl)),
112+
db.select({ value: chat.customizations }).from(chat).where(isNotNull(chat.customizations)),
113+
db.select({ value: oauthClient.icon }).from(oauthClient).where(isNotNull(oauthClient.icon)),
132114
]
133115
for (const source of sources) for (const row of await source) collectStrings(row.value, values)
134116
const organizations = await db
@@ -140,9 +122,9 @@ async function loadUrlReferencedKeys(needles: readonly string[] | null): Promise
140122
.from(organization)
141123
.where(
142124
or(
143-
matching(organization.logo, isNotNull(organization.logo)),
144-
matching(organization.whitelabelSettings, isNotNull(organization.whitelabelSettings)),
145-
matching(organization.metadata, isNotNull(organization.metadata))
125+
isNotNull(organization.logo),
126+
isNotNull(organization.whitelabelSettings),
127+
isNotNull(organization.metadata)
146128
)
147129
)
148130
for (const row of organizations) collectStrings([row.logo, row.whitelabel, row.metadata], values)
@@ -158,8 +140,8 @@ async function loadUrlReferencedKeys(needles: readonly string[] | null): Promise
158140
* `workspace_file_version.key`, `document.storage_key`, `execution_large_values.key`; or
159141
* - it sits under a world-readable root (`isPublicStorageKey`, the same prefixes the file-serve
160142
* route's `isPublicByKeyPrefix` answers without a session) and a stored asset URL names it (see
161-
* {@link loadUrlReferencedKeys}). Only rows containing the key's longest URL-safe run are read,
162-
* each run at most once per guard. No URL column is read for any other key:
143+
* {@link loadUrlReferencedKeys}), read at most once per guard and only once a public key comes
144+
* up. No URL column is read for any other key:
163145
* `verifyFileAccess` serves it only through a row, by its key path under a workspace that still
164146
* exists (which exact-key cleanup keeps; see `keysUnderLiveWorkspaces`), or by object metadata
165147
* to a signed-in holder of a hand-pasted URL to a private file this deletion means to erase.
@@ -171,25 +153,7 @@ async function loadUrlReferencedKeys(needles: readonly string[] | null): Promise
171153
* attachments that is the intended erasure. A failed lookup throws, so nothing is deleted.
172154
*/
173155
export function createStorageReferenceGuard(): (keys: readonly string[]) => Promise<Set<string>> {
174-
const loadedNeedles = new Set<string>()
175-
const named = new Set<string>()
176-
let loadedAll = false
177-
const loadNamedKeys = async (publicKeys: readonly string[]) => {
178-
if (loadedAll || publicKeys.length === 0) return
179-
const needles = publicKeys.map(urlNeedle)
180-
if (needles.includes(null)) {
181-
for (const key of await loadUrlReferencedKeys(null)) named.add(key)
182-
loadedAll = true
183-
return
184-
}
185-
const pending = [...new Set(needles.flatMap((needle) => needle ?? []))].filter(
186-
(needle) => !loadedNeedles.has(needle)
187-
)
188-
for (const batch of chunkArray(pending, NEEDLES_PER_QUERY)) {
189-
for (const key of await loadUrlReferencedKeys(batch)) named.add(key)
190-
for (const needle of batch) loadedNeedles.add(needle)
191-
}
192-
}
156+
let named: Set<string> | null = null
193157
return async (keys) => {
194158
const referenced = new Set<string>()
195159
if (keys.length === 0) return referenced
@@ -218,8 +182,10 @@ export function createStorageReferenceGuard(): (keys: readonly string[]) => Prom
218182
]
219183
for (const row of rows) if (row.key) referenced.add(row.key)
220184
const publicKeys = list.filter(isPublicStorageKey)
221-
await loadNamedKeys(publicKeys)
222-
for (const key of publicKeys) if (named.has(key)) referenced.add(key)
185+
if (publicKeys.length > 0) {
186+
named ??= await loadUrlReferencedKeys()
187+
for (const key of publicKeys) if (named.has(key)) referenced.add(key)
188+
}
223189
return referenced
224190
}
225191
}

‎apps/sim/lib/users/storage-sweep.integration.ts‎

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,14 @@ async function insertAdminWorkspace(id: string, ownerId: string) {
103103
})
104104
}
105105

106+
/** A legal URL for `key` that shares no literal run with it, as a hand-written link can be. */
107+
function percentEncodeEveryCharacter(key: string): string {
108+
return Array.from(
109+
key,
110+
(character) => `%${character.charCodeAt(0).toString(16).padStart(2, '0')}`
111+
).join('')
112+
}
113+
106114
function fileRow(values: {
107115
key: string
108116
userId: string
@@ -112,7 +120,7 @@ function fileRow(values: {
112120
}) {
113121
return {
114122
id: generateId(),
115-
originalName: 'fixture.txt',
123+
originalName: `fixture-${generateId()}.txt`,
116124
contentType: 'text/plain',
117125
sizeBytes: 1,
118126
chatId: values.chatId ?? null,
@@ -154,6 +162,7 @@ async function seed() {
154162
legacyWorkspacelessKb: `kb/${generateId()}-legacy.txt`,
155163
profilePictureRowLogo: `profile-pictures/${generateId()}/org-logo.png`,
156164
logoNamedByUrl: `workspace-logos/${doomedId}/${generateId()}-org-logo.png`,
165+
logoNamedByEncodedUrl: `workspace-logos/${doomedId}/${generateId()}-workspace-logo.png`,
157166
}
158167
await db.insert(workspaceFiles).values([
159168
fileRow({
@@ -217,6 +226,12 @@ async function seed() {
217226
.update(user)
218227
.set({ image: `/api/files/serve/${encodeURIComponent(keys.logoNamedByUrl)}` })
219228
.where(eq(user.id, ownerId))
229+
await db
230+
.update(workspace)
231+
.set({
232+
logoUrl: `/api/files/serve/s3/${percentEncodeEveryCharacter(keys.logoNamedByEncodedUrl)}`,
233+
})
234+
.where(eq(workspace.id, survivingId))
220235
for (const key of Object.values(keys)) await put(key)
221236
return { ownerId, departingId, doomedId, survivingId, keys }
222237
}
@@ -304,6 +319,7 @@ describe('account deletion storage sweep', () => {
304319
f.keys.legacyWorkspacelessKb,
305320
f.keys.profilePictureRowLogo,
306321
f.keys.logoNamedByUrl,
322+
f.keys.logoNamedByEncodedUrl,
307323
])
308324
expect(await readFile(join(storageRoot, key), 'utf8')).toBe('x')
309325
expect(

‎apps/sim/lib/users/storage-sweep.test.ts‎

Lines changed: 31 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -186,19 +186,31 @@ describe('deleted workspace storage sweep', () => {
186186
expect(await exists('workspace/ws1/nested/b.txt.upload-metadata.json')).toBe(false)
187187
})
188188

189-
it('spends an attempt instead of continuing when a full local page removes nothing', async () => {
190-
for (let index = 0; index < 520; index++)
191-
await put(`workspace/ws1/legacy/legacy file ${index}.txt`)
189+
it('pages past a full page of kept objects to the ones behind it', async () => {
190+
const keptKeys = Array.from(
191+
{ length: 520 },
192+
(_, index) => `workspace/ws1/a-kept/${String(index).padStart(3, '0')}.bin`
193+
)
194+
const unkeptKeys = Array.from(
195+
{ length: 10 },
196+
(_, index) => `workspace/ws1/z-unkept/${index}.bin`
197+
)
198+
for (const key of [...keptKeys, ...unkeptKeys]) await put(key)
199+
queueTableRows(
200+
schemaMock.document,
201+
keptKeys.map((key) => ({ key }))
202+
)
203+
queueTableRows(
204+
schemaMock.document,
205+
keptKeys.map((key) => ({ key }))
206+
)
192207

193-
const result = await sweep(
194-
{ version: 1, workspaceId: 'ws1' },
195-
outboxContext({ deadlineAt: Date.now() + 120_000 })
208+
await expect(sweep({ version: 1, workspaceId: 'ws1' }, outboxContext())).resolves.toBe(
209+
undefined
196210
)
197211

198-
expect(result).toEqual({
199-
outcome: 'deferred',
200-
reason: 'Workspace storage sweep removed nothing from a full listing page',
201-
})
212+
for (const key of unkeptKeys) expect(await exists(key)).toBe(false)
213+
for (const key of keptKeys) expect(await exists(key)).toBe(true)
202214
})
203215

204216
it.each([
@@ -398,20 +410,22 @@ describe('queueing account storage cleanup', () => {
398410
expect(keys).toContain('copilot/abc/v1004.txt')
399411
})
400412

401-
it('only inserts inside the deletion transaction', async () => {
413+
it('queues a file committed between the early read and the deletion locks, once', async () => {
414+
const early = { id: 'file-1', key: 'copilot/abc/early.txt', context: 'copilot' }
415+
queueTableRows(schemaMock.workspaceFiles, [early])
416+
const objects = await collectAccountStorageObjects({ userId: 'user-1', workspaceIds: [] })
402417
queueTableRows(schemaMock.workspaceFiles, [
403-
{ id: 'file-1', key: 'copilot/abc/file.txt', context: 'copilot' },
418+
early,
419+
{ id: 'file-2', key: 'copilot/abc/late.txt', context: 'copilot' },
404420
])
405-
const objects = await collectAccountStorageObjects({ userId: 'user-1', workspaceIds: ['ws1'] })
406-
const insertOnly = { insert: dbChainMock.db.insert } as unknown as DbTransaction
407421

408-
await enqueueAccountStorageCleanupInTx(insertOnly, {
422+
await enqueueAccountStorageCleanupInTx(dbChainMock.db as unknown as DbTransaction, {
409423
userId: 'user-1',
410-
workspaceIds: ['ws1'],
424+
workspaceIds: [],
411425
objects,
412426
})
413427

414-
expect(queuedObjectKeys()).toEqual(['copilot/abc/file.txt'])
428+
expect(queuedObjectKeys()).toEqual(['copilot/abc/early.txt', 'copilot/abc/late.txt'])
415429
})
416430

417431
it('never queues a profile picture row', async () => {

0 commit comments

Comments
 (0)