Gzip.lean

import Zip.Native.Inflate
import Zip.Native.InflateTreeFree
import Zip.Native.InflateFast
import Zip.Native.DeflateDynamic
import Zip.Native.Crc32
import Zip.Native.Adler32
import ZipCommon.Binary

/-!
  Pure Lean gzip (RFC 1952) and zlib (RFC 1950) compression and decompression.

  Compression wraps native DEFLATE output with gzip/zlib framing headers,
  trailers, and checksums. Decompression parses the framing, inflates the
  DEFLATE stream, and verifies checksums.
-/

namespace Zip.Native

namespace GzipDecode

/-- Absolute ceiling on the exact-size fastloop's speculative presize allocation,
    mirroring `Zip.Archive.nativePresizeCap`. A member whose declared decompressed
    size (gzip trailer `ISIZE`) exceeds this keeps the push decoder rather than
    pre-extending a large buffer up front. -/
def presizeCap : Nat := 64 * 1024 * 1024

/-- Scan forward from `pos` in `data` for the next zero byte (NUL).
    Returns the index of the zero byte, or `data.size` if none is found. -/
def scanToZero (data : ByteArray) (pos : Nat) : Nat :=
  if h : pos < data.size then
    if data[pos] == 0 then pos
    else scanToZero data (pos + 1)
  else pos
termination_by data.size - pos

/-- Decompress a gzip stream (RFC 1952). Supports concatenated members.
    Returns the decompressed data.

    `maxOutputSize` (default 1 GiB) caps the *total* output across all
    concatenated members as a zip-bomb guard. Unlike the FFI path, where
    `maxDecompressedSize := 0` means unlimited, here `0` rejects any
    non-empty output (the inner inflate guards compare
    `output.size + len > maxOutputSize`). The outer-loop guard raises an
    `Except` error containing `"Gzip: total output exceeds maximum size"`;
    the inner per-member `Inflate.inflateRaw` call also enforces the
    bound and may surface `"Inflate: output exceeds maximum size"` first.
    See `SECURITY_INVENTORY.md` *Decompression Limit Inventory*. -/
