|
| 1 | +/** |
| 2 | + * @vitest-environment node |
| 3 | + */ |
| 4 | +import { Readable } from 'stream' |
| 5 | +import type { SFTPWrapper } from 'ssh2' |
| 6 | +import { describe, expect, it, vi } from 'vitest' |
| 7 | +import { isPayloadSizeLimitError } from '@/lib/core/utils/stream-limits' |
| 8 | +import { MAX_SFTP_READ_BYTES, readSftpFileCapped } from '@/app/api/tools/sftp/utils' |
| 9 | + |
| 10 | +/** |
| 11 | + * Builds a fake SFTP wrapper whose read stream emits `chunkCount` chunks of |
| 12 | + * `chunkSize` bytes — the shape of a malicious server that understates the |
| 13 | + * file size in its stat reply and then streams unbounded data. |
| 14 | + */ |
| 15 | +function fakeSftp(chunkSize: number, chunkCount: number) { |
| 16 | + let emitted = 0 |
| 17 | + const stream = new Readable({ |
| 18 | + read() { |
| 19 | + if (emitted >= chunkCount) { |
| 20 | + this.push(null) |
| 21 | + return |
| 22 | + } |
| 23 | + emitted++ |
| 24 | + this.push(Buffer.alloc(chunkSize, 0x41)) |
| 25 | + }, |
| 26 | + }) |
| 27 | + const createReadStream = vi.fn(() => stream) |
| 28 | + return { sftp: { createReadStream } as unknown as SFTPWrapper, stream, createReadStream } |
| 29 | +} |
| 30 | + |
| 31 | +describe('readSftpFileCapped', () => { |
| 32 | + it('resolves with the full contents when under the cap', async () => { |
| 33 | + const { sftp, createReadStream } = fakeSftp(4, 3) |
| 34 | + |
| 35 | + const buffer = await readSftpFileCapped(sftp, '/file', 1024, 'file') |
| 36 | + |
| 37 | + expect(buffer.toString()).toBe('A'.repeat(12)) |
| 38 | + expect(createReadStream).toHaveBeenCalledWith('/file') |
| 39 | + }) |
| 40 | + |
| 41 | + it('rejects and destroys the stream once received bytes exceed the cap', async () => { |
| 42 | + const { sftp, stream } = fakeSftp(8, 1_000_000) |
| 43 | + |
| 44 | + await expect(readSftpFileCapped(sftp, '/bomb', 16, 'file')).rejects.toSatisfy( |
| 45 | + isPayloadSizeLimitError |
| 46 | + ) |
| 47 | + expect(stream.destroyed).toBe(true) |
| 48 | + }) |
| 49 | + |
| 50 | + it('enforces the cap on actual bytes even when the file was reported as tiny', async () => { |
| 51 | + const { sftp, stream } = fakeSftp(1024, 1_000_000) |
| 52 | + |
| 53 | + await expect(readSftpFileCapped(sftp, '/bomb', 4096, 'file')).rejects.toThrow( |
| 54 | + /exceeds maximum size of 4096 bytes/ |
| 55 | + ) |
| 56 | + expect(stream.destroyed).toBe(true) |
| 57 | + }) |
| 58 | + |
| 59 | + it('survives the late stream error ssh2 emits when the channel closes after an abort', async () => { |
| 60 | + const { sftp, stream } = fakeSftp(8, 1_000_000) |
| 61 | + |
| 62 | + await expect(readSftpFileCapped(sftp, '/bomb', 16, 'file')).rejects.toSatisfy( |
| 63 | + isPayloadSizeLimitError |
| 64 | + ) |
| 65 | + |
| 66 | + expect(() => stream.emit('error', new Error('No response from server'))).not.toThrow() |
| 67 | + }) |
| 68 | + |
| 69 | + it('caps remote reads at 50MB', () => { |
| 70 | + expect(MAX_SFTP_READ_BYTES).toBe(50 * 1024 * 1024) |
| 71 | + }) |
| 72 | +}) |
0 commit comments