Lambdaの/tmpを使わずにS3上の数十GBのZIPを解凍してみた

/** ZIPの [start, end) の範囲を読む関数。S3ならRange GETで実装する */
export type ReadRange = (start: number, end: number) => Promise<Buffer>;

/** 目録から得たファイルの情報 */
export type ZipEntry = {
  /** ZIP内のパス */
  name: string;
  /** 圧縮方式（0: STORED, 8: DEFLATE） */
  method: number;
  /** 汎用ビットフラグ */
  flags: number;
  /** 展開後データのCRC-32 */
  crc32: number;
  /** ZIP内で占めるバイト数 */
  compressedSize: number;
  /** 展開後のバイト数 */
  uncompressedSize: number;
  /** ローカルファイルヘッダの位置（ZIP先頭からのバイト数） */
  localHeaderOffset: number;
};

export const METHOD_STORED = 0;
export const METHOD_DEFLATE = 8;

const SIG_EOCD = 0x06054b50;
const SIG_ZIP64_EOCD_LOCATOR = 0x07064b50;
const SIG_ZIP64_EOCD = 0x06064b50;
const SIG_CENTRAL_FILE_HEADER = 0x02014b50;
const SIG_LOCAL_FILE_HEADER = 0x04034b50;

const EOCD_MIN_SIZE = 22;
const MAX_COMMENT_SIZE = 0xffff;
const ZIP64_EOCD_LOCATOR_SIZE = 20;
const ZIP64_EOCD_MIN_SIZE = 56;
const CENTRAL_FILE_HEADER_SIZE = 46;
const LOCAL_FILE_HEADER_SIZE = 30;
const ZIP64_EXTRA_FIELD_ID = 0x0001;
const UINT16_MAX = 0xffff;
const UINT32_MAX = 0xffffffff;
const FLAG_ENCRYPTED = 0x0001;

const readUInt64 = (buf: Buffer, offset: number): number => {
  const value = Number(buf.readBigUInt64LE(offset));
  if (!Number.isSafeInteger(value)) {
    throw new Error(`64bit value exceeds MAX_SAFE_INTEGER at ${offset}`);
  }
  return value;
};

// 終端レコード（EOCD）を探す。後ろに最大65535バイトのコメントが付くので、末尾をまとめて読んで後ろから探す
const findEocd = async (
  readRange: ReadRange,
  zipSize: number
): Promise<{ tail: Buffer; tailStart: number; eocdPos: number }> => {
  if (zipSize < EOCD_MIN_SIZE) {
    throw new Error('not a zip: too small');
  }
  const tailStart = Math.max(0, zipSize - EOCD_MIN_SIZE - MAX_COMMENT_SIZE);
  const tail = await readRange(tailStart, zipSize);
  for (let pos = tail.length - EOCD_MIN_SIZE; pos >= 0; pos--) {
    if (tail.readUInt32LE(pos) === SIG_EOCD) {
      return { tail, tailStart, eocdPos: pos };
    }
  }
  throw new Error('not a zip: end of central directory not found');
};

// 目録の位置・サイズ・件数を得る。EOCDの値が上限値ならZIP64 EOCDを読みにいく
const locateCentralDirectory = async (
  readRange: ReadRange,
  zipSize: number
): Promise<{ offset: number; size: number; count: number }> => {
  const { tail, tailStart, eocdPos } = await findEocd(readRange, zipSize);

  const count = tail.readUInt16LE(eocdPos + 10);
  const size = tail.readUInt32LE(eocdPos + 12);
  const offset = tail.readUInt32LE(eocdPos + 16);
  const needsZip64 =
    count === UINT16_MAX || size === UINT32_MAX || offset === UINT32_MAX;

  const locatorPos = eocdPos - ZIP64_EOCD_LOCATOR_SIZE;
  const hasLocator =
    locatorPos >= 0 && tail.readUInt32LE(locatorPos) === SIG_ZIP64_EOCD_LOCATOR;

  if (!hasLocator) {
    if (needsZip64) {
      throw new Error('broken zip: zip64 locator not found');
    }
    return { offset, size, count };
  }

  // ロケータにはZIP64 EOCDの絶対位置が書いてある。末尾バッファ内に収まっていれば読み直さない
  const zip64EocdOffset = readUInt64(tail, locatorPos + 8);
  const posInTail = zip64EocdOffset - tailStart;
  const zip64Eocd =
    posInTail >= 0
      ? tail.subarray(posInTail, posInTail + ZIP64_EOCD_MIN_SIZE)
      : await readRange(zip64EocdOffset, zip64EocdOffset + ZIP64_EOCD_MIN_SIZE);

  if (zip64Eocd.readUInt32LE(0) !== SIG_ZIP64_EOCD) {
    throw new Error('broken zip: zip64 end of central directory not found');
  }
  return {
    count: readUInt64(zip64Eocd, 32),
    size: readUInt64(zip64Eocd, 40),
    offset: readUInt64(zip64Eocd, 48),
  };
};