def decompress (data : ByteArray) (maxOutputSize : Nat := 1024 * 1024 * 1024) :
    Except String ByteArray := do
  if data.size < 10 then throw "Gzip: input too short for gzip header"
  let mut pos : Nat := 0
  let mut result : ByteArray := .empty
  -- Tracks the very first member. `result.size == 0` is *not* a first-member test:
  -- an empty member (or a run of them) leaves `result` empty, so it would let every
  -- empty-prefix member re-attempt the speculative presize. A dedicated flag caps
  -- the fastloop attempt at exactly one per stream.
  let mut firstMember : Bool := true
  -- Process concatenated gzip members
  for _ in [:1000] do
    if pos ≥ data.size then return result
    -- Parse header (need 10 bytes: ID1, ID2, CM, FLG, MTIME[4], XFL, OS)
    if hHdr : pos + 10 ≤ data.size then
      -- Read all four header bytes up front so the bound `hHdr` is in scope
      -- (do-notation rebinds `mut pos` after `unless`, losing the proof).
      let id1 := data[pos]
      let id2 := data[pos + 1]
      let cm := data[pos + 2]
      let flg := data[pos + 3]
      unless id1 == 0x1f && id2 == 0x8b do
        if pos == 0 then throw "Gzip: invalid magic bytes"
        -- End of concatenated stream
        return result
      unless cm == 8 do throw "Gzip: unsupported compression method"
      -- Skip MTIME (4), XFL (1), OS (1)
      pos := pos + 10
      -- FEXTRA
      if flg &&& 0x04 != 0 then
        if pos + 2 > data.size then throw "Gzip: truncated FEXTRA length"
        let xlen := (Binary.readUInt16LE data pos).toNat
        pos := pos + 2 + xlen
      -- FNAME (null-terminated)
      if flg &&& 0x08 != 0 then
        pos := scanToZero data pos
        pos := pos + 1 -- skip NUL
      -- FCOMMENT (null-terminated)
      if flg &&& 0x10 != 0 then
        pos := scanToZero data pos
        pos := pos + 1
      -- FHCRC (2-byte header CRC)
      if flg &&& 0x02 != 0 then pos := pos + 2
      if pos > data.size then throw "Gzip: header extends past end of input"
      -- Inflate (cap each member to remaining budget so total stays within maxOutputSize)
      let memberMax := maxOutputSize - result.size
      -- Exact-size fastloop for the single-member case (tar.gz, the dominant real
      -- workload). For a stream with exactly one member, the trailing 8 bytes are
      -- this member's trailer, so `ISIZE = readUInt32LE data (size - 4)` is the
      -- exact decompressed length. We attempt the verified branch-free `uset`
      -- fastloop (`inflateRawSized … (exact := true)`) only on the *first* member
      -- (`firstMember`) — for a later member the trailing ISIZE is not this member's
      -- size — and only when the hint is bounded by `presizeCap` and the member
      -- budget (`exact` conjuncts). The `memberMax < UInt32.size` conjunct keeps
      -- ISIZE (which is the size *mod 2^32*) a genuinely exact hint: with the output
      -- bounded below 4 GiB, `ISIZE = decompressed.size` outright, no 2^32 residue
      -- ambiguity, so a matching hint fires the fast path rather than wasting a
      -- decode that would reject. `inflateRawSized` is proven equal to `inflateRaw`
      -- for every input in the `USize` addressability regime
      -- (`Zip.Native.inflateRawSized_agrees`, under `data.size < USize.size` — always
      -- true for an in-memory `ByteArray` on a 64-bit target, the regime every native
      -- decode proof here assumes — and `memberMax < USize.size` from the conjunct,
      -- with `pos ≤ data.size` from the guard above). So a wrong hint — a concatenated
      -- multi-member stream, trailing padding, or a malicious ISIZE — makes the
      -- fastloop's exact-size contract reject and fall back to the push `inflateRaw`,
      -- never changing the decoded bytes or the returned `endPos`. The CRC32 / ISIZE
      -- trailer checks below stay integrity checks, not soundness backstops.
      let (decompressed, endPos) ←
        if firstMember then
          let isize := (Binary.readUInt32LE data (data.size - 4)).toNat
          let sizeHint := min isize presizeCap
          let exact := isize == sizeHint && isize ≤ memberMax
            && memberMax < UInt32.size && memberMax < USize.size
          Inflate.inflateRawSized data pos memberMax (sizeHint := sizeHint) (exact := exact)
        else
          Inflate.inflateRaw data pos memberMax
      firstMember := false
      pos := endPos
      -- Parse trailer: CRC32 (4 bytes LE) + ISIZE (4 bytes LE)
      if pos + 8 > data.size then throw "Gzip: truncated trailer"
      let expectedCrc := Binary.readUInt32LE data pos
      let expectedSize := Binary.readUInt32LE data (pos + 4)
      pos := pos + 8
      -- Verify CRC32
      let actualCrc := Crc32.Native.crc32 0 decompressed
      unless actualCrc == expectedCrc do
        throw s!"Gzip: CRC32 mismatch: expected {expectedCrc}, got {actualCrc}"
      -- Verify size (mod 2^32)
      let actualSize := decompressed.size.toUInt32
      unless actualSize == expectedSize do
        throw s!"Gzip: size mismatch: expected {expectedSize}, got {actualSize}"
      result := result ++ decompressed
      if result.size > maxOutputSize then
        throw "Gzip: total output exceeds maximum size"
    else
      throw "Gzip: truncated header"
  throw "Gzip: too many concatenated members"

end GzipDecode

namespace GzipEncode

/-- Compress data to gzip format (RFC 1952).
    Level 0 = stored, 1 = fixed Huffman, 2–4 = lazy LZ77, 5+ = dynamic Huffman. -/
