diablo2-web/src/netproto/d2gs/framing-s2c.ts

import { ProtocolError } from '../core/errors.ts'
import {
  D2gsHuffmanCodec,
  defaultD2gsHuffmanCodec,
  parseD2gsBlockHeader,
  unpackCustomHuffmanLengths,
} from './compression.ts'
import { S2C_PACKET_SIZES, resolveS2cPacketSize as resolveS2cSizeStruct } from './tables/s2c-sizes.ts'

export { S2C_PACKET_SIZES }
export const D2GS_S2C_PACKET_SIZES = S2C_PACKET_SIZES

/**
 * Convenience helper returning the resolved S->C packet byte length, or `null` if incomplete,
 * or throwing `ProtocolError` if invalid.
 */
export function resolveS2cPacketSize(buf: Uint8Array, offset = 0): number | null {
  const res = resolveS2cSizeStruct(buf, offset)
  if (res.status === 'incomplete') return null
  if (res.status === 'invalid') {
    throw new ProtocolError(`Invalid D2GS S->C packet: ${res.reason}`, res.opcode, offset)
  }
  return res.size
}

export interface D2gsFramedS2cPacket {
  readonly packetId: number
  readonly bytes: Uint8Array
  readonly isSystemPacket: boolean
}

export interface D2gsFramePushResult {
  readonly packets: readonly D2gsFramedS2cPacket[]
  readonly compressedBlocksProcessed: number
}

/**
 * Stateful D2GS S->C stream framer matching `D2Net.dll` (`0x6fbf7290..0x6fbf7470`, `0x6fbf6d60..0x6fbf6ec0`):
 * 1. Starts in uncompressed mode (`compressedMode = false`) for the initial `[0xAF, ...]` server greeting.
 * 2. When packet `0xAF` is parsed with `bytes[1] !== 0`:
 *    - Transitions `compressedMode = true` for all subsequent stream reads.
 *    - If `bytes[1] === 0x81` (`130` bytes), rebuilds the Huffman table via `Fog.10219` (`unpackCustomHuffmanLengths`).
 *    - If `bytes[1] === 0x01` (`2` bytes `[0xAF, 0x01]`), keeps the static `Fog.dll` (`0x6ff7f210`) Huffman table.
 * 3. In compressed mode, parses 1- or 2-byte `totalFrameSize` block headers (`compressedPayloadLen = totalFrameSize - headerSize`),
 *    decompresses each block with `Fog.10224`, appends to the packet reassembly buffer, and splits complete packets via `resolveS2cPacketSize`.
 * 4. Fails fast with `ProtocolError` on unknown/unused opcodes (`S2C_PACKET_SIZES[opcode] === 0` or `opcode >= 0xB5`) or invalid sizes (`<= 0` or `> 0x204`).
 */
export class D2gsS2cFramer {
  private tcpBuf: Uint8Array = new Uint8Array(0)
  private packetBuf: Uint8Array = new Uint8Array(0)
  private compressedMode = false
  private huffmanCodec: D2gsHuffmanCodec

  constructor(initialCompressedMode = false, huffmanCodec: D2gsHuffmanCodec = defaultD2gsHuffmanCodec) {
    this.compressedMode = initialCompressedMode
    this.huffmanCodec = huffmanCodec
  }

  get isCompressedMode(): boolean {
    return this.compressedMode
  }

  get isCompressed(): boolean {
    return this.compressedMode
  }

  get bufferedTcpBytes(): number {
    return this.tcpBuf.length
  }

  get bufferedPacketBytes(): number {
    return this.packetBuf.length
  }

  reset(initialCompressedMode = false): void {
    this.tcpBuf = new Uint8Array(0)
    this.packetBuf = new Uint8Array(0)
    this.compressedMode = initialCompressedMode
    this.huffmanCodec = defaultD2gsHuffmanCodec
  }