// ZIP64拡張フィールドには「ヘッダ側が0xFFFFFFFFの項目だけ」が
// 元サイズ → 圧縮後サイズ → ローカルヘッダ位置 の順に入る
const applyZip64Extra = (
  extra: Buffer,
  header: {
    uncompressedSize: number;
    compressedSize: number;
    localHeaderOffset: number;
  }
): {
  uncompressedSize: number;
  compressedSize: number;
  localHeaderOffset: number;
} => {
  const result = { ...header };
  const needs = [
    header.uncompressedSize,
    header.compressedSize,
    header.localHeaderOffset,
  ].some((v) => v === UINT32_MAX);
  if (!needs) {
    return result;
  }

  let pos = 0;
  while (pos + 4 <= extra.length) {
    const id = extra.readUInt16LE(pos);
    const len = extra.readUInt16LE(pos + 2);
    if (id === ZIP64_EXTRA_FIELD_ID) {
      let p = pos + 4;
      if (header.uncompressedSize === UINT32_MAX) {
        result.uncompressedSize = readUInt64(extra, p);
        p += 8;
      }
      if (header.compressedSize === UINT32_MAX) {
        result.compressedSize = readUInt64(extra, p);
        p += 8;
      }
      if (header.localHeaderOffset === UINT32_MAX) {
        result.localHeaderOffset = readUInt64(extra, p);
      }
      return result;
    }
    pos += 4 + len;
  }
  throw new Error('broken zip: zip64 extra field not found');
};

// 目録を読み、全ファイルの情報を返す。ファイル名はUTF-8として解釈する
export const readCentralDirectory = async (
  readRange: ReadRange,
  zipSize: number
): Promise<ZipEntry[]> => {
  const { offset, size, count } = await locateCentralDirectory(
    readRange,
    zipSize
  );
  if (offset + size > zipSize) {
    throw new Error('broken zip: central directory out of range');
  }
  const cd = await readRange(offset, offset + size);

  const entries: ZipEntry[] = [];
  let pos = 0;
  for (let i = 0; i < count; i++) {
    if (
      pos + CENTRAL_FILE_HEADER_SIZE > cd.length ||
      cd.readUInt32LE(pos) !== SIG_CENTRAL_FILE_HEADER
    ) {
      throw new Error(`broken zip: central file header #${i} not found`);
    }
    const nameLen = cd.readUInt16LE(pos + 28);
    const extraLen = cd.readUInt16LE(pos + 30);
    const commentLen = cd.readUInt16LE(pos + 32);
    const nameStart = pos + CENTRAL_FILE_HEADER_SIZE;
    const extraStart = nameStart + nameLen;

    const sizes = applyZip64Extra(
      cd.subarray(extraStart, extraStart + extraLen),
      {
        compressedSize: cd.readUInt32LE(pos + 20),
        uncompressedSize: cd.readUInt32LE(pos + 24),
        localHeaderOffset: cd.readUInt32LE(pos + 42),
      }
    );

    entries.push({
      name: cd.subarray(nameStart, extraStart).toString('utf8'),
      method: cd.readUInt16LE(pos + 10),
      flags: cd.readUInt16LE(pos + 8),
      crc32: cd.readUInt32LE(pos + 16),
      ...sizes,
    });
    pos = extraStart + extraLen + commentLen;
  }
  return entries;
};

