local t = require('tests.harness')
local archive = require('arti-pack.archive')
local deflate = require('tests.deflate')
local zzlib = require('arti-pack.vendor.zzlib')

local function tmpdir()
  local dir = vim.fn.tempname()
  vim.fn.mkdir(dir, 'p')
  return dir
end

local function le(n, bytes)
  local out = {}
  for _ = 1, bytes do
    table.insert(out, string.char(n % 256))
    n = math.floor(n / 256)
  end
  return table.concat(out)
end

-- Fixtures are built here rather than shelled out to zip/tar so the hostile
-- entry names below are exactly the bytes under test. `deflate = true` runs the
-- payload through a real compressor, so DEFLATE entries are not approximated
-- and the tests cannot skip themselves when no zip binary is around.
local function zipf(entries)
  local locals, central, offset = {}, {}, 0
  for _, e in ipairs(entries) do
    local data = e.data or ''
    local method, payload = 0, data
    if e.deflate then
      method, payload = 8, deflate.compress(data, e.mode)
    end
    method = e.method or method
    if e.mangle then
      payload = e.mangle(payload)
    end
    local crc = e.crc or zzlib.crc32(data)
    local csize, usize = #payload, e.usize or #data
    local head = 'PK\3\4'
      .. le(20, 2)
      .. le(e.flags or 0, 2)
      .. le(method, 2)
      .. le(0, 4)
      .. le(crc, 4)
      .. le(csize, 4)
      .. le(usize, 4)
      .. le(#e.name, 2)
      .. le(0, 2)
      .. e.name
    table.insert(locals, head .. payload)
    table.insert(
      central,
      'PK\1\2'
        .. le(e.made_by or (3 * 256 + 20), 2)
        .. le(20, 2)
        .. le(e.flags or 0, 2)
        .. le(method, 2)
        .. le(0, 4)
        .. le(crc, 4)
        .. le(csize, 4)
        .. le(usize, 4)
        .. le(#e.name, 2)
        .. le(0, 2)
        .. le(0, 2)
        .. le(0, 2)
        .. le(0, 2)
        .. le(e.attr or 0, 4)
        .. le(e.local_offset or offset, 4)
        .. e.name
    )
    offset = offset + #head + #payload
  end
  local body = table.concat(locals)
  local cd = table.concat(central)
  local comment = entries.comment or ''
  local eocd = 'PK\5\6'
    .. le(0, 2)
    .. le(0, 2)
    .. le(#entries, 2)
    .. le(#entries, 2)
    .. le(#cd, 4)
    .. le(#body, 4)
    .. le(#comment, 2)
    .. comment
  return body .. cd .. eocd
end

local function stored_blocks(data)
  local blocks, pos = {}, 1
  repeat
    local chunk = data:sub(pos, pos + 65534)
    pos = pos + #chunk
    local final = pos > #data and 1 or 0
    table.insert(blocks, string.char(final) .. le(#chunk, 2) .. le(65535 - #chunk, 2) .. chunk)
  until pos > #data
  return table.concat(blocks)
end

-- A gzip member. Without a mode the DEFLATE payload is all stored blocks; with
-- one it is genuinely compressed, which is what drives the huffman decoder.
local function gz(data, mode)
  local payload = mode and deflate.compress(data, mode) or stored_blocks(data)
  return '\31\139\8\0\0\0\0\0\0\3' .. payload .. le(zzlib.crc32(data), 4) .. le(#data, 4)
end

local function tar_header(name, size, kind, mode)
  local h = name .. string.rep('\0', 100 - #name)
  h = h .. string.format('%07o\0', mode or 420)
  h = h .. string.format('%07o\0', 0) .. string.format('%07o\0', 0)
  h = h .. string.format('%011o\0', size) .. string.format('%011o\0', 0)
  h = h .. string.rep(' ', 8) .. kind .. string.rep('\0', 100)
  h = h .. 'ustar\0' .. '00' .. string.rep('\0', 32) .. string.rep('\0', 32)
  h = h .. string.rep('\0', 8) .. string.rep('\0', 8) .. string.rep('\0', 155)
  h = h .. string.rep('\0', 512 - #h)
  local sum = 0
  for i = 1, 512 do
    sum = sum + h:byte(i)
  end
  return h:sub(1, 148) .. string.format('%06o\0 ', sum) .. h:sub(157)
end

local function tarf(entries)
  local out = {}
  for _, e in ipairs(entries) do
    local data = e.data or ''
    table.insert(out, tar_header(e.name, #data, e.kind or '0', e.mode))
    if #data > 0 then
      table.insert(out, data .. string.rep('\0', (512 - #data % 512) % 512))
    end
  end
  table.insert(out, string.rep('\0', 1024))
  return table.concat(out)
end

local function write(path, data)
  local fd = assert(io.open(path, 'wb'))
  fd:write(data)
  fd:close()
  return path
end

local function read(path)
  local fd = io.open(path, 'rb')
  if not fd then
    return nil
  end
  local data = fd:read('*a')
  fd:close()
  return data
end

local function extract(bytes, ext)
  local dir = tmpdir()
  local src = write(dir .. '/archive', bytes)
  local dest = dir .. '/out'
  local ok, err = archive.extract(src, dest, ext or 'zip')
  return ok, err, dest
end

-- Deterministic pseudo-random bytes, so a compressed fixture is the same on
-- every run but still spans the whole literal alphabet.
local function noise(n, seed)
  local out, state = {}, seed
  for _ = 1, n do
    state = (state * 1103515245 + 12345) % 2147483648
    table.insert(out, string.char(math.floor(state / 65536) % 256))
  end
  return table.concat(out)
end

-- Every byte value, then repeated text: literal codes of both widths plus real
-- back-references.
local function sample()
  local parts = {}
  for i = 0, 255 do
    table.insert(parts, string.char(i))
  end
  for i = 1, 40 do
    table.insert(parts, string.format('local module_%d = require("arti-pack.mod")\n', i))
  end
  return table.concat(parts)
end

local function patch32(bytes, at, value)
  return bytes:sub(1, at - 1) .. le(value, 4) .. bytes:sub(at + 4)
end

local function flip(bytes, at, count)
  local out = bytes:sub(1, at - 1)
  for i = at, at + count - 1 do
    out = out .. string.char((bytes:byte(i) + 137) % 256)
  end
  return out .. bytes:sub(at + count)
end

t.describe('archive.safe_path', function()
  t.it('accepts an ordinary nested path', function()
    t.eq('plugin/lua/init.lua', archive.safe_path('plugin/lua/init.lua'))
  end)

  t.it('collapses redundant separators and dot components', function()
    t.eq('a/b.lua', archive.safe_path('./a//./b.lua'))
  end)

  t.it('rejects a parent-directory component', function()
    local ok, err = archive.safe_path('a/../../etc/passwd')
    t.falsy(ok)
    t.matches('escapes', err)
  end)

  t.it('rejects a leading parent-directory component', function()
    t.falsy(archive.safe_path('../evil'))
  end)

  t.it('rejects an absolute path', function()
    local ok, err = archive.safe_path('/etc/passwd')
    t.falsy(ok)
    t.matches('absolute', err)
  end)

  t.it('rejects a windows drive path', function()
    local ok, err = archive.safe_path('C:/windows/system32')
    t.falsy(ok)
    t.matches('drive', err)
  end)

  t.it('rejects a backslash separator', function()
    t.falsy(archive.safe_path('a\\..\\..\\evil'))
  end)

  t.it('rejects an embedded NUL', function()
    t.falsy(archive.safe_path('a\0b'))
  end)

  t.it('rejects a name that resolves to nothing', function()
    t.falsy(archive.safe_path('./'))
  end)
end)

t.describe('archive.extract zip', function()
  t.it('writes nested files, empty files and directories', function()
    local ok, err, dest = extract(zipf({
      { name = 'root/' },
      { name = 'root/init.lua', data = 'return {}' },
      { name = 'root/empty.txt', data = '' },
      { name = 'root/lua/deep/mod.lua', data = 'deep' },
      { name = 'root/noext', data = 'x' },
      { name = 'root/with space.txt', data = 'y' },
      { name = 'root/ünïcødé.txt', data = 'z' },
    }))
    t.truthy(ok, err)
    t.eq('return {}', read(dest .. '/root/init.lua'))
    t.eq('', read(dest .. '/root/empty.txt'))
    t.eq('deep', read(dest .. '/root/lua/deep/mod.lua'))
    t.eq('x', read(dest .. '/root/noext'))
    t.eq('y', read(dest .. '/root/with space.txt'))
    t.eq('z', read(dest .. '/root/ünïcødé.txt'))
  end)

  t.it('reads an archive that carries a trailing comment', function()
    local bytes = zipf({ { name = 'root/a.lua', data = 'hi' }, comment = 'built by CI PK\5\6 not really' })
    local ok, err, dest = extract(bytes)
    t.truthy(ok, err)
    t.eq('hi', read(dest .. '/root/a.lua'))
  end)

  t.it('carries the executable bit but drops setuid', function()
    local uv = vim.uv or vim.loop
    local ok, err, dest = extract(zipf({
      { name = 'root/run.sh', data = '#!/bin/sh\n', attr = (0x8000 + 0x800 + 0x1ed) * 65536 },
    }))
    t.truthy(ok, err)
    t.eq(0x1ed, uv.fs_stat(dest .. '/root/run.sh').mode % 0x1000)
  end)

  t.it('refuses an entry that walks out of the destination', function()
    local ok, err = extract(zipf({ { name = '../evil.lua', data = 'pwned' } }))
    t.falsy(ok)
    t.matches('escapes', err)
  end)

  t.it('refuses an absolute entry', function()
    local ok, err = extract(zipf({ { name = '/tmp/evil.lua', data = 'pwned' } }))
    t.falsy(ok)
    t.matches('absolute', err)
  end)

  t.it('refuses a symlink entry', function()
    local ok, err = extract(zipf({ { name = 'root/link', data = '/etc', attr = (0xa000 + 0x1ff) * 65536 } }))
    t.falsy(ok)
    t.matches('not a regular file', err)
  end)

  t.it('refuses an encrypted entry', function()
    local ok, err = extract(zipf({ { name = 'root/a.lua', data = 'x', flags = 1 } }))
    t.falsy(ok)
    t.matches('encrypted', err)
  end)

  t.it('fails loudly on a CRC mismatch', function()
    local ok, err = extract(zipf({ { name = 'root/a.lua', data = 'hello', crc = 12345 } }))
    t.falsy(ok)
    t.matches('CRC mismatch', err)
  end)

  t.it('fails on a truncated archive', function()
    local bytes = zipf({ { name = 'root/a.lua', data = 'hello' } })
    local ok, err = extract(bytes:sub(1, #bytes - 40))
    t.falsy(ok)
    t.matches('end%-of%-central%-directory', err)
  end)

  t.it('fails on random bytes with a zip magic', function()
    local ok = extract('PK\3\4' .. string.rep('\200', 400))
    t.falsy(ok)
  end)
  t.it('rejects an unsupported compression method', function()
    local ok, err = extract(zipf({ { name = 'root/a.lua', data = 'x', method = 12 } }))
    t.falsy(ok)
    t.matches('unsupported compression method 12', err)
  end)
end)

t.describe('archive.extract zip malformed directory', function()
  -- A central directory record that stops short of its own header used to
  -- raise instead of reporting.
  local function truncate_cd(keep)
    local bytes = zipf({ { name = 'root/a.lua', data = 'hi' } })
    local cd = bytes:find('PK\1\2', 1, true)
    local eocd = bytes:find('PK\5\6', cd, true)
    return bytes:sub(1, cd + keep - 1) .. bytes:sub(eocd)
  end

  t.it('reports a record cut off before its extra-field lengths', function()
    -- The directory offset is attacker-controlled, so point it at a record
    -- planted in the trailing comment: a zero name length lets the parse run
    -- cleanly up to the extra-field and comment lengths, which is where it
    -- used to raise.
    local comment = 'PK\1\2' .. string.rep('\0', 26)
    local bytes = zipf({ { name = 'root/a.lua', data = 'hi' }, comment = comment })
    local ok, err = extract(patch32(bytes, #bytes - #comment - 21 + 16, #bytes - #comment))
    t.falsy(ok)
    t.matches('truncated central directory', err)
  end)

  t.it('reports a record cut off just after its signature', function()
    local ok, err = extract(truncate_cd(4))
    t.falsy(ok)
    t.matches('truncated central directory', err)
  end)

  t.it('reports a file that another entry then treats as a directory', function()
    local ok, err = extract(zipf({
      { name = 'root/a', data = 'x' },
      { name = 'root/a/b', data = 'y' },
    }))
    t.falsy(ok)
    t.matches('Cannot create directory', err)
  end)

  t.it('reports a directory entry that collides with an earlier file', function()
    local ok, err = extract(zipf({ { name = 'root/a', data = 'x' }, { name = 'root/a/' } }))
    t.falsy(ok)
    t.matches('Cannot create directory', err)
  end)

  t.it('reports a local header that runs off the end of the archive', function()
    -- The trailing comment plants a local-header signature in the last four
    -- bytes; the entry is pointed at it.
    local bytes = zipf({ { name = 'root/a.lua', data = 'hi' }, comment = 'PK\3\4' })
    local cd = bytes:find('PK\1\2', 1, true)
    local ok, err = extract(patch32(bytes, cd + 42, #bytes - 4))
    t.falsy(ok)
    t.matches('truncated local header', err)
  end)
end)

t.describe('archive.extract zip deflate', function()
  local payload = sample()

  t.it('round-trips a dynamic-huffman entry byte for byte', function()
    local ok, err, dest = extract(zipf({ { name = 'root/a.lua', data = payload, deflate = true } }))
    t.truthy(ok, err)
    t.eq(payload, read(dest .. '/root/a.lua'))
  end)

  t.it('round-trips a fixed-huffman entry byte for byte', function()
    local bytes = zipf({ { name = 'root/a.lua', data = payload, deflate = true, mode = 'fixed' } })
    local ok, err, dest = extract(bytes)
    t.truthy(ok, err)
    t.eq(payload, read(dest .. '/root/a.lua'))
  end)

  t.it('round-trips deflated entries alongside stored ones', function()
    local overlapping = string.rep('ab', 3000)
    local ok, err, dest = extract(zipf({
      { name = 'root/' },
      { name = 'root/stored.lua', data = 'return {}' },
      { name = 'root/packed.lua', data = payload, deflate = true },
      { name = 'root/nested/overlap.txt', data = overlapping, deflate = true },
    }))
    t.truthy(ok, err)
    t.eq('return {}', read(dest .. '/root/stored.lua'))
    t.eq(payload, read(dest .. '/root/packed.lua'))
    t.eq(overlapping, read(dest .. '/root/nested/overlap.txt'))
  end)

  t.it('round-trips a deflated entry that decompresses to nothing', function()
    local ok, err, dest = extract(zipf({ { name = 'root/empty.lua', data = '', deflate = true } }))
    t.truthy(ok, err)
    t.eq('', read(dest .. '/root/empty.lua'))
  end)

  t.it('round-trips a large deflated entry', function()
    local big = string.rep(noise(1024, 11), 400)
    local bytes = zipf({ { name = 'root/big.bin', data = big, deflate = true } })
    t.truthy(#bytes < #big / 4, 'the fixture is not actually compressed')
    local ok, err, dest = extract(bytes)
    t.truthy(ok, err)
    local out = read(dest .. '/root/big.bin')
    t.eq(#big, #out)
    t.truthy(out == big, 'the large entry did not round-trip exactly')
  end)

  t.it('fails loudly when a deflated entry carries the wrong CRC', function()
    local ok, err, dest = extract(zipf({ { name = 'root/a.lua', data = payload, deflate = true, crc = 12345 } }))
    t.falsy(ok)
    t.matches('checksum', err)
    t.falsy(read(dest .. '/root/a.lua'), 'nothing should be written')
  end)

  t.it('fails loudly when a deflated entry disagrees with the stated size', function()
    local bytes = zipf({ { name = 'root/a.lua', data = payload, deflate = true, usize = #payload + 1 } })
    local ok, err, dest = extract(bytes)
    t.falsy(ok)
    t.matches('does not match the directory', err)
    t.falsy(read(dest .. '/root/a.lua'), 'nothing should be written')
  end)

  t.it('fails loudly on a truncated deflate stream', function()
    local ok, err, dest = extract(zipf({
      {
        name = 'root/a.lua',
        data = payload,
        deflate = true,
        mangle = function(s)
          return s:sub(1, math.floor(#s / 2))
        end,
      },
    }))
    t.falsy(ok)
    t.matches('root/a%.lua', err)
    t.falsy(read(dest .. '/root/a.lua'), 'a partial inflate must not reach disk')
  end)

  t.it('fails loudly on a deflate stream corrupted in the middle', function()
    local ok, err, dest = extract(zipf({
      {
        name = 'root/a.lua',
        data = payload,
        deflate = true,
        mangle = function(s)
          return flip(s, math.floor(#s / 2), 4)
        end,
      },
    }))
    t.falsy(ok)
    t.matches('root/a%.lua', err)
    t.falsy(read(dest .. '/root/a.lua'), 'a corrupt inflate must not reach disk')
  end)

  t.it('refuses a deflated entry that walks out of the destination', function()
    local ok, err, dest = extract(zipf({ { name = '../evil.lua', data = payload, deflate = true } }))
    t.falsy(ok)
    t.matches('escapes', err)
    t.falsy(read(dest .. '/../evil.lua'), 'nothing should be written')
  end)

  t.it('refuses a deflated symlink entry', function()
    local ok, err = extract(zipf({
      { name = 'root/link', data = string.rep('/etc/passwd', 40), deflate = true, attr = (0xa000 + 0x1ff) * 65536 },
    }))
    t.falsy(ok)
    t.matches('not a regular file', err)
  end)
end)

t.describe('archive.extract tar.gz', function()
  t.it('writes files and directories', function()
    local ok, err, dest = extract(
      gz(tarf({
        { name = 'root/', kind = '5' },
        { name = 'root/init.lua', data = 'return {}' },
        { name = 'root/empty.txt', data = '' },
        { name = 'root/lua/deep/mod.lua', data = 'deep' },
      })),
      'tar.gz'
    )
    t.truthy(ok, err)
    t.eq('return {}', read(dest .. '/root/init.lua'))
    t.eq('', read(dest .. '/root/empty.txt'))
    t.eq('deep', read(dest .. '/root/lua/deep/mod.lua'))
  end)

  t.it('skips a pax global header', function()
    local ok, err, dest = extract(
      gz(tarf({
        { name = 'pax_global_header', kind = 'g', data = '20 comment=deadbeef\n' },
        { name = 'root/init.lua', data = 'ok' },
      })),
      'tar.gz'
    )
    t.truthy(ok, err)
    t.eq('ok', read(dest .. '/root/init.lua'))
  end)

  t.it('honours a pax path override', function()
    local path = 'root/a-very-long-name.lua'
    -- A pax record's length field counts itself, so it has to reach a fixed point.
    local len = #path + 7
    while #tostring(len) + #path + 7 ~= len do
      len = #tostring(len) + #path + 7
    end
    local record = string.format('%d path=%s\n', len, path)
    local ok, err, dest = extract(
      gz(tarf({
        { name = 'short', kind = 'x', data = record },
        { name = 'short', data = 'long' },
      })),
      'tar.gz'
    )
    t.truthy(ok, err)
    t.eq('long', read(dest .. '/' .. path))
  end)

  t.it('refuses an entry that walks out of the destination', function()
    local ok, err = extract(gz(tarf({ { name = '../evil.lua', data = 'pwned' } })), 'tar.gz')
    t.falsy(ok)
    t.matches('escapes', err)
  end)

  t.it('refuses a symlink entry', function()
    local ok, err = extract(gz(tarf({ { name = 'root/link', kind = '2' } })), 'tar.gz')
    t.falsy(ok)
    t.matches('unsupported tar entry type', err)
  end)

  t.it('fails loudly on a corrupt gzip payload', function()
    local bytes = gz(tarf({ { name = 'root/a.lua', data = 'hello' } }))
    local flipped = bytes:sub(1, 40) .. string.char((bytes:byte(41) + 1) % 256) .. bytes:sub(42)
    local ok, err = extract(flipped, 'tar.gz')
    t.falsy(ok)
    t.matches('gunzip failed', err)
  end)

  t.it('reports a tar directory that collides with an earlier file', function()
    local bytes = gz(tarf({ { name = 'root/a', data = 'x' }, { name = 'root/a/', kind = '5' } }))
    local ok, err = extract(bytes, 'tar.gz')
    t.falsy(ok)
    t.matches('Cannot create directory', err)
  end)

  t.it('fails on a corrupt tar header', function()
    local tar = tarf({ { name = 'root/a.lua', data = 'hello' } })
    local ok, err = extract(gz('X' .. tar:sub(2)), 'tar.gz')
    t.falsy(ok)
    t.matches('checksum mismatch', err)
  end)
end)

-- The stored-block fixtures above never reach the huffman decoder, so the
-- compressed cases are covered separately.
t.describe('archive.extract tar.gz deflate', function()
  local payload = sample()

  t.it('inflates a dynamic-huffman member', function()
    local bytes = gz(tarf({ { name = 'root/a.lua', data = payload } }), 'dynamic')
    local ok, err, dest = extract(bytes, 'tar.gz')
    t.truthy(ok, err)
    t.eq(payload, read(dest .. '/root/a.lua'))
  end)

  t.it('inflates a fixed-huffman member', function()
    local bytes = gz(tarf({ { name = 'root/a.lua', data = payload } }), 'fixed')
    local ok, err, dest = extract(bytes, 'tar.gz')
    t.truthy(ok, err)
    t.eq(payload, read(dest .. '/root/a.lua'))
  end)

  t.it('inflates an empty member', function()
    local bytes = gz(tarf({ { name = 'root/empty.lua', data = '' } }), 'dynamic')
    local ok, err, dest = extract(bytes, 'tar.gz')
    t.truthy(ok, err)
    t.eq('', read(dest .. '/root/empty.lua'))
  end)

  t.it('inflates a large member', function()
    local big = string.rep(noise(1024, 29), 300)
    local bytes = gz(tarf({ { name = 'root/big.bin', data = big } }), 'dynamic')
    t.truthy(#bytes < #big / 4, 'the fixture is not actually compressed')
    local ok, err, dest = extract(bytes, 'tar.gz')
    t.truthy(ok, err)
    local out = read(dest .. '/root/big.bin')
    t.eq(#big, #out)
    t.truthy(out == big, 'the large member did not round-trip exactly')
  end)

  t.it('fails loudly when the gzip checksum disagrees', function()
    local bytes = gz(tarf({ { name = 'root/a.lua', data = payload } }), 'dynamic')
    local ok, err, dest = extract(patch32(bytes, #bytes - 7, 12345), 'tar.gz')
    t.falsy(ok)
    t.matches('gunzip failed', err)
    t.falsy(read(dest .. '/root/a.lua'), 'nothing should be written')
  end)

  t.it('fails loudly on a truncated compressed member', function()
    local bytes = gz(tarf({ { name = 'root/a.lua', data = payload } }), 'dynamic')
    local ok, err, dest = extract(bytes:sub(1, 10 + math.floor((#bytes - 18) / 2)), 'tar.gz')
    t.falsy(ok)
    t.matches('gunzip failed', err)
    t.falsy(read(dest .. '/root/a.lua'), 'a partial inflate must not reach disk')
  end)

  t.it('fails loudly on a member corrupted in the middle', function()
    local bytes = gz(tarf({ { name = 'root/a.lua', data = payload } }), 'dynamic')
    local ok, err, dest = extract(flip(bytes, math.floor(#bytes / 2), 4), 'tar.gz')
    t.falsy(ok)
    t.matches('gunzip failed', err)
    t.falsy(read(dest .. '/root/a.lua'), 'a corrupt inflate must not reach disk')
  end)

  t.it('refuses a hostile path inside a compressed member', function()
    local bytes = gz(tarf({ { name = '../evil.lua', data = payload } }), 'dynamic')
    local ok, err, dest = extract(bytes, 'tar.gz')
    t.falsy(ok)
    t.matches('escapes', err)
    t.falsy(read(dest .. '/../evil.lua'), 'nothing should be written')
  end)

  t.it('stops a pax record that claims a zero length', function()
    local bytes = gz(
      tarf({
        { name = 'root/pax', kind = 'x', data = '0 path=root/a.lua\n' },
        { name = 'root/a.lua', data = 'ok' },
      }),
      'dynamic'
    )
    local ok, err, dest = extract(bytes, 'tar.gz')
    t.truthy(ok, err)
    t.eq('ok', read(dest .. '/root/a.lua'))
  end)
end)