  push(chunk: Uint8Array): D2gsFramedS2cPacket[] {
    if (chunk.length > 0) {
      this.tcpBuf = concatBytes(this.tcpBuf, chunk)
    }

    const out: D2gsFramedS2cPacket[] = []

    // Phase 1: If still in uncompressed mode (waiting for 0xAF greeting), split directly from tcpBuf
    while (!this.compressedMode && this.tcpBuf.length > 0) {
      const opcode = this.tcpBuf[0]!
      if (opcode >= S2C_PACKET_SIZES.length || S2C_PACKET_SIZES[opcode] === 0) {
        throw new ProtocolError(
          `Unknown or invalid D2GS S->C opcode 0x${opcode.toString(16).padStart(2, '0')} in uncompressed stream`,
          opcode,
          0,
        )
      }
      const pktSize = resolveS2cPacketSize(this.tcpBuf, 0)
      if (pktSize === null) {
        // Need more bytes to resolve variable-length header or complete fixed-size packet
        break
      }
      if (pktSize <= 0 || pktSize > 0x204) {
        throw new ProtocolError(
          `Invalid D2GS S->C packet size ${pktSize} for opcode 0x${opcode.toString(16).padStart(2, '0')}`,
          opcode,
          0,
        )
      }
      const pktBytes = this.tcpBuf.slice(0, pktSize)
      this.tcpBuf = this.tcpBuf.subarray(pktSize)
      this.handleSpecialPacket(pktBytes)
      out.push({
        packetId: opcode,
        bytes: pktBytes,
        isSystemPacket: opcode >= 0xaf,
      })
    }

    // Phase 2: In compressed mode, parse 1/2-byte block headers, decompress blocks, and split packets
    if (this.compressedMode) {
      while (this.tcpBuf.length > 0) {
        const header = parseD2gsBlockHeader(this.tcpBuf, 0)
        if (!header) {
          break
        }
        if (this.tcpBuf.length < header.totalFrameSize) {
          break
        }
        const compressedPayload = this.tcpBuf.subarray(header.headerSize, header.totalFrameSize)
        this.tcpBuf = this.tcpBuf.subarray(header.totalFrameSize)

        const decompressed = this.huffmanCodec.decompress(compressedPayload)
        if (decompressed.length > 0) {
          this.packetBuf = concatBytes(this.packetBuf, decompressed)
        }
        this.drainPacketBuffer(out)
      }
    }

    return out
  }

  /**
   * Splits already-decompressed S->C packet bytes into individual packets, failing fast with `ProtocolError`
   * on unknown opcodes or invalid packet sizes.
   */
  splitDecompressedBlock(decompressed: Uint8Array): D2gsFramedS2cPacket[] {
    if (decompressed.length > 0) {
      this.packetBuf = concatBytes(this.packetBuf, decompressed)
    }
    const out: D2gsFramedS2cPacket[] = []
    this.drainPacketBuffer(out)
    return out
  }

  private drainPacketBuffer(out: D2gsFramedS2cPacket[]): void {
    let offset = 0
    while (offset < this.packetBuf.length) {
      const opcode = this.packetBuf[offset]!
      if (opcode >= S2C_PACKET_SIZES.length || S2C_PACKET_SIZES[opcode] === 0) {
        const consumed = offset
        this.packetBuf = this.packetBuf.subarray(offset)
        throw new ProtocolError(
          `Unknown or invalid D2GS S->C opcode 0x${opcode.toString(16).padStart(2, '0')} at offset ${consumed}`,
          opcode,
          consumed,
        )
      }
      let pktSize: number | null
      try {
        pktSize = resolveS2cPacketSize(this.packetBuf, offset)
      } catch (err) {
        this.packetBuf = this.packetBuf.subarray(offset)
        throw err
      }
      if (pktSize === null) {
        break
      }
      if (pktSize <= 0 || pktSize > 0x204) {
        const consumed = offset
        this.packetBuf = this.packetBuf.subarray(offset)
        throw new ProtocolError(
          `Invalid D2GS S->C packet size ${pktSize} for opcode 0x${opcode.toString(16).padStart(2, '0')} at offset ${consumed}`,
          opcode,
          consumed,
        )
      }
      const pktBytes = this.packetBuf.slice(offset, offset + pktSize)
      offset += pktSize
      this.handleSpecialPacket(pktBytes)
      out.push({
        packetId: opcode,
        bytes: pktBytes,
        isSystemPacket: opcode >= 0xaf,
      })
    }
    if (offset > 0) {
      this.packetBuf = this.packetBuf.subarray(offset)
    }
  }

  private handleSpecialPacket(pktBytes: Uint8Array): void {
    const opcode = pktBytes[0]!
    if (opcode === 0xaf && pktBytes.length >= 2) {
      const flag = pktBytes[1]!
      if (flag !== 0) {
        this.compressedMode = true
        if (flag === 0x81) {
          const customLengths = unpackCustomHuffmanLengths(pktBytes)
          this.huffmanCodec = new D2gsHuffmanCodec(customLengths)
        }
      }
    }
  }
}

function concatBytes(a: Uint8Array, b: Uint8Array): Uint8Array {
  if (a.length === 0) return b.slice()
  if (b.length === 0) return a
  const out = new Uint8Array(a.length + b.length)
  out.set(a, 0)
  out.set(b, a.length)
  return out
}