diff --git a/server/src/nest/storage/drivers/s3.driver.ts b/server/src/nest/storage/drivers/s3.driver.ts index 5be8fa525..fa97ab07c 100644 --- a/server/src/nest/storage/drivers/s3.driver.ts +++ b/server/src/nest/storage/drivers/s3.driver.ts @@ -1,10 +1,12 @@ /// import awsLite from '@aws-lite/client'; -import type { Readable } from 'node:stream'; +import { PassThrough, type Readable } from 'node:stream'; +import fs from 'node:fs'; import { assertValidKey, assertValidPrefix, isValidKey } from '../storage-keys'; import { StorageBackendError, StorageNotFoundError, + isLocalTempFile, type ByteRange, type LocalTempFile, type ObjectStat, @@ -129,8 +131,92 @@ export class S3Driver implements StorageDriver { return null; } - async put(_key: string, _source: Readable | LocalTempFile, _opts?: PutOptions): Promise { - throw new StorageBackendError('not implemented'); + async put(key: string, source: Readable | LocalTempFile, opts?: PutOptions): Promise { + assertValidKey(key); + const s3 = await this.client(); + const base: Record = { Bucket: this.bucket, Key: this.keyPrefix + key }; + if (opts?.contentType) base.ContentType = opts.contentType; + + if (isLocalTempFile(source)) { + // Ownership transfer: consume the tmp file on success AND failure. + try { + const { size } = await fs.promises.stat(source.tmpPath); + if (size < MULTIPART_THRESHOLD) { + await this.deadlined(s3.PutObject({ ...base, File: source.tmpPath }), `put '${key}'`); + } else { + await this.uploadMonitored(s3, base, fs.createReadStream(source.tmpPath), key); + } + } catch (err) { + throw this.wrap(err, `put failed for '${key}'`); + } finally { + await fs.promises.rm(source.tmpPath, { force: true }); + } + return; + } + + // Bounded peek (spec, Client): a stream ending within the threshold — + // including zero bytes — becomes one PutObject with the buffered Body (the + // same memory Upload would buffer for part 1, and one round trip instead + // of 3+; aws-lite's Upload also breaks on empty streams). A larger stream + // flows into Upload with the prefix re-attached. An unbounded stream never + // reaches PutObject, which would buffer it wholesale. + const peeked = await peekStream(source, MULTIPART_THRESHOLD); // source errors surface untouched + // `readableEnded` guards a race: a source whose 'end' fires between the + // peek resolving (size > threshold) and pipe() attaching would never end + // the PassThrough — but ended means the chunks already hold everything. + if (peeked.ended || source.readableEnded) { + try { + await this.deadlined(s3.PutObject({ ...base, Body: Buffer.concat(peeked.chunks) }), `put '${key}'`); + } catch (err) { + throw this.wrap(err, `put failed for '${key}'`); + } + return; + } + const combined = new PassThrough(); + for (const chunk of peeked.chunks) combined.write(chunk); + source.pipe(combined); + source.on('error', (err) => combined.destroy(err)); + try { + await this.uploadMonitored(s3, base, combined, key, source); + } catch (err) { + throw this.wrap(err, `put failed for '${key}'`); + } + } + + /** + * Multipart Upload under an INACTIVITY deadline: a fixed whole-operation + * deadline would break large backups, so the timer resets on data movement. + * Expiry destroys the monitored stream, driving Upload's own error path — + * in-flight parts drain, then AbortMultipartUpload fires (verified upstream). + */ + private async uploadMonitored( + s3: S3Api, + base: Record, + body: Readable, + key: string, + sourceToDestroy?: Readable, + ): Promise { + let timer: NodeJS.Timeout | undefined; + const expire = () => { + const err = new StorageBackendError( + `put '${key}' stalled for ${this.timeoutMs}ms on '${this.id}' — multipart upload aborted`, + ); + body.destroy(err); + sourceToDestroy?.destroy(err); + }; + const reset = () => { + clearTimeout(timer); + timer = setTimeout(expire, this.timeoutMs); + timer.unref?.(); + }; + body.on('data', reset); // body is consumed by Upload; extra listener observes only + reset(); + try { + await s3.Upload({ ...base, Body: body, ChunkSize: MULTIPART_THRESHOLD }); + } finally { + clearTimeout(timer); + body.off('data', reset); + } } async getStream(key: string, range?: ByteRange): Promise<{ stream: Readable; stat: ObjectStat }> { @@ -265,3 +351,50 @@ function totalFromContentRange(contentRange: string | undefined, key: string, id } return Number(total); } + +/** + * Reads up to `limit` bytes in paused mode. Resolves `ended: true` when the + * stream finished within the limit (chunks = the whole payload), else + * `ended: false` with the peeked prefix still unconsumed downstream. Rejects + * with the source's own error untouched (LocalDriver parity). + */ +function peekStream(source: Readable, limit: number): Promise<{ chunks: Buffer[]; size: number; ended: boolean }> { + return new Promise((resolve, reject) => { + const chunks: Buffer[] = []; + let size = 0; + const cleanup = () => { + source.off('readable', onReadable); + source.off('end', onEnd); + source.off('error', onError); + }; + const onReadable = () => { + for (;;) { + const raw = source.read() as Buffer | string | null; + if (raw === null) return; + // A source not explicitly in binary mode (e.g. an object-mode + // Readable.from(string), as used by callers/tests) can yield string + // chunks; PutObject's Body and Buffer.concat both need real Buffers. + const chunk = Buffer.isBuffer(raw) ? raw : Buffer.from(raw); + chunks.push(chunk); + size += chunk.length; + if (size > limit) { + cleanup(); + resolve({ chunks, size, ended: false }); + return; + } + } + }; + const onEnd = () => { + cleanup(); + resolve({ chunks, size, ended: true }); + }; + const onError = (err: Error) => { + cleanup(); + reject(err); + }; + source.on('readable', onReadable); + source.on('end', onEnd); + source.on('error', onError); + onReadable(); + }); +} diff --git a/server/tests/unit/nest/storage/s3.driver.test.ts b/server/tests/unit/nest/storage/s3.driver.test.ts index 54aa57f4a..d3202a9fd 100644 --- a/server/tests/unit/nest/storage/s3.driver.test.ts +++ b/server/tests/unit/nest/storage/s3.driver.test.ts @@ -1,5 +1,8 @@ import { describe, it, expect, vi } from 'vitest'; import { Readable } from 'node:stream'; +import fs from 'node:fs'; +import os from 'node:os'; +import path from 'node:path'; import { S3Driver, MULTIPART_THRESHOLD, @@ -233,3 +236,119 @@ describe('defaultClientFactory', () => { } }); }); + +describe('S3Driver put — bounded peek routing', () => { + it('routes a small stream (ends within threshold) to a single PutObject with buffered Body', async () => { + const api = makeMockApi({ PutObject: vi.fn().mockResolvedValue({}) }); + await makeDriver(api).put('a/x.bin', Readable.from('hello'), { contentType: 'text/plain' }); + const input = (api.PutObject as ReturnType).mock.calls[0][0]; + expect(Buffer.isBuffer(input.Body) && input.Body.toString()).toBe('hello'); + expect(input).toMatchObject({ Key: 'a/x.bin', ContentType: 'text/plain' }); + expect(api.Upload).not.toHaveBeenCalled(); + }); + it('routes a zero-byte stream to PutObject (Upload breaks on empty streams upstream)', async () => { + const api = makeMockApi({ PutObject: vi.fn().mockResolvedValue({}) }); + await makeDriver(api).put('a/empty.bin', Readable.from([])); + const input = (api.PutObject as ReturnType).mock.calls[0][0]; + expect(Buffer.isBuffer(input.Body) && input.Body.length).toBe(0); + }); + it('routes a stream exceeding the threshold to Upload with the peeked prefix re-attached', async () => { + let uploaded: Buffer | null = null; + const api = makeMockApi({ + Upload: vi.fn().mockImplementation(async ({ Body }: { Body: Readable }) => { + uploaded = await drain(Body); + }), + }); + // Three chunks so the source cannot end during the peek (the peek stops + // after chunk 2 crosses the threshold; chunk 3 is still pending) — the + // Upload path is deterministic, not timing-dependent. + const chunk = Buffer.alloc(6 * 1024 * 1024, 7); + const big = Buffer.concat([chunk, chunk, chunk]); + await makeDriver(api).put('a/big.bin', Readable.from([chunk, chunk, chunk])); + expect(api.PutObject).not.toHaveBeenCalled(); + expect(uploaded!.length).toBe(big.length); + expect(uploaded!.equals(big)).toBe(true); + expect((api.Upload as ReturnType).mock.calls[0][0]).toMatchObject({ + Key: 'a/big.bin', + ChunkSize: MULTIPART_THRESHOLD, + }); + }); + it('routes a single-chunk oversized stream that ends at the peek boundary to PutObject (readableEnded guard)', async () => { + // One huge chunk: 'end' can fire between the peek resolving and pipe() + // attaching — the driver must detect readableEnded and PutObject the + // already-buffered payload rather than hang a never-ending PassThrough. + const api = makeMockApi({ + PutObject: vi.fn().mockResolvedValue({}), + Upload: vi.fn().mockResolvedValue({}), + }); + const big = Buffer.alloc(MULTIPART_THRESHOLD + 3, 7); + const oneChunk = new Readable({ read() {} }); + oneChunk.push(big); + oneChunk.push(null); + await makeDriver(api).put('a/one-chunk.bin', oneChunk); + // Either path must complete with the full payload; no hang, no data loss. + const putCalls = (api.PutObject as ReturnType).mock.calls; + const uploadCalls = (api.Upload as ReturnType).mock.calls; + expect(putCalls.length + uploadCalls.length).toBe(1); + if (putCalls.length === 1) { + expect((putCalls[0][0].Body as Buffer).length).toBe(big.length); + } + }); + it('surfaces a source-stream error before any client call', async () => { + const api = makeMockApi(); + const boom = new Readable({ + read() { + this.destroy(new Error('source exploded')); + }, + }); + await expect(makeDriver(api).put('a/broken.bin', boom)).rejects.toThrow('source exploded'); + expect(api.PutObject).not.toHaveBeenCalled(); + expect(api.Upload).not.toHaveBeenCalled(); + }); + it('expires a stalled Upload via the inactivity deadline (PassThrough destroyed)', async () => { + const api = makeMockApi({ + Upload: vi.fn().mockImplementation( + ({ Body }: { Body: Readable }) => + new Promise((_, reject) => { + Body.on('error', reject); // aws-lite Upload fails when its source errors + }), + ), + }); + const stalled = new Readable({ read() {} }); + stalled.push(Buffer.alloc(MULTIPART_THRESHOLD + 1)); // exceed the peek, then go silent + await expect(makeDriver(api, { timeoutMs: 60 }).put('a/stall.bin', stalled)).rejects.toBeInstanceOf( + StorageBackendError, + ); + }); +}); + +describe('S3Driver put — LocalTempFile ownership', () => { + async function makeTmp(bytes: number): Promise { + const dir = await fs.promises.mkdtemp(path.join(os.tmpdir(), 'trek-s3-')); + const p = path.join(dir, 'src.bin'); + await fs.promises.writeFile(p, Buffer.alloc(bytes, 1)); + return p; + } + it('uploads a small temp file via PutObject File and consumes it', async () => { + const api = makeMockApi({ PutObject: vi.fn().mockResolvedValue({}) }); + const tmp = await makeTmp(10); + await makeDriver(api).put('a/t.bin', { tmpPath: tmp }, { contentType: 'image/png' }); + expect(api.PutObject).toHaveBeenCalledWith( + expect.objectContaining({ File: tmp, Key: 'a/t.bin', ContentType: 'image/png' }), + ); + expect(fs.existsSync(tmp)).toBe(false); + }); + it('routes a threshold-sized temp file to Upload', async () => { + const api = makeMockApi({ Upload: vi.fn().mockImplementation(async ({ Body }: { Body: Readable }) => drain(Body)) }); + const tmp = await makeTmp(MULTIPART_THRESHOLD); + await makeDriver(api).put('a/t.bin', { tmpPath: tmp }); + expect(api.Upload).toHaveBeenCalled(); + expect(fs.existsSync(tmp)).toBe(false); + }); + it('consumes the temp file on the failure path too (MirrorDriver precedent)', async () => { + const api = makeMockApi({ PutObject: vi.fn().mockRejectedValue(awsError(500)) }); + const tmp = await makeTmp(10); + await expect(makeDriver(api).put('a/t.bin', { tmpPath: tmp })).rejects.toBeInstanceOf(StorageBackendError); + expect(fs.existsSync(tmp)).toBe(false); + }); +});