On-the-fly local file headers parsing (#12)

Author: mickael-menu

Sources/ZIPFoundation/Archive+Progress.swift

@@ -16,12 +16,12 @@ extension Archive {
     ///
     /// - Parameter entry: The entry that will be removed.
     /// - Returns: The number of the work units.
-    public func totalUnitCountForRemoving(_ entry: Entry) -> Int64 {
-        return Int64(self.offsetToStartOfCentralDirectory - entry.localSize)
+    func totalUnitCountForRemoving(_ entry: Entry, localFileHeader: LocalFileHeader) throws -> Int64 {
+        return try Int64(self.offsetToStartOfCentralDirectory - entry.localSize(with: localFileHeader))
     }
 
-    func makeProgressForRemoving(_ entry: Entry) -> Progress {
-        return Progress(totalUnitCount: self.totalUnitCountForRemoving(entry))
+    func makeProgressForRemoving(_ entry: Entry, localFileHeader: LocalFileHeader) throws -> Progress {
+        return Progress(totalUnitCount: try totalUnitCountForRemoving(entry, localFileHeader: localFileHeader))
     }
 
     /// The number of the work units that have to be performed when

Sources/ZIPFoundation/Archive+Reading.swift

@@ -91,12 +91,12 @@ extension Archive {
         guard bufferSize > 0 else {
             throw ArchiveError.invalidBufferSize
         }
-        let localFileHeader = entry.localFileHeader
-        guard entry.dataOffset <= .max else { throw ArchiveError.invalidLocalHeaderDataOffset }
+        let localFileHeader = try await localFileHeader(for: entry)
+        let dataOffset = try entry.dataOffset(with: localFileHeader)
 
         let transaction = try await dataSource.openRead()
         var checksum = CRC32(0)
-        try await transaction.seek(to: entry.dataOffset)
+        try await transaction.seek(to: dataOffset)
         progress?.totalUnitCount = self.totalUnitCountForReading(entry)
         switch entry.type {
         case .file:
@@ -128,7 +128,6 @@ extension Archive {
             try await consumer(Data())
             progress?.completedUnitCount = self.totalUnitCountForReading(entry)
         case .symlink:
-            let localFileHeader = entry.localFileHeader
             let size = Int(localFileHeader.compressedSize)
             let data = try await transaction.read(length: size)
             checksum = data.crc32(checksum: 0)
@@ -162,21 +161,18 @@ extension Archive {
         guard range.lowerBound >= 0, range.upperBound <= entry.uncompressedSize else {
             throw ArchiveError.rangeOutOfBounds
         }
-        let localFileHeader = entry.localFileHeader
-        guard entry.dataOffset <= .max else {
-            throw ArchiveError.invalidLocalHeaderDataOffset
-        }
+        let localFileHeader = try await localFileHeader(for: entry)
 
         guard let compressionMethod = CompressionMethod(rawValue: localFileHeader.compressionMethod) else {
             throw ArchiveError.invalidCompressionMethod
         }
 
         switch compressionMethod {
         case .none:
-            try await extractStoredRange(range, of: entry, bufferSize: bufferSize, consumer: consumer)
+            try await extractStoredRange(range, of: entry, localFileHeader: localFileHeader, bufferSize: bufferSize, consumer: consumer)
 
         case .deflate:
-            try await extractCompressedRange(range, of: entry, bufferSize: bufferSize, consumer: consumer)
+            try await extractCompressedRange(range, of: entry, localFileHeader: localFileHeader, bufferSize: bufferSize, consumer: consumer)
         }
     }
 
@@ -185,11 +181,13 @@ extension Archive {
     private func extractStoredRange(
         _ range: Range<UInt64>,
         of entry: Entry,
+        localFileHeader: LocalFileHeader,
         bufferSize: Int,
         consumer: Consumer
     ) async throws {
+        let offset = try entry.dataOffset(with: localFileHeader)
         let transaction = try await dataSource.openRead()
-        try await transaction.seek(to: entry.dataOffset + range.lowerBound)
+        try await transaction.seek(to: offset + range.lowerBound)
 
         _ = try await Data.consumePart(
             of: Int64(range.count),
@@ -207,14 +205,16 @@ extension Archive {
     private func extractCompressedRange(
         _ range: Range<UInt64>,
         of entry: Entry,
+        localFileHeader: LocalFileHeader,
         bufferSize: Int,
         consumer: Consumer
     ) async throws {
+        let offset = try entry.dataOffset(with: localFileHeader)
         let transaction = try await dataSource.openRead()
         let bytesReadCounter = SharedMutableValue<UInt64>()
 
         do {
-            try await transaction.seek(to: entry.dataOffset)
+            try await transaction.seek(to: offset)
 
             _ = try await readCompressed(
                 transaction: transaction,

Sources/ZIPFoundation/Archive+Writing.swift

@@ -97,6 +97,8 @@ extension Archive {
                               compressionMethod: compressionMethod, bufferSize: bufferSize,
                               progress: progress, provider: provider)
         }
+        
+        didWrite()
     }
 
     /// Write files, directories or symlinks to the receiver.
@@ -191,6 +193,8 @@ extension Archive {
             try await rollback(transaction, UInt64(fileHeaderStart), (existingData, existingSize), bufferSize, eocdRecord, zip64EOCD)
             throw ArchiveError.cancelledOperation
         }
+        
+        didWrite()
     }
 
     /// Remove a ZIP `Entry` from the receiver.
@@ -208,14 +212,16 @@ extension Archive {
             throw ArchiveError.unwritableArchive
         }
 
+        let lfh = try await localFileHeader(for: entry)
         let transaction = try await dataSource.openWrite()
         let (tempArchive, tempDir) = try await self.makeTempArchive()
         let tempTransaction = try await tempArchive.dataSource.openWrite()
         defer { tempDir.map { try? FileManager().removeItem(at: $0) } }
-        progress?.totalUnitCount = self.totalUnitCountForRemoving(entry)
+        progress?.totalUnitCount = try self.totalUnitCountForRemoving(entry, localFileHeader: lfh)
         var centralDirectoryData = Data()
         var offset: UInt64 = 0
-        for try await currentEntry in self {
+        for currentEntry in try await entries() {
+            let currentEntryLFH = try await localFileHeader(for: currentEntry)
             let cds = currentEntry.centralDirectoryStructure
             if currentEntry != entry {
                 let entryStart = cds.effectiveRelativeOffsetOfLocalHeader
@@ -228,13 +234,15 @@ extension Archive {
                     try await tempTransaction.write(data)
                     progress?.completedUnitCount += Int64(data.count)
                 }
-                guard currentEntry.localSize <= .max else { throw ArchiveError.invalidLocalHeaderSize }
-                _ = try await Data.consumePart(of: Int64(currentEntry.localSize), chunkSize: bufferSize,
+                let localSize = try currentEntry.localSize(with: currentEntryLFH)
+                _ = try await Data.consumePart(of: Int64(localSize), chunkSize: bufferSize,
                                                provider: provider, consumer: consumer)
                 let updatedCentralDirectory = updateOffsetInCentralDirectory(centralDirectoryStructure: cds,
                                                                              updatedOffset: entryStart - offset)
                 centralDirectoryData.append(updatedCentralDirectory.data)
-            } else { offset = currentEntry.localSize }
+            } else {
+                offset = try currentEntry.localSize(with: currentEntryLFH)
+            }
         }
 
         let startOfCentralDirectory = try await tempTransaction.position()
@@ -252,6 +260,8 @@ extension Archive {
         self.setEndOfCentralDirectory(ecodStructure)
         try await tempTransaction.flush()
         try await self.replaceCurrentArchive(with: tempArchive)
+        
+        didWrite()
     }
 
     func replaceCurrentArchive(with archive: Archive) async throws {
@@ -270,6 +280,8 @@ extension Archive {
         _ = try fileManager.moveItem(at: archiveURL, to: url)
 #endif
         self.dataSource = try await FileDataSource(url: url, isWritable: true)
+        
+        didWrite()
     }
 }
 

Sources/ZIPFoundation/Archive.swift

@@ -52,9 +52,7 @@ let centralDirectoryStructSignature = 0x02014b50
 ///     var archiveURL = URL(fileURLWithPath: "/path/file.zip")
 ///     var archive = Archive(url: archiveURL, accessMode: .update)
 ///     try archive?.addEntry("test.txt", relativeTo: baseURL, compressionMethod: .deflate)
-public actor Archive: AsyncSequence {
-    public typealias Element = Entry
-    
+public actor Archive {
 
     typealias LocalFileHeader = Entry.LocalFileHeader
     typealias DataDescriptor = Entry.DefaultDataDescriptor
@@ -79,6 +77,8 @@ public actor Archive: AsyncSequence {
         case invalidBufferSize
         /// Thrown when uncompressedSize/compressedSize exceeds `Int64.max` (Imposed by file API).
         case invalidEntrySize
+        /// Thrown when the local header cannot be found.
+        case localHeaderNotFound
         /// Thrown when the offset of local header data exceeds `Int64.max` (Imposed by file API).
         case invalidLocalHeaderDataOffset
         /// Thrown when the size of local header exceeds `Int64.max` (Imposed by file API).
@@ -206,80 +206,42 @@ public actor Archive: AsyncSequence {
         self.zip64EndOfCentralDirectory = config.zip64EndOfCentralDirectory
     }
 
-    public nonisolated func makeAsyncIterator() -> Iterator {
-        Iterator(archive: self)
-    }
+    private var entriesTask: Task<[Entry], Error>?
 
-    public actor Iterator: AsyncIteratorProtocol {
-        
-        private struct Input {
-            let transaction: DataSourceTransaction
-            let totalNumberOfEntriesInCD: UInt64
+    /// Returns the list of entries in the archive.
+    public func entries() async throws -> [Entry] {
+        if entriesTask == nil {
+            entriesTask = Task { try await readEntries() }
         }
 
-        private let archive: Archive
-        private var directoryIndex: UInt64 = 0
-        private var index = 0
-        
-        fileprivate init(archive: Archive) {
-            self.archive = archive
+        return try await entriesTask!.value
+    }
+    
+    private func readEntries() async throws -> [Entry] {
+        guard totalNumberOfEntriesInCentralDirectory > 0 else {
+            return []
         }
 
-        private var _initializeTask: Task<DataSourceTransaction, Error>?
+        var entries: [Entry] = []
+        var directoryIndex = offsetToStartOfCentralDirectory
+        let transaction = try await dataSource.openRead()
 
-        private func initialize() async throws -> DataSourceTransaction {
-            if _initializeTask == nil {
-                _initializeTask = Task {
-                    directoryIndex = await archive.offsetToStartOfCentralDirectory
-                    return try await archive.dataSource.openRead()
-                }
+        for _ in 0..<totalNumberOfEntriesInCentralDirectory {
+            guard let centralDirStruct: CentralDirectoryStructure = try await transaction.readStruct(at: directoryIndex) else {
+                continue
             }
 
-            return try await _initializeTask!.value
-        }
-        
-        public func next() async throws -> Entry? {
-            let totalNumberOfEntries = await archive.totalNumberOfEntriesInCentralDirectory
-            guard index < totalNumberOfEntries else {
-                return nil
-            }
-
-            let dataSource = try await initialize()
-
-            do {
-                guard let centralDirStruct: CentralDirectoryStructure = try await dataSource.readStruct(at: directoryIndex) else {
-                    return nil
-                }
-                let offset = UInt64(centralDirStruct.effectiveRelativeOffsetOfLocalHeader)
-                guard let localFileHeader: LocalFileHeader = try await dataSource.readStruct(at: offset) else { return nil }
-                var dataDescriptor: DataDescriptor?
-                var zip64DataDescriptor: ZIP64DataDescriptor?
-                if centralDirStruct.usesDataDescriptor {
-                    let additionalSize = UInt64(localFileHeader.fileNameLength) + UInt64(localFileHeader.extraFieldLength)
-                    let isCompressed = centralDirStruct.compressionMethod != CompressionMethod.none.rawValue
-                    let dataSize = isCompressed
-                    ? centralDirStruct.effectiveCompressedSize
-                    : centralDirStruct.effectiveUncompressedSize
-                    let descriptorPosition = offset + UInt64(LocalFileHeader.size) + additionalSize + dataSize
-                    if centralDirStruct.isZIP64 {
-                        zip64DataDescriptor = try await dataSource.readStruct(at: descriptorPosition)
-                    } else {
-                        dataDescriptor = try await dataSource.readStruct(at: descriptorPosition)
-                    }
-                }
-                defer {
-                    directoryIndex += UInt64(CentralDirectoryStructure.size)
-                    directoryIndex += UInt64(centralDirStruct.fileNameLength)
-                    directoryIndex += UInt64(centralDirStruct.extraFieldLength)
-                    directoryIndex += UInt64(centralDirStruct.fileCommentLength)
-                    index += 1
-                }
-                return Entry(centralDirectoryStructure: centralDirStruct, localFileHeader: localFileHeader,
-                             dataDescriptor: dataDescriptor, zip64DataDescriptor: zip64DataDescriptor)
-            } catch {
-                return nil
+            if let entry = Entry(centralDirectoryStructure: centralDirStruct) {
+                entries.append(entry)
             }
+            
+            directoryIndex += UInt64(CentralDirectoryStructure.size)
+            directoryIndex += UInt64(centralDirStruct.fileNameLength)
+            directoryIndex += UInt64(centralDirStruct.extraFieldLength)
+            directoryIndex += UInt64(centralDirStruct.fileCommentLength)
         }
+        
+        return entries
     }
 
     /// Retrieve the ZIP `Entry` with the given `path` from the receiver.
@@ -291,10 +253,64 @@ public actor Archive: AsyncSequence {
     /// - Parameter path: A relative file path identifying the corresponding `Entry`.
     /// - Returns: An `Entry` with the given `path`. Otherwise, `nil`.
     public func get(_ path: String) async throws -> Entry? {
-        if let encoding = self.pathEncoding {
-            return try await self.first { $0.path(using: encoding) == path }
+        let result = try await entries().first {
+            if let encoding = self.pathEncoding {
+                return $0.path(using: encoding) == path
+            } else {
+                return $0.path == path
+            }
         }
-        return try await self.first { $0.path == path }
+        return result
+    }
+    
+    /// Cache for local file headers.
+    private var localFileHeaders: [String: Task<LocalFileHeader, Error>] = [:]
+    
+    /// Retrieves the local file header for the given `entry`.
+    func localFileHeader(for entry: Entry) async throws -> LocalFileHeader {
+        if let task = localFileHeaders[entry.path] {
+            return try await task.value
+        }
+        
+        let task = Task {
+            let transaction = try await dataSource.openRead()
+            let centralDirStruct = entry.centralDirectoryStructure
+            let offset = UInt64(centralDirStruct.effectiveRelativeOffsetOfLocalHeader)
+            guard
+                var localFileHeader: LocalFileHeader = try await transaction.readStruct(at: offset)
+            else {
+                throw Archive.ArchiveError.localHeaderNotFound
+            }
+            
+            /// We only load the data descriptors if we are in writing mode, because
+            /// they might need to be written over. In read mode it is is
+            /// superfluous as the same infos are in the central directory structure.
+            if centralDirStruct.usesDataDescriptor && accessMode != .read {
+                let additionalSize = UInt64(localFileHeader.fileNameLength) + UInt64(localFileHeader.extraFieldLength)
+                let isCompressed = centralDirStruct.compressionMethod != CompressionMethod.none.rawValue
+                let dataSize = isCompressed
+                ? centralDirStruct.effectiveCompressedSize
+                : centralDirStruct.effectiveUncompressedSize
+                let descriptorPosition = offset + UInt64(LocalFileHeader.size) + additionalSize + dataSize
+                if centralDirStruct.isZIP64 {
+                    localFileHeader.zip64DataDescriptor = try await transaction.readStruct(at: descriptorPosition)
+                } else {
+                    localFileHeader.dataDescriptor = try await transaction.readStruct(at: descriptorPosition)
+                }
+            }
+            
+            return localFileHeader
+        }
+        localFileHeaders[entry.path] = task
+        
+        return try await task.value
+    }
+    
+    /// Called when the archive was modified.
+    func didWrite() {
+        // Clears the caches.
+        entriesTask = nil
+        localFileHeaders = [:]
     }
 
     // MARK: - Helpers