kotej/Sources/KotejEngine/ZipReader.swift

import Foundation
import Compression

/// Reads ZIP containers (which includes JAR, EAR and WAR) without unpacking them.
///
/// The central directory already records every entry's CRC32 and sizes, so two
/// archives can be compared entry by entry *without decompressing anything*.
/// Bytes are only inflated when a specific entry is actually opened — that's what
/// keeps comparing large EARs fast.
public struct ZipReader {
    /// One entry in the archive's central directory.
    public struct Entry {
        public let path: String
        public let crc32: UInt32
        public let compressedSize: UInt64
        public let uncompressedSize: UInt64
        public let compressionMethod: UInt16
        public let localHeaderOffset: UInt64
        public let isDirectory: Bool
    }

    public enum ZipError: Error, LocalizedError {
        case notAZip
        case unsupported(String)
        case corrupt(String)

        public var errorDescription: String? {
            switch self {
            case .notAZip: return "Not a ZIP archive (no end-of-central-directory record)."
            case .unsupported(let what): return "Unsupported ZIP feature: \(what)."
            case .corrupt(let why): return "Corrupt ZIP: \(why)."
            }
        }
    }

    public let url: URL
    public let entries: [Entry]

    /// File extensions handled as archives, i.e. browsable like directories.
    public static let archiveExtensions: Set<String> = ["zip", "jar", "ear", "war", "aar", "ipa"]

    public static func isArchive(_ url: URL) -> Bool {
        if archiveExtensions.contains(url.pathExtension.lowercased()) { return true }
        // A ZIP renamed to anything else is still a ZIP, and deployment artefacts
        // do this constantly, so fall back to the signature.
        return looksLikeArchive(fileAt: url)
    }

    /// Extensions that are definitely not containers, used to skip the signature
    /// check on the thousands of small entries inside a real EAR.
    static let neverArchiveExtensions: Set<String> = [
        "class", "java", "swift", "kt", "go", "c", "h", "m", "js", "ts", "css", "scss",
        "html", "htm", "xml", "xsd", "xsl", "xslt", "wsdl", "json", "yml", "yaml",
        "properties", "conf", "cfg", "ini", "toml", "md", "txt", "log", "csv", "sql",
        "sh", "bat", "mf", "sf", "rsa", "dsa", "png", "jpg", "jpeg", "gif", "svg",
        "ico", "pdf", "ttf", "otf", "woff", "woff2", "so", "dylib", "plist",
    ]

    /// The local-file-header signature, plus the empty and spanned variants.
    private static let signatures: [[UInt8]] = [
        [0x50, 0x4B, 0x03, 0x04],
        [0x50, 0x4B, 0x05, 0x06],
        [0x50, 0x4B, 0x07, 0x08],
    ]

    /// Cheap check: read only the first four bytes.
    public static func looksLikeArchive(fileAt url: URL) -> Bool {
        guard let handle = try? FileHandle(forReadingFrom: url) else { return false }
        defer { try? handle.close() }
        guard let head = try? handle.read(upToCount: 4) else { return false }
        return looksLikeArchive(head)
    }

    public static func looksLikeArchive(_ head: Data?) -> Bool {
        guard let head, head.count >= 4 else { return false }
        let bytes = Array(head.prefix(4))
        return signatures.contains(bytes)
    }

    /// Whether an entry inside an archive is worth opening as one. Known archive
    /// extensions always are; anything with a clearly non-container extension
    /// never is; the rest are only inflated when small enough that guessing wrong
    /// is cheap, since checking the signature means decompressing it.
    static func mayBeNestedArchive(name: String, uncompressedSize: UInt64,
                                   inflateLimit: UInt64 = 64 * 1024 * 1024) -> Bool {
        let ext = (name as NSString).pathExtension.lowercased()
        if archiveExtensions.contains(ext) { return true }
        if !ext.isEmpty && neverArchiveExtensions.contains(ext) { return false }
        return uncompressedSize >= 22 && uncompressedSize <= inflateLimit
    }