// データ本体の開始位置を求める。ローカルヘッダの拡張フィールド長は目録側と一致するとは限らないため、
// ローカルヘッダの固定部（30バイト）を読んで計算する
export const resolveDataOffset = async (
  readRange: ReadRange,
  entry: ZipEntry
): Promise<number> => {
  if (entry.flags & FLAG_ENCRYPTED) {
    throw new Error(`encrypted entry is not supported: ${entry.name}`);
  }
  const header = await readRange(
    entry.localHeaderOffset,
    entry.localHeaderOffset + LOCAL_FILE_HEADER_SIZE
  );
  if (header.readUInt32LE(0) !== SIG_LOCAL_FILE_HEADER) {
    throw new Error(`broken zip: local file header not found: ${entry.name}`);
  }
  const nameLen = header.readUInt16LE(26);
  const extraLen = header.readUInt16LE(28);
  return entry.localHeaderOffset + LOCAL_FILE_HEADER_SIZE + nameLen + extraLen;
};

import {
  GetObjectCommand,
  HeadObjectCommand,
  type S3Client,
} from '@aws-sdk/client-s3';
import {
  METHOD_STORED,
  type ReadRange,
  readCentralDirectory,
  resolveDataOffset,
  type ZipEntry,
} from './central-directory.ts';
import { extractByCopy } from './copy-extract.ts';
import { extractByStreaming } from './stream-extract.ts';

const MiB = 1024 * 1024;
// どちらも今回の計測で使った値
const COPY_OPTIONS = { partSize: 64 * MiB, concurrency: 48 };
const STREAM_OPTIONS = { partSize: 32 * MiB, concurrency: 8, windowSize: 64 * MiB };

/** 1ファイルを展開するための入力 */
export type ExtractInput = {
  /** S3クライアント */
  client: S3Client;
  /** バケット名（ZIPと展開先で共通） */
  bucket: string;
  /** ZIPのキー */
  zipKey: string;
  /** 目録から得たファイルの情報 */
  entry: ZipEntry;
  /** データ本体の開始位置 */
  dataOffset: number;
  /** 展開先のキー */
  destKey: string;
};

// 目録やローカルヘッダのような小さな範囲を、Range GETで丸ごと読む
const createS3ReadRange = (
  client: S3Client,
  bucket: string,
  key: string
): ReadRange => {
  return async (start, end) => {
    const res = await client.send(
      new GetObjectCommand({
        Bucket: bucket,
        Key: key,
        Range: `bytes=${start}-${end - 1}`,
      })
    );
    if (!res.Body) {
      throw new Error(`empty body: s3://${bucket}/${key}`);
    }
    return Buffer.from(await res.Body.transformToByteArray());
  };
};

export const extractZip = async (
  client: S3Client,
  bucket: string,
  zipKey: string,
  destPrefix: string
): Promise<void> => {
  const head = await client.send(
    new HeadObjectCommand({ Bucket: bucket, Key: zipKey })
  );
  const readRange = createS3ReadRange(client, bucket, zipKey);
  // 名前が `/` で終わるのはディレクトリのエントリなので飛ばす
  const entries = (
    await readCentralDirectory(readRange, head.ContentLength ?? 0)
  ).filter((e) => !e.name.endsWith('/'));

  for (const entry of entries) {
    const input: ExtractInput = {
      client,
      bucket,
      zipKey,
      entry,
      dataOffset: await resolveDataOffset(readRange, entry),
      destKey: `${destPrefix}${entry.name}`,
    };
    if (entry.method === METHOD_STORED) {
      await extractByCopy(input, COPY_OPTIONS);
    } else {
      await extractByStreaming(input, STREAM_OPTIONS);
    }
  }
};

import type { Readable } from 'node:stream';
import { GetObjectCommand, type S3Client } from '@aws-sdk/client-s3';

