|
| 1 | +import Foundation |
| 2 | +import zlib |
| 3 | + |
| 4 | +struct SaveError: LocalizedError { |
| 5 | + let message: String |
| 6 | + init(_ message: String) { self.message = message } |
| 7 | + var errorDescription: String? { message } |
| 8 | +} |
| 9 | +struct SaveEntry { |
| 10 | + let path: String |
| 11 | + let data: Data |
| 12 | + var modified: Date = Date() |
| 13 | +} |
| 14 | +extension Data { |
| 15 | + func u16(_ at: Int) -> Int { Int(self[at]) | Int(self[at + 1]) << 8 } |
| 16 | + func u32(_ at: Int) -> UInt32 { (0..<4).reduce(0) { $0 | UInt32(self[at + $1]) << ($1 * 8) } } |
| 17 | + mutating func put16(_ at: Int, _ n: Int) { self[at] = UInt8(truncatingIfNeeded: n); self[at + 1] = UInt8(truncatingIfNeeded: n >> 8) } |
| 18 | + mutating func put32(_ at: Int, _ n: UInt32) { for i in 0..<4 { self[at + i] = UInt8(truncatingIfNeeded: n >> (i * 8)) } } |
| 19 | +} |
| 20 | + |
| 21 | +enum SaveArchive { |
| 22 | + static let maxTotal = 256 << 20, maxEntry = 64 << 20, maxEntries = 4096 |
| 23 | + static func validPath(_ path: String) -> Bool { |
| 24 | + let parts = path.split(separator: "/", omittingEmptySubsequences: false) |
| 25 | + return !path.isEmpty && path.utf8.count <= 1024 && !path.contains(":") && !path.contains("\\") && |
| 26 | + !path.unicodeScalars.contains { $0.value < 32 || $0.value == 127 } && |
| 27 | + parts.allSatisfy { !$0.isEmpty && $0 != "." && $0 != ".." && $0.utf8.count <= 255 } |
| 28 | + } |
| 29 | + static func validatePaths(_ entries: [SaveEntry]) throws { |
| 30 | + var seen = Set<String>(), bytes = 0 |
| 31 | + guard entries.count <= maxEntries else { throw SaveError("Too many archive entries.") } |
| 32 | + for e in entries { |
| 33 | + bytes += e.data.count |
| 34 | + let path = e.path.precomposedStringWithCanonicalMapping.lowercased() |
| 35 | + guard validPath(e.path), seen.insert(path).inserted, e.data.count <= maxEntry, bytes <= maxTotal else { |
| 36 | + throw SaveError("Unsafe, duplicate or oversized archive entry.") |
| 37 | + } |
| 38 | + } |
| 39 | + for name in seen { |
| 40 | + var parts = name.split(separator: "/"); parts.removeLast() |
| 41 | + while !parts.isEmpty { |
| 42 | + guard !seen.contains(parts.joined(separator: "/")) else { throw SaveError("Archive file and directory paths overlap.") } |
| 43 | + parts.removeLast() |
| 44 | + } |
| 45 | + } |
| 46 | + } |
| 47 | + static func crc(_ data: Data) -> UInt32 { |
| 48 | + data.withUnsafeBytes { UInt32(crc32(0, $0.bindMemory(to: Bytef.self).baseAddress, uInt(data.count))) } |
| 49 | + } |
| 50 | + static func deflate(_ input: Data, inflate: Bool, limit: Int = maxTotal) throws -> Data { |
| 51 | + var stream = z_stream() |
| 52 | + let status = inflate ? inflateInit2_(&stream, -MAX_WBITS, ZLIB_VERSION, Int32(MemoryLayout<z_stream>.size)) : |
| 53 | + deflateInit2_(&stream, Z_DEFAULT_COMPRESSION, Z_DEFLATED, -MAX_WBITS, 8, Z_DEFAULT_STRATEGY, ZLIB_VERSION, Int32(MemoryLayout<z_stream>.size)) |
| 54 | + guard status == Z_OK else { throw SaveError("Could not open compressed data.") } |
| 55 | + defer { if inflate { inflateEnd(&stream) } else { deflateEnd(&stream) } } |
| 56 | + return try input.withUnsafeBytes { source in |
| 57 | + stream.next_in = UnsafeMutablePointer(mutating: source.bindMemory(to: Bytef.self).baseAddress) |
| 58 | + stream.avail_in = uInt(input.count) |
| 59 | + var result = Data(), chunk = [UInt8](repeating: 0, count: 32768) |
| 60 | + while true { |
| 61 | + let rc = chunk.withUnsafeMutableBytes { output -> Int32 in |
| 62 | + stream.next_out = output.bindMemory(to: Bytef.self).baseAddress |
| 63 | + stream.avail_out = 32768 |
| 64 | + return inflate ? zlib.inflate(&stream, Z_NO_FLUSH) : zlib.deflate(&stream, Z_FINISH) |
| 65 | + } |
| 66 | + let count = chunk.count - Int(stream.avail_out) |
| 67 | + guard result.count + count <= limit else { throw SaveError("Compressed data exceeds the size limit.") } |
| 68 | + result.append(contentsOf: chunk.prefix(count)) |
| 69 | + if rc == Z_STREAM_END { |
| 70 | + guard stream.avail_in == 0 else { throw SaveError("Trailing compressed data.") } |
| 71 | + return result |
| 72 | + } |
| 73 | + guard rc == Z_OK, count > 0 else { throw SaveError("Damaged compressed data.") } |
| 74 | + } |
| 75 | + } |
| 76 | + } |
| 77 | + static func encode(_ entries: [SaveEntry]) throws -> Data { |
| 78 | + try validatePaths(entries) |
| 79 | + var local = Data(), central = Data() |
| 80 | + for entry in entries { |
| 81 | + let name = Data(entry.path.utf8), checksum = crc(entry.data), offset = local.count |
| 82 | + let c = Calendar(identifier: .gregorian).dateComponents([.year,.month,.day,.hour,.minute,.second], from: entry.modified) |
| 83 | + let date = (max(1980, min(2107, c.year!)) - 1980) << 9 | c.month! << 5 | c.day! |
| 84 | + let time = c.hour! << 11 | c.minute! << 5 | c.second! / 2 |
| 85 | + var h = Data(repeating: 0, count: 30) |
| 86 | + h.put32(0, 0x04034b50); h.put16(4,20); h.put16(6,0x800); h.put16(10,time); h.put16(12,date) |
| 87 | + h.put32(14,checksum); h.put32(18,UInt32(entry.data.count)); h.put32(22,UInt32(entry.data.count)); h.put16(26,name.count) |
| 88 | + local.append(h); local.append(name); local.append(entry.data) |
| 89 | + var cd = Data(repeating: 0, count: 46) |
| 90 | + cd.put32(0,0x02014b50); cd.put16(4,0x31e); cd.put16(6,20); cd.put16(8,0x800); cd.put16(12,time); cd.put16(14,date) |
| 91 | + cd.put32(16,checksum); cd.put32(20,UInt32(entry.data.count)); cd.put32(24,UInt32(entry.data.count)); cd.put16(28,name.count) |
| 92 | + cd.put32(38,0o100644 << 16); cd.put32(42,UInt32(offset)); central.append(cd); central.append(name) |
| 93 | + } |
| 94 | + var end = Data(repeating: 0, count: 22) |
| 95 | + end.put32(0,0x06054b50); end.put16(8,entries.count); end.put16(10,entries.count) |
| 96 | + end.put32(12,UInt32(central.count)); end.put32(16,UInt32(local.count)) |
| 97 | + local.append(central); local.append(end); return local |
| 98 | + } |
| 99 | + static func decode(_ data: Data) throws -> [SaveEntry] { |
| 100 | + func bad() -> SaveError { SaveError("Invalid or unsafe ZIP archive.") } |
| 101 | + guard data.count >= 22, data.count <= maxTotal + (8 << 20) else { throw bad() } |
| 102 | + guard let end = stride(from: data.count - 22, through: max(0,data.count - 65557), by: -1).first(where: { |
| 103 | + data.u32($0) == 0x06054b50 && $0 + 22 + data.u16($0 + 20) == data.count |
| 104 | + }) else { throw bad() } |
| 105 | + let count = data.u16(end + 10), start = Int(data.u32(end + 16)) |
| 106 | + guard data.u32(end + 4) == 0, count == data.u16(end + 8), count <= maxEntries, |
| 107 | + start + Int(data.u32(end + 12)) == end else { throw bad() } |
| 108 | + var cursor = start, total = 0, entries = [SaveEntry](), ranges = [Range<Int>]() |
| 109 | + for _ in 0..<count { |
| 110 | + guard cursor >= 0, cursor + 46 <= end, data.u32(cursor) == 0x02014b50 else { throw bad() } |
| 111 | + let flags = data.u16(cursor + 8), method = data.u16(cursor + 10), packed = Int(data.u32(cursor + 20)), size = Int(data.u32(cursor + 24)) |
| 112 | + let nl = data.u16(cursor + 28), extra = data.u16(cursor + 30), comment = data.u16(cursor + 32), local = Int(data.u32(cursor + 42)) |
| 113 | + let mode = (data.u32(cursor + 38) >> 16) & 0xf000 |
| 114 | + guard flags & ~0x808 == 0, method == 0 || method == 8, data.u16(cursor + 34) == 0, |
| 115 | + mode == 0 || mode == 0x8000 || mode == 0x4000, |
| 116 | + cursor + 46 + nl + extra + comment <= end, |
| 117 | + let name = String(data: data[(cursor + 46)..<(cursor + 46 + nl)], encoding: .utf8), |
| 118 | + local + 30 <= start, data.u32(local) == 0x04034b50, |
| 119 | + data.u16(local + 6) == flags, data.u16(local + 8) == method, data.u16(local + 26) == nl else { throw bad() } |
| 120 | + let payload = local + 30 + nl + data.u16(local + 28) |
| 121 | + guard payload + packed <= start, Data(data[(local + 30)..<(local + 30 + nl)]) == Data(name.utf8), |
| 122 | + size <= maxEntry, total + size <= maxTotal else { throw bad() } |
| 123 | + let range = local..<(payload + packed) |
| 124 | + guard !ranges.contains(where: { $0.overlaps(range) }) else { throw bad() }; ranges.append(range) |
| 125 | + if flags & 8 == 0 { |
| 126 | + guard data.u32(local + 14) == data.u32(cursor + 16), Int(data.u32(local + 18)) == packed, |
| 127 | + Int(data.u32(local + 22)) == size else { throw bad() } |
| 128 | + } |
| 129 | + total += size |
| 130 | + let raw = Data(data[payload..<(payload + packed)]) |
| 131 | + let content = method == 0 ? raw : try deflate(raw, inflate: true, limit: size) |
| 132 | + guard content.count == size, crc(content) == data.u32(cursor + 16) else { throw bad() } |
| 133 | + if name.hasSuffix("/") { |
| 134 | + guard size == 0, validPath(String(name.dropLast())) else { throw bad() } |
| 135 | + } else { |
| 136 | + guard mode != 0x4000 else { throw bad() } |
| 137 | + let day = data.u16(cursor + 14), time = data.u16(cursor + 12) |
| 138 | + let date = Calendar(identifier: .gregorian).date(from: DateComponents(year: 1980 + (day >> 9), month: (day >> 5) & 15, day: day & 31, hour: time >> 11, minute: (time >> 5) & 63, second: (time & 31) * 2)) ?? Date(timeIntervalSince1970: 0) |
| 139 | + entries.append(SaveEntry(path: name, data: content, modified: date)) |
| 140 | + } |
| 141 | + cursor += 46 + nl + extra + comment |
| 142 | + } |
| 143 | + guard cursor == end else { throw bad() }; try validatePaths(entries); return entries |
| 144 | + } |
| 145 | +} |
0 commit comments