gongxiangp2/lib/__tests__/material-upload-chunk-limit.test.ts

/** @jest-environment node */

import { mkdtemp, readFile, readdir, rm, stat } from 'fs/promises'
import os from 'os'
import path from 'path'
import { NextRequest } from 'next/server'

import { POST } from '@/app/api/materials/upload/chunk/route'
import { materialUploadChunkRequestLimiter } from '@/lib/material-upload-request-guard'
import { getAuthenticatedUploadUserId } from '@/lib/upload-auth'
import { getUploadSession, updateChunkWithStats } from '@/lib/upload-sessions'
import {
  beginUploadCompletionSeal,
  cancelUploadSessionSeal,
  completeUploadCompletionSeal,
} from '@/lib/material-upload-session-seal'
import { acquireMaterialUploadStorageLock } from '@/lib/material-upload-cross-process-lock'

jest.mock('@/lib/upload-auth', () => ({
  getAuthenticatedUploadUserId: jest.fn(),
}))

jest.mock('@/lib/upload-sessions', () => ({
  getUploadSession: jest.fn(),
  updateChunkWithStats: jest.fn(),
}))

const MEBIBYTE = 1024 * 1024

function streamFromChunks(chunks: Uint8Array[]): ReadableStream<Uint8Array> {
  let index = 0
  return new ReadableStream<Uint8Array>({
    pull(controller) {
      const chunk = chunks[index]
      index += 1
      if (!chunk) {
        controller.close()
        return
      }
      controller.enqueue(chunk)
    },
  })
}

function createRawChunkRequest(
  chunks: Uint8Array[],
  declaredLength: number,
  includeSizeHint = true,
) {
  const headers = new Headers({
    'content-length': String(declaredLength),
    'content-type': 'application/octet-stream',
    'x-file-name': 'video.mp4',
    'x-upload-id': 'upload-test-1',
    'x-chunk-index': '0',
  })
  if (includeSizeHint) {
    headers.set('x-chunk-size', String(declaredLength))
  }

  return new NextRequest('http://localhost/api/materials/upload/chunk', {
    method: 'POST',
    headers,
    body: streamFromChunks(chunks),
    duplex: 'half',
  } as any)
}