// 範囲を windowSize ごとのRange GETに分け、前の範囲を読み終えてから次を開く。
// 先に開いたレスポンスを読まずに待たせると、接続が切られるため
export const streamS3Range = async function* (
  client: S3Client,
  bucket: string,
  key: string,
  start: number,
  end: number,
  options: { windowSize: number; maxRetries: number }
): AsyncGenerator<Buffer> {
  const { windowSize, maxRetries } = options;
  let position = start;
  let consecutiveFailures = 0;

  while (position < end) {
    const windowEnd = Math.min(position + windowSize, end);
    const res = await client.send(
      new GetObjectCommand({
        Bucket: bucket,
        Key: key,
        Range: `bytes=${position}-${windowEnd - 1}`,
      })
    );
    const body = res.Body as Readable | undefined;
    if (!body) {
      throw new Error(`empty body: s3://${bucket}/${key}`);
    }

    let interruption: unknown;
    try {
      for await (const chunk of body) {
        const buf = chunk as Buffer;
        position += buf.length;
        consecutiveFailures = 0;
        yield buf;
      }
      if (position < windowEnd) {
        interruption = new Error(`short read at ${position}`);
      }
    } catch (error) {
      interruption = error;
    } finally {
      body.destroy();
    }

    // 途中で切れたら、読み終えた位置（position）から読み直す
    if (interruption === undefined) {
      continue;
    }
    consecutiveFailures++;
    if (consecutiveFailures > maxRetries) {
      throw interruption;
    }
  }
};

import { crc32 } from 'node:zlib';
import {
  AbortMultipartUploadCommand,
  type CompletedPart,
  CompleteMultipartUploadCommand,
  CreateMultipartUploadCommand,
  PutObjectCommand,
  type S3Client,
  UploadPartCommand,
} from '@aws-sdk/client-s3';
import { crc32ToBase64 } from './checksum.ts';

/** 順に届くデータをS3へ書き込むライタ */
export type MultipartWriter = {
  /** データを書き込む。送信待ちのパートが上限に達していれば、空くまで待つ */
  write: (chunk: Buffer) => Promise<void>;
  /** 残りを送ってアップロードを完了する。書き込んだバイト数とCRC-32を返す */
  complete: () => Promise<{ bytes: number; crc32: number }>;
  /** 途中まで送ったパートを破棄する */
  abort: () => Promise<void>;
};

// メモリ使用量は最大で partSize × (concurrency + 1)。ファイルサイズには比例しない。
// 1パートに満たない小さなデータはPutObjectで送る
export const createMultipartWriter = (
  client: S3Client,
  bucket: string,
  key: string,
  options: { partSize: number; concurrency: number }
): MultipartWriter => {
  const { partSize, concurrency } = options;
  let uploadId: string | undefined;
  let buffer = Buffer.allocUnsafe(partSize);
  let filled = 0;
  let partNumber = 0;
  let bytes = 0;
  let totalCrc = 0;
  const completedParts: CompletedPart[] = [];
  const inFlight = new Set<Promise<void>>();
  let failure: unknown;

  // 1パートを送信キューに入れ、キューが上限に達していれば1つ空くまで待つ（ここで背圧がかかる）
  const dispatch = async (body: Buffer): Promise<void> => {
    if (!uploadId) {
      const created = await client.send(
        new CreateMultipartUploadCommand({
          Bucket: bucket,
          Key: key,
          ChecksumAlgorithm: 'CRC32',
          ChecksumType: 'FULL_OBJECT',
        })
      );
      uploadId = created.UploadId;
    }
    const currentPart = ++partNumber;
    const task: Promise<void> = client
      .send(
        new UploadPartCommand({
          Bucket: bucket,
          Key: key,
          UploadId: uploadId,
          PartNumber: currentPart,
          Body: body,
          ChecksumCRC32: crc32ToBase64(crc32(body)),
        })
      )
      .then((res) => {
        completedParts.push({
          PartNumber: currentPart,
          ETag: res.ETag,
          ChecksumCRC32: res.ChecksumCRC32,
        });
      })
      .catch((error: unknown) => {
        failure ??= error;
      })
      .finally(() => {
        inFlight.delete(task);
      });
    inFlight.add(task);

    if (inFlight.size >= concurrency) {
      await Promise.race(inFlight);
    }
    if (failure) {
      throw failure;
    }
  };

  const write = async (chunk: Buffer): Promise<void> => {
    totalCrc = crc32(chunk, totalCrc);
    bytes += chunk.length;
    let offset = 0;
    while (offset < chunk.length) {
      const copied = chunk.copy(buffer, filled, offset);
      filled += copied;
      offset += copied;
      if (filled === partSize) {
        const full = buffer;
        buffer = Buffer.allocUnsafe(partSize);
        filled = 0;
        await dispatch(full);
      }
    }
  };

  const complete = async (): Promise<{ bytes: number; crc32: number }> => {
    const expected = crc32ToBase64(totalCrc);

    if (partNumber === 0) {
      await client.send(
        new PutObjectCommand({
          Bucket: bucket,
          Key: key,
          Body: buffer.subarray(0, filled),
          ChecksumCRC32: expected,
        })
      );
      return { bytes, crc32: totalCrc };
    }

    if (filled > 0) {
      await dispatch(buffer.subarray(0, filled));
    }
    await Promise.all(inFlight);
    if (failure) {
      throw failure;
    }

    completedParts.sort((a, b) => (a.PartNumber ?? 0) - (b.PartNumber ?? 0));
    await client.send(
      new CompleteMultipartUploadCommand({
        Bucket: bucket,
        Key: key,
        UploadId: uploadId,
        MultipartUpload: { Parts: completedParts },
        ChecksumType: 'FULL_OBJECT',
        ChecksumCRC32: expected,
      })
    );
    return { bytes, crc32: totalCrc };
  };

  const abort = async (): Promise<void> => {
    await Promise.allSettled(inFlight);
    if (!uploadId) {
      return;
    }
    await client.send(
      new AbortMultipartUploadCommand({
        Bucket: bucket,
        Key: key,
        UploadId: uploadId,
      })
    );
  };

  return { write, complete, abort };
};