def compress (data : ByteArray) (level : UInt8 := 1) : ByteArray :=
  let deflated := Deflate.deflateRaw data level
  -- Gzip header: ID1=0x1f, ID2=0x8b, CM=8, FLG=0, MTIME=0, XFL, OS=255
  let xfl : UInt8 := if level == 0 then 0x00 else if level ≥ 5 then 0x02 else 0x04
  let header := ByteArray.mk #[0x1f, 0x8b, 8, 0, 0, 0, 0, 0, xfl, 0xFF]
  -- CRC32 of uncompressed data (4 bytes LE)
  let crc := Crc32.Native.crc32 0 data
  -- ISIZE: original size mod 2^32 (4 bytes LE)
  let isize := data.size.toUInt32
  let trailer := ByteArray.mk #[
    (crc &&& 0xFF).toUInt8, ((crc >>> 8) &&& 0xFF).toUInt8,
    ((crc >>> 16) &&& 0xFF).toUInt8, ((crc >>> 24) &&& 0xFF).toUInt8,
    (isize &&& 0xFF).toUInt8, ((isize >>> 8) &&& 0xFF).toUInt8,
    ((isize >>> 16) &&& 0xFF).toUInt8, ((isize >>> 24) &&& 0xFF).toUInt8
  ]
  header ++ deflated ++ trailer

end GzipEncode

namespace ZlibDecode

/-- Decompress a zlib stream (RFC 1950).
    Returns the decompressed data.

    `maxOutputSize` (default 1 GiB) is forwarded to the inner
    `Inflate.inflateRaw`; this layer adds no separate guard. Unlike the
    FFI path, where `maxDecompressedSize := 0` means unlimited, here `0`
    rejects any non-empty output (the inflate guards compare
    `output.size + len > maxOutputSize`). Overflow raises an `Except`
    error containing `"Inflate: output exceeds maximum size"`.
    See `SECURITY_INVENTORY.md` *Decompression Limit Inventory*. -/
def decompress (data : ByteArray) (maxOutputSize : Nat := 1024 * 1024 * 1024) :
    Except String ByteArray := do
  if hSz : data.size < 6 then throw "Zlib: input too short"
  else
    -- Parse header: CMF (1 byte) + FLG (1 byte)
    let cmf := data[0]
    let flg := data[1]
    -- Check header checksum
    let check := cmf.toUInt16 * 256 + flg.toUInt16
    unless check % 31 == 0 do throw "Zlib: header check failed"
    -- CM must be 8 (deflate)
    unless cmf &&& 0x0F == 8 do throw "Zlib: unsupported compression method"
    -- CINFO (window size) must be ≤ 7
    let cinfo := cmf >>> 4
    unless cinfo ≤ 7 do throw s!"Zlib: invalid window size {cinfo}"
    let mut pos : Nat := 2
    -- FDICT: preset dictionary (not supported)
    if flg &&& 0x20 != 0 then
      throw "Zlib: preset dictionaries not supported"
    -- Inflate
    let (decompressed, endPos) ← Inflate.inflateRaw data pos maxOutputSize
    pos := endPos
    -- Parse trailer: Adler32 (4 bytes big-endian)
    if hT : pos + 4 ≤ data.size then
      let b0 := data[pos].toUInt32
      let b1 := data[pos + 1].toUInt32
      let b2 := data[pos + 2].toUInt32
      let b3 := data[pos + 3].toUInt32
      let expectedAdler := (b0 <<< 24) ||| (b1 <<< 16) ||| (b2 <<< 8) ||| b3
      -- Verify Adler32
      let actualAdler := Adler32.Native.adler32 1 decompressed
      unless actualAdler == expectedAdler do
        throw s!"Zlib: Adler32 mismatch: expected {expectedAdler}, got {actualAdler}"
      return decompressed
    else
      throw "Zlib: truncated trailer"

end ZlibDecode