describe('material upload chunk hard limit', () => {
  let uploadRoot: string

  beforeEach(async () => {
    jest.clearAllMocks()
    uploadRoot = await mkdtemp(path.join(os.tmpdir(), 'material-chunk-limit-'))
    process.env.UPLOAD_DIR = uploadRoot
    ;(getAuthenticatedUploadUserId as jest.Mock).mockResolvedValue('user-1')
    ;(getUploadSession as jest.Mock).mockResolvedValue({
      fileName: 'video.mp4',
      fileSize: 40 * MEBIBYTE,
      totalChunks: 2,
      chunkSize: 20 * MEBIBYTE,
      chunks: new Map(),
      createdAt: Date.now(),
      userId: 'user-1',
    })
    ;(updateChunkWithStats as jest.Mock).mockResolvedValue(undefined)
  })

  afterEach(async () => {
    delete process.env.UPLOAD_DIR
    await rm(uploadRoot, { recursive: true, force: true })
  })

  test('20MB raw chunk streams to disk and records the exact size', async () => {
    const request = createRawChunkRequest(
      [new Uint8Array(10 * MEBIBYTE), new Uint8Array(10 * MEBIBYTE)],
      20 * MEBIBYTE,
    )

    const response = await POST(request)
    const body = await response.json()
    const chunkPath = path.join(uploadRoot, 'temp', 'upload-test-1', 'chunk-0')

    expect(response.status).toBe(200)
    expect(body.receivedBytes).toBe(20 * MEBIBYTE)
    expect((await stat(chunkPath)).size).toBe(20 * MEBIBYTE)
    expect(updateChunkWithStats).toHaveBeenCalledWith(
      'upload-test-1',
      0,
      chunkPath,
      20 * MEBIBYTE,
    )
  })

  test('流式实际内容超过20MB时立即失败并删除已经写入的半文件', async () => {
    const request = createRawChunkRequest(
      [new Uint8Array(10 * MEBIBYTE), new Uint8Array(10 * MEBIBYTE), new Uint8Array(1)],
      20 * MEBIBYTE,
    )

    const response = await POST(request)
    const body = await response.json()
    const chunkPath = path.join(uploadRoot, 'temp', 'upload-test-1', 'chunk-0')

    expect(response.status).toBe(413)
    expect(body.code).toBe('CHUNK_TOO_LARGE')
    await expect(stat(chunkPath)).rejects.toMatchObject({ code: 'ENOENT' })
    expect(updateChunkWithStats).not.toHaveBeenCalled()
  })

  test('分片大小必须与会话声明的当前位置精确一致', async () => {
    ;(getUploadSession as jest.Mock).mockResolvedValueOnce({
      fileName: 'small-video.mp4',
      fileSize: 10 * MEBIBYTE,
      totalChunks: 1,
      chunkSize: 20 * MEBIBYTE,
      chunks: new Map(),
      createdAt: Date.now(),
      userId: 'user-1',
    })
    const request = createRawChunkRequest(
      [new Uint8Array(20 * MEBIBYTE)],
      20 * MEBIBYTE,
    )

    const response = await POST(request)
    const body = await response.json()

    expect(response.status).toBe(400)
    expect(body.code).toBe('CHUNK_SIZE_MISMATCH')
    await expect(
      stat(path.join(uploadRoot, 'temp', 'upload-test-1', 'chunk-0')),
    ).rejects.toMatchObject({ code: 'ENOENT' })
    expect(updateChunkWithStats).not.toHaveBeenCalled()
  })

  test('multipart 请求在 formData 解析前先等待八路并发槽', async () => {
    const releases: Array<() => void> = []
    for (let index = 0; index < 8; index += 1) {
      releases.push(await materialUploadChunkRequestLimiter.acquire())
    }

    const formData = jest.fn(async () => new FormData())
    const request = {
      headers: new Headers({
        'content-length': '128',
        'content-type': 'multipart/form-data; boundary=test',
      }),
      formData,
    } as unknown as NextRequest

    const responsePromise = POST(request)
    await new Promise((resolve) => setTimeout(resolve, 10))
    expect(formData).not.toHaveBeenCalled()

    releases.shift()?.()
    const response = await responsePromise
    expect(formData).toHaveBeenCalledTimes(1)
    expect(response.status).toBe(400)

    releases.forEach((release) => release())
    expect(materialUploadChunkRequestLimiter.snapshot()).toEqual({ active: 0, waiting: 0 })
  })

  test('同一分片并发重试只会原子替换完整文件，不会交错写坏', async () => {
    ;(getUploadSession as jest.Mock).mockResolvedValue({
      fileName: 'tiny-video.mp4',
      fileSize: 4,
      totalChunks: 1,
      chunkSize: 4,
      chunks: new Map(),
      createdAt: Date.now(),
      userId: 'user-1',
    })

    let firstController: ReadableStreamDefaultController<Uint8Array> | null = null
    const firstBody = new ReadableStream<Uint8Array>({
      start(controller) {
        firstController = controller
        controller.enqueue(new Uint8Array([1, 1]))
      },
    })
    const firstRequest = new NextRequest('http://localhost/api/materials/upload/chunk', {
      method: 'POST',
      headers: {
        'content-length': '4',
        'content-type': 'application/octet-stream',
        'x-file-name': 'tiny-video.mp4',
        'x-upload-id': 'upload-test-1',
        'x-chunk-index': '0',
        'x-chunk-size': '4',
      },
      body: firstBody,
      duplex: 'half',
    } as any)

    const firstResponsePromise = POST(firstRequest)
    const tempDir = path.join(uploadRoot, 'temp', 'upload-test-1')
    await waitFor(async () => (
      (await readdir(tempDir).catch(() => [])).some((name) => name.includes('.partial-'))
    ))

    const secondResponse = await POST(createRawChunkRequest(
      [new Uint8Array([2, 2, 2, 2])],
      4,
    ))
    expect(secondResponse.status).toBe(200)

    firstController!.enqueue(new Uint8Array([1, 1]))
    firstController!.close()
    const firstResponse = await firstResponsePromise
    expect(firstResponse.status).toBe(200)

    const finalChunk = await readFile(path.join(tempDir, 'chunk-0'))
    expect(Array.from(finalChunk)).toEqual([1, 1, 1, 1])
    expect((await readdir(tempDir)).filter((name) => name.includes('.partial-'))).toEqual([])
  })

  test('complete 封存后迟到分片不得发布 chunk 或复活陈旧会话', async () => {
    ;(getUploadSession as jest.Mock).mockResolvedValue({
      fileName: 'tiny-video.mp4',
      fileSize: 4,
      totalChunks: 1,
      chunkSize: 4,
      chunks: new Map(),
      createdAt: Date.now(),
      userId: 'user-1',
    })

    let lateController: ReadableStreamDefaultController<Uint8Array> | null = null
    const lateBody = new ReadableStream<Uint8Array>({
      start(controller) {
        lateController = controller
        controller.enqueue(new Uint8Array([1, 1]))
      },
    })
    const lateRequest = new NextRequest('http://localhost/api/materials/upload/chunk', {
      method: 'POST',
      headers: {
        'content-length': '4',
        'content-type': 'application/octet-stream',
        'x-file-name': 'tiny-video.mp4',
        'x-upload-id': 'upload-test-1',
        'x-chunk-index': '0',
        'x-chunk-size': '4',
      },
      body: lateBody,
      duplex: 'half',
    } as any)

    const lateResponsePromise = POST(lateRequest)
    const tempDir = path.join(uploadRoot, 'temp', 'upload-test-1')
    await waitFor(async () => (
      (await readdir(tempDir).catch(() => [])).some((name) => name.includes('.partial-'))
    ))

    // 模拟 complete 已在跨容器独占锁内完成。迟到请求仍持有之前读取的
    // 陈旧 session，但写完网络流后必须看到共享盘上的 completed seal。
    const seal = await beginUploadCompletionSeal('upload-test-1')
    expect(seal.state).toBe('sealing')
    if (seal.state !== 'sealing') {
      throw new Error('测试无法创建上传完成标记')
    }
    await completeUploadCompletionSeal(seal.handle)

    lateController!.enqueue(new Uint8Array([1, 1]))
    lateController!.close()
    const lateResponse = await lateResponsePromise
    const body = await lateResponse.json()

    expect(lateResponse.status).toBe(409)
    expect(body.code).toBe('UPLOAD_SESSION_SEALED')
    await expect(readFile(path.join(tempDir, 'chunk-0'))).rejects.toMatchObject({
      code: 'ENOENT',
    })
    expect((await readdir(tempDir)).filter((name) => name.includes('.partial-'))).toEqual([])
    expect(updateChunkWithStats).not.toHaveBeenCalled()
  })

  test('取消 tombstone 写入后，在途迟到分片不得发布或复活会话', async () => {
    ;(getUploadSession as jest.Mock).mockResolvedValue({
      fileName: 'tiny-video.mp4',
      fileSize: 4,
      totalChunks: 1,
      chunkSize: 4,
      chunks: new Map(),
      createdAt: Date.now(),
      userId: 'user-1',
    })

    let lateController: ReadableStreamDefaultController<Uint8Array> | null = null
    const lateBody = new ReadableStream<Uint8Array>({
      start(controller) {
        lateController = controller
        controller.enqueue(new Uint8Array([3, 3]))
      },
    })
    const lateRequest = new NextRequest('http://localhost/api/materials/upload/chunk', {
      method: 'POST',
      headers: {
        'content-length': '4',
        'content-type': 'application/octet-stream',
        'x-file-name': 'tiny-video.mp4',
        'x-upload-id': 'upload-test-1',
        'x-chunk-index': '0',
        'x-chunk-size': '4',
      },
      body: lateBody,
      duplex: 'half',
    } as any)

    const lateResponsePromise = POST(lateRequest)
    const tempDir = path.join(uploadRoot, 'temp', 'upload-test-1')
    await waitFor(async () => (
      (await readdir(tempDir).catch(() => [])).some((name) => name.includes('.partial-'))
    ))

    const release = await acquireMaterialUploadStorageLock('exclusive')
    await cancelUploadSessionSeal('upload-test-1')
    await release()

    lateController!.enqueue(new Uint8Array([3, 3]))
    lateController!.close()
    const response = await lateResponsePromise
    const body = await response.json()

    expect(response.status).toBe(409)
    expect(body.code).toBe('UPLOAD_SESSION_CANCELLED')
    await expect(readFile(path.join(tempDir, 'chunk-0'))).rejects.toMatchObject({
      code: 'ENOENT',
    })
    expect((await readdir(tempDir)).filter((name) => name.includes('.partial-'))).toEqual([])
    expect(updateChunkWithStats).not.toHaveBeenCalled()
  })
})

async function waitFor(predicate: () => Promise<boolean>, timeoutMs = 2_000): Promise<void> {
  const deadline = Date.now() + timeoutMs
  while (Date.now() < deadline) {
    if (await predicate()) {
      return
    }
    await new Promise((resolve) => setTimeout(resolve, 10))
  }
  throw new Error('等待条件超时')
}