import { Readable } from 'node:stream';
import { pipeline } from 'node:stream/promises';
import { createInflateRaw } from 'node:zlib';
import { METHOD_DEFLATE, METHOD_STORED } from './central-directory.ts';
import type { ExtractInput } from './extract-zip.ts';
import { createMultipartWriter } from './multipart-writer.ts';
import { streamS3Range } from './range-stream.ts';

const MAX_RANGE_RETRIES = 3;

export const extractByStreaming = async (
  input: ExtractInput,
  options: { partSize: number; concurrency: number; windowSize: number }
): Promise<void> => {
  const { client, bucket, zipKey, entry, dataOffset, destKey } = input;
  if (entry.method !== METHOD_DEFLATE && entry.method !== METHOD_STORED) {
    throw new Error(`unsupported method ${entry.method}: ${entry.name}`);
  }

  const writer = createMultipartWriter(client, bucket, destKey, options);
  const source = Readable.from(
    streamS3Range(
      client,
      bucket,
      zipKey,
      dataOffset,
      dataOffset + entry.compressedSize,
      { windowSize: options.windowSize, maxRetries: MAX_RANGE_RETRIES }
    )
  );

  // ライタの送信キューが詰まっている間は、次を読まない
  const sink = async (chunks: AsyncIterable<Buffer>): Promise<void> => {
    for await (const chunk of chunks) {
      await writer.write(chunk);
    }
  };

  try {
    if (entry.method === METHOD_DEFLATE) {
      await pipeline(source, createInflateRaw(), sink);
    } else {
      await pipeline(source, sink);
    }
    const written = await writer.complete();

    // S3側の検証に加えて、Lambda側でも目録の値と照合する
    if (written.bytes !== entry.uncompressedSize) {
      throw new Error(
        `size mismatch: ${entry.name} expected ${entry.uncompressedSize}, got ${written.bytes}`
      );
    }
    if (written.crc32 !== entry.crc32) {
      throw new Error(`crc32 mismatch: ${entry.name}`);
    }
  } catch (error) {
    await writer.abort().catch(() => undefined);
    throw error;
  }
};

import {
  AbortMultipartUploadCommand,
  CompleteMultipartUploadCommand,
  CreateMultipartUploadCommand,
  PutObjectCommand,
  UploadPartCopyCommand,
} from '@aws-sdk/client-s3';
import { crc32ToBase64 } from './checksum.ts';
import { METHOD_STORED } from './central-directory.ts';
import type { ExtractInput } from './extract-zip.ts';