    public init(url: URL) throws {
        self.url = url
        let data = try Data(contentsOf: url, options: .mappedIfSafe)
        self.entries = try Self.readCentralDirectory(data)
    }

    // MARK: Central directory

    private static let endOfCentralDirectorySignature: UInt32 = 0x06054b50
    private static let zip64LocatorSignature: UInt32 = 0x07064b50
    private static let zip64EndSignature: UInt32 = 0x06064b50
    private static let centralFileHeaderSignature: UInt32 = 0x02014b50

    private static func readCentralDirectory(_ data: Data) throws -> [Entry] {
        guard let eocdOffset = findEndOfCentralDirectory(data) else { throw ZipError.notAZip }

        var entryCount = Int(data.u16(eocdOffset + 10))
        var directoryOffset = Int(data.u32(eocdOffset + 16))
        let directorySize = Int(data.u32(eocdOffset + 12))

        // ZIP64: the 32-bit fields saturate, and the real values live in the
        // ZIP64 record that the locator points at.
        if directoryOffset == 0xFFFF_FFFF || entryCount == 0xFFFF || directorySize == 0xFFFF_FFFF {
            guard let locator = findZip64Locator(data, before: eocdOffset) else {
                throw ZipError.unsupported("ZIP64 without locator")
            }
            let zip64End = Int(data.u64(locator + 8))
            guard zip64End + 56 <= data.count, data.u32(zip64End) == zip64EndSignature else {
                throw ZipError.corrupt("bad ZIP64 end record")
            }
            entryCount = Int(data.u64(zip64End + 32))
            directoryOffset = Int(data.u64(zip64End + 48))
        }

        var entries: [Entry] = []
        entries.reserveCapacity(entryCount)
        var offset = directoryOffset

        for _ in 0..<entryCount {
            guard offset + 46 <= data.count, data.u32(offset) == centralFileHeaderSignature else {
                throw ZipError.corrupt("bad central directory header")
            }
            let method = data.u16(offset + 10)
            let crc = data.u32(offset + 16)
            var compressed = UInt64(data.u32(offset + 20))
            var uncompressed = UInt64(data.u32(offset + 24))
            let nameLength = Int(data.u16(offset + 28))
            let extraLength = Int(data.u16(offset + 30))
            let commentLength = Int(data.u16(offset + 32))
            var localOffset = UInt64(data.u32(offset + 42))

            let nameStart = offset + 46
            guard nameStart + nameLength <= data.count else { throw ZipError.corrupt("truncated name") }
            let name = String(decoding: data[nameStart..<nameStart + nameLength], as: UTF8.self)

            // Oversized values are carried in the ZIP64 extra field.
            if uncompressed == 0xFFFF_FFFF || compressed == 0xFFFF_FFFF || localOffset == 0xFFFF_FFFF {
                let extraStart = nameStart + nameLength
                readZip64Extra(data, start: extraStart, length: extraLength,
                               uncompressed: &uncompressed, compressed: &compressed, localOffset: &localOffset)
            }

            entries.append(Entry(path: normalise(name),
                                 crc32: crc,
                                 compressedSize: compressed,
                                 uncompressedSize: uncompressed,
                                 compressionMethod: method,
                                 localHeaderOffset: localOffset,
                                 isDirectory: name.hasSuffix("/")))
            offset = nameStart + nameLength + extraLength + commentLength
        }
        return entries
    }

    /// Entries may use either separator and some tools prefix "./".
    private static func normalise(_ name: String) -> String {
        var path = name.replacingOccurrences(of: "\\", with: "/")
        while path.hasPrefix("./") { path.removeFirst(2) }
        while path.hasSuffix("/") { path.removeLast() }
        return path
    }

    private static func findEndOfCentralDirectory(_ data: Data) -> Int? {
        // The record is at the end but may be followed by a comment (max 64 KiB).
        let minimumSize = 22
        guard data.count >= minimumSize else { return nil }
        let searchLimit = max(0, data.count - minimumSize - 0xFFFF)
        var offset = data.count - minimumSize
        while offset >= searchLimit {
            if data.u32(offset) == endOfCentralDirectorySignature { return offset }
            offset -= 1
        }
        return nil
    }