namespace ZlibEncode

/-- Compress data to zlib format (RFC 1950).
    Level 0 = stored, 1 = fixed Huffman, 2–4 = lazy LZ77, 5+ = dynamic Huffman. -/
def compress (data : ByteArray) (level : UInt8 := 1) : ByteArray :=
  let deflated := Deflate.deflateRaw data level
  -- CMF: CM=8 (deflate), CINFO=7 (32K window)
  let cmf : UInt8 := 0x78
  -- FLG: FLEVEL (bits 6-7) + FCHECK (bits 0-4) such that (CMF*256 + FLG) % 31 == 0
  -- FLEVEL: 0=fastest, 1=fast (1-4), 2=default (5-8), 3=maximum (9)
  let flevel : UInt8 := if level == 0 then 0x00
    else if level < 5 then 0x40
    else if level < 9 then 0x80
    else 0xC0
  let fcheck := (31 - ((cmf.toNat * 256 + flevel.toNat) % 31)) % 31
  let flg := flevel ||| fcheck.toUInt8
  let header := ByteArray.mk #[cmf, flg]
  -- Adler32 of uncompressed data (4 bytes big-endian)
  let adler := Adler32.Native.adler32 1 data
  let trailer := ByteArray.mk #[
    ((adler >>> 24) &&& 0xFF).toUInt8, ((adler >>> 16) &&& 0xFF).toUInt8,
    ((adler >>> 8) &&& 0xFF).toUInt8, (adler &&& 0xFF).toUInt8
  ]
  header ++ deflated ++ trailer

end ZlibEncode

/-- Format detected from the first bytes of a compressed stream. -/
inductive CompressFormat where
  | gzip
  | zlib
  | rawDeflate

/-- Detect the compression format from the first bytes.
    - Gzip: starts with 0x1f 0x8b
    - Zlib: first byte has CM=8 (low nibble), and header check passes
    - Raw deflate: fallback -/
def detectFormat (data : ByteArray) : CompressFormat :=
  if h : data.size ≥ 2 then
    if data[0] == 0x1f && data[1] == 0x8b then
      .gzip
    else
      let cmf := data[0]
      let flg := data[1]
      let check := cmf.toUInt16 * 256 + flg.toUInt16
      if cmf &&& 0x0F == 8 && check % 31 == 0 then .zlib
      else .rawDeflate
  else
    .rawDeflate

/-- Decompress data by auto-detecting the format (gzip, zlib, or raw deflate).

    `maxOutputSize` (default 1 GiB) is forwarded to whichever of
    `GzipDecode.decompress`, `ZlibDecode.decompress`, or `Inflate.inflate`
    the dispatch picks based on `detectFormat`. The surfaced error
    substring depends on the dispatch: `"Gzip: total output exceeds
    maximum size"` (gzip outer guard), or `"Inflate: output exceeds
    maximum size"` (zlib, raw deflate, or any inner inflate guard).
    Unlike the FFI path, where `maxDecompressedSize := 0` means unlimited,
    here `0` rejects any non-empty output.
    See `SECURITY_INVENTORY.md` *Decompression Limit Inventory*. -/
def decompressAuto (data : ByteArray) (maxOutputSize : Nat := 1024 * 1024 * 1024) :
    Except String ByteArray :=
  match detectFormat data with
  | .gzip => GzipDecode.decompress data maxOutputSize
  | .zlib => ZlibDecode.decompress data maxOutputSize
  | .rawDeflate => Inflate.inflate data maxOutputSize

/-- Compress data with format selection.
    Default: gzip format, level 1 (fixed Huffman). -/
def compressAuto (data : ByteArray)
    (format : CompressFormat := .gzip) (level : UInt8 := 1) :
    ByteArray :=
  match format with
  | .gzip => GzipEncode.compress data level
  | .zlib => ZlibEncode.compress data level
  | .rawDeflate => Deflate.deflateRaw data level

end Zip.Native