const MiB = 1024 * 1024;
/** UploadPartCopyのパートサイズの下限（最後のパートを除く） */
const MIN_PART_SIZE = 5 * MiB;
/** UploadPartCopyのパートサイズの上限 */
const MAX_PART_SIZE = 5 * 1024 * MiB;
/** マルチパートアップロードのパート数の上限 */
const MAX_PARTS = 10_000;

// 同時実行数を制限しながら、各要素に非同期処理を適用する（結果は入力と同じ順）
const mapWithConcurrency = async <T, R>(
  items: readonly T[],
  limit: number,
  fn: (item: T, index: number) => Promise<R>
): Promise<R[]> => {
  const results = new Array<R>(items.length);
  let next = 0;

  const worker = async (): Promise<void> => {
    while (next < items.length) {
      const index = next++;
      results[index] = await fn(items[index] as T, index);
    }
  };

  await Promise.all(
    Array.from({ length: Math.min(limit, items.length) }, () => worker())
  );
  return results;
};

// コピー元のキーは、パス区切りの `/` を残してURLエンコードする
const toCopySource = (bucket: string, key: string): string => {
  return `${bucket}/${key.split('/').map(encodeURIComponent).join('/')}`;
};

export const extractByCopy = async (
  input: ExtractInput,
  options: { partSize: number; concurrency: number }
): Promise<void> => {
  const { client, bucket, zipKey, entry, dataOffset, destKey } = input;
  const { partSize, concurrency } = options;
  if (entry.method !== METHOD_STORED) {
    throw new Error(`copy strategy requires STORED: ${entry.name}`);
  }
  if (partSize < MIN_PART_SIZE || partSize > MAX_PART_SIZE) {
    throw new Error(`partSize must be between 5MiB and 5GiB: ${partSize}`);
  }

  const size = entry.compressedSize;
  const expectedChecksum = crc32ToBase64(entry.crc32);

  // 空ファイルはパートを作れないので、空のオブジェクトを置く
  if (size === 0) {
    await client.send(
      new PutObjectCommand({
        Bucket: bucket,
        Key: destKey,
        Body: Buffer.alloc(0),
        ChecksumCRC32: expectedChecksum,
      })
    );
    return;
  }

  const ranges: Array<{ start: number; end: number }> = [];
  for (let start = 0; start < size; start += partSize) {
    ranges.push({ start, end: Math.min(start + partSize, size) });
  }
  if (ranges.length > MAX_PARTS) {
    throw new Error(`too many parts: ${ranges.length}`);
  }

  const created = await client.send(
    new CreateMultipartUploadCommand({
      Bucket: bucket,
      Key: destKey,
      ChecksumAlgorithm: 'CRC32',
      ChecksumType: 'FULL_OBJECT',
    })
  );
  const uploadId = created.UploadId;

  try {
    const parts = await mapWithConcurrency(
      ranges,
      concurrency,
      async ({ start, end }, index) => {
        const res = await client.send(
          new UploadPartCopyCommand({
            Bucket: bucket,
            Key: destKey,
            UploadId: uploadId,
            PartNumber: index + 1,
            CopySource: toCopySource(bucket, zipKey),
            CopySourceRange: `bytes=${dataOffset + start}-${dataOffset + end - 1}`,
          })
        );
        return {
          PartNumber: index + 1,
          ETag: res.CopyPartResult?.ETag,
          ChecksumCRC32: res.CopyPartResult?.ChecksumCRC32,
        };
      }
    );

    // 目録のCRC-32を全体のチェックサムとして渡し、S3側でコピー結果と照合させる
    await client.send(
      new CompleteMultipartUploadCommand({
        Bucket: bucket,
        Key: destKey,
        UploadId: uploadId,
        MultipartUpload: { Parts: parts },
        ChecksumType: 'FULL_OBJECT',
        ChecksumCRC32: expectedChecksum,
      })
    );
  } catch (error) {
    await client
      .send(
        new AbortMultipartUploadCommand({
          Bucket: bucket,
          Key: destKey,
          UploadId: uploadId,
        })
      )
      .catch(() => undefined);
    throw error;
  }
};