    private static func findZip64Locator(_ data: Data, before eocd: Int) -> Int? {
        let locator = eocd - 20
        guard locator >= 0, data.u32(locator) == zip64LocatorSignature else { return nil }
        return locator
    }

    private static func readZip64Extra(_ data: Data, start: Int, length: Int,
                                       uncompressed: inout UInt64, compressed: inout UInt64,
                                       localOffset: inout UInt64) {
        var cursor = start
        let end = min(start + length, data.count)
        while cursor + 4 <= end {
            let headerID = data.u16(cursor)
            let size = Int(data.u16(cursor + 2))
            let body = cursor + 4
            if headerID == 0x0001 {
                var field = body
                if uncompressed == 0xFFFF_FFFF, field + 8 <= end { uncompressed = data.u64(field); field += 8 }
                if compressed == 0xFFFF_FFFF, field + 8 <= end { compressed = data.u64(field); field += 8 }
                if localOffset == 0xFFFF_FFFF, field + 8 <= end { localOffset = data.u64(field) }
                return
            }
            cursor = body + size
        }
    }

    // MARK: Reading entry bytes

    /// Inflates one entry. Only called when the bytes are genuinely needed —
    /// comparisons normally stop at the CRC32 from the central directory.
    public func data(for entry: Entry) throws -> Data {
        let archive = try Data(contentsOf: url, options: .mappedIfSafe)
        let headerOffset = Int(entry.localHeaderOffset)
        guard headerOffset + 30 <= archive.count, archive.u32(headerOffset) == 0x04034b50 else {
            throw ZipError.corrupt("bad local header for \(entry.path)")
        }
        let nameLength = Int(archive.u16(headerOffset + 26))
        let extraLength = Int(archive.u16(headerOffset + 28))
        let start = headerOffset + 30 + nameLength + extraLength
        let end = start + Int(entry.compressedSize)
        guard end <= archive.count else { throw ZipError.corrupt("truncated data for \(entry.path)") }
        let payload = archive.subdata(in: start..<end)

        switch entry.compressionMethod {
        case 0:
            return payload
        case 8:
            return try inflate(payload, expectedSize: Int(entry.uncompressedSize))
        default:
            throw ZipError.unsupported("compression method \(entry.compressionMethod)")
        }
    }

    /// Raw DEFLATE via the system Compression framework (no third-party zlib).
    private func inflate(_ payload: Data, expectedSize: Int) throws -> Data {
        guard expectedSize > 0 else { return Data() }
        var output = Data(count: expectedSize)

        let written: Int = try output.withUnsafeMutableBytes { destination in
            try payload.withUnsafeBytes { source in
                guard let destinationBase = destination.bindMemory(to: UInt8.self).baseAddress,
                      let sourceBase = source.bindMemory(to: UInt8.self).baseAddress else {
                    throw ZipError.corrupt("empty buffer")
                }
                return compression_decode_buffer(destinationBase, expectedSize,
                                                 sourceBase, payload.count,
                                                 nil, COMPRESSION_ZLIB)
            }
        }
        guard written == expectedSize else {
            throw ZipError.corrupt("inflate produced \(written) of \(expectedSize) bytes")
        }
        return output
    }
}

// MARK: - Little-endian reads

private extension Data {
    func u16(_ offset: Int) -> UInt16 {
        guard offset + 2 <= count else { return 0 }
        return UInt16(self[startIndex + offset]) | UInt16(self[startIndex + offset + 1]) << 8
    }

    func u32(_ offset: Int) -> UInt32 {
        guard offset + 4 <= count else { return 0 }
        var value: UInt32 = 0
        for i in (0..<4).reversed() { value = value << 8 | UInt32(self[startIndex + offset + i]) }
        return value
    }

    func u64(_ offset: Int) -> UInt64 {
        guard offset + 8 <= count else { return 0 }
        var value: UInt64 = 0
        for i in (0..<8).reversed() { value = value << 8 | UInt64(self[startIndex + offset + i]) }
        return value
    }
}