gongxiangp2/lib/__tests__/upload-sessions-security.test.ts

/** @jest-environment node */

import path from 'path'
import os from 'os'
import {
  mkdtemp,
  mkdir,
  readFile,
  rm,
  symlink,
  writeFile,
} from 'fs/promises'
import {
  getUploadSession,
  getUploadedFile,
  saveUploadSession,
  uploadSessions,
  type UploadSession,
} from '@/lib/upload-sessions'
import { cleanupTempFiles } from '@/lib/task-material-processor'

function createSession(fileName: string, fileSize: number): UploadSession {
  return {
    fileName,
    fileSize,
    totalChunks: 1,
    chunks: new Map(),
    createdAt: Date.now(),
    userId: 'user-1',
  }
}

describe('upload-sessions 最终落盘安全', () => {
  let uploadRoot = ''
  const previousUploadDir = process.env.UPLOAD_DIR

  beforeEach(async () => {
    uploadSessions.clear()
    uploadRoot = await mkdtemp(path.join(os.tmpdir(), 'upload-session-security-'))
    process.env.UPLOAD_DIR = uploadRoot
  })

  afterEach(async () => {
    uploadSessions.clear()
    if (uploadRoot) {
      await rm(uploadRoot, { recursive: true, force: true })
    }
    if (previousUploadDir === undefined) {
      delete process.env.UPLOAD_DIR
    } else {
      process.env.UPLOAD_DIR = previousUploadDir
    }
  })

  async function writeSingleChunk(uploadId: string, content: Buffer) {
    const sessionDir = path.join(uploadRoot, 'temp', uploadId)
    await mkdir(sessionDir, { recursive: true })
    await writeFile(path.join(sessionDir, 'chunk-0'), content)
  }

  test('合法会话只能在 assembled 目录生成普通文件', async () => {
    const uploadId = 'upload-safe-file'
    const content = Buffer.from('safe-content')
    await writeSingleChunk(uploadId, content)
    await saveUploadSession(uploadId, createSession('中文源文件.psd', content.length))

    const outputPath = await getUploadedFile(uploadId)

    expect(outputPath).toBe(
      path.join(uploadRoot, 'assembled', `${uploadId}-中文源文件.psd`),
    )
    await expect(readFile(outputPath!)).resolves.toEqual(content)
  })

  test('物理分片超过会话声明大小时在读入内存前拒绝组装', async () => {
    const uploadId = 'upload-oversized-physical-chunk'
    const declaredContent = Buffer.from('safe')
    const oversizedChunk = Buffer.concat([declaredContent, Buffer.alloc(1024 * 1024, 1)])
    await writeSingleChunk(uploadId, oversizedChunk)
    await saveUploadSession(uploadId, createSession('source.jpg', declaredContent.length))

    await expect(getUploadedFile(uploadId)).resolves.toBeNull()
    await expect(
      readFile(path.join(uploadRoot, 'assembled', `${uploadId}-source.jpg`)),
    ).rejects.toMatchObject({ code: 'ENOENT' })
  })

  test.each([
    '../../../sentinel.txt',
    '..\\..\\sentinel.txt',
    '/app/server.js',
    'C:\\Windows\\system.ini',
    '\\\\server\\share\\outside.psd',
    'bad\u0000name.psd',
  ])('历史恶意会话在最终落盘时仍被拒绝：%s', async (fileName) => {
    const uploadId = `upload-malicious-${Math.random().toString(16).slice(2)}`
    const sentinelPath = path.join(uploadRoot, 'sentinel.txt')
    await writeFile(sentinelPath, 'sentinel')
    await writeSingleChunk(uploadId, Buffer.from('attack'))

    // 模拟修复前已经进入内存或 Redis 的旧会话，绕过新的 saveUploadSession 入口校验。
    uploadSessions.set(uploadId, createSession(fileName, 6))

    await expect(getUploadedFile(uploadId)).resolves.toBeNull()
    await expect(readFile(sentinelPath, 'utf8')).resolves.toBe('sentinel')
  })

  test('saveUploadSession 自身拒绝不安全文件名', async () => {
    await expect(
      saveUploadSession(
        'upload-invalid-session',
        createSession('../outside.psd', 1),
      ),
    ).rejects.toThrow('上传会话包含不安全的标识或文件名')
  })

  test('业务完成清理会同时移除进程内会话，不能只删除 Redis', async () => {
    const uploadId = 'upload-unified-session-cleanup'
    await writeSingleChunk(uploadId, Buffer.from('safe'))
    await saveUploadSession(uploadId, createSession('source.jpg', 4))
    await expect(getUploadSession(uploadId)).resolves.not.toBeNull()

    // 受控清理只能消费已完成的上传：先原子合并并写入 completed
    // tombstone，再删除 assembled/temp 和内存 + Redis 会话。
    await expect(getUploadedFile(uploadId)).resolves.toEqual(
      path.join(uploadRoot, 'assembled', `${uploadId}-source.jpg`),
    )

    await cleanupTempFiles([], [uploadId])

    await expect(getUploadSession(uploadId)).resolves.toBeNull()
  })

  test('最终目标是符号链接时不会跟随或覆盖链接目标', async () => {
    const uploadId = 'upload-symlink-target'
    const content = Buffer.from('attack')
    await writeSingleChunk(uploadId, content)
    await saveUploadSession(uploadId, createSession('source.psd', content.length))

    const assembledDir = path.join(uploadRoot, 'assembled')
    const sentinelDir = path.join(uploadRoot, 'sentinel-directory')
    const sentinelPath = path.join(sentinelDir, 'sentinel.txt')
    await mkdir(assembledDir)
    await mkdir(sentinelDir)
    await writeFile(sentinelPath, 'sentinel')
    await symlink(
      sentinelDir,
      path.join(assembledDir, `${uploadId}-source.psd`),
      'junction',
    )

    await expect(getUploadedFile(uploadId)).resolves.toBeNull()
    await expect(readFile(sentinelPath, 'utf8')).resolves.toBe('sentinel')
  })
})