Mp4 Restructure account for large mdat box size (#78)

dimitris-c · May 9, 2024 · a8865bb · a8865bb
1 parent dd2e790
commit a8865bb
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 32 deletions.
diff --git a/AudioStreaming/Streaming/Audio Source/Mp4/Mp4Restructure.swift b/AudioStreaming/Streaming/Audio Source/Mp4/Mp4Restructure.swift
@@ -108,17 +108,23 @@ final class Mp4Restructure {
         return (initialData, mdatOffset)
     }
 
-    func checkIsOptimized(data: Data) -> (optimized: Bool, offset: Int?)? {
+    func checkIsOptimized(data: Data) throws -> (optimized: Bool, offset: Int?)? {
         while atomOffset < UInt64(data.count) {
-            let atomSize = Int(readUInt32FromData(data: data, offset: atomOffset))
-            let atomType = Int(readUInt32FromData(data: data, offset: atomOffset + 4))
+            var atomSize = try Int(getInteger(data: data, offset: atomOffset) as UInt32)
+            let atomType = try Int(getInteger(data: data, offset: atomOffset + 4) as UInt32)
             switch atomType {
             case Atoms.ftyp:
                 let ftypData = data[Int(atomOffset) ..< atomSize]
                 let ftyp = MP4Atom(type: atomType, size: atomSize, offset: atomOffset, data: ftypData)
                 self.ftyp = ftyp
                 atoms.append(ftyp)
             case Atoms.mdat:
+                // ref: https://developer.apple.com/documentation/quicktime-file-format/movie_data_atom
+                // This atom can be quite large, and may exceed 2^32 bytes, in which case the size field will be set to 1, 
+                // and the header will contain a 64-bit extended size field.
+                if atomSize == 1 {
+                    atomSize =  Int(try getInteger(data: data, offset: atomOffset + 8) as UInt64)
+                }
                 let mdat = MP4Atom(type: atomType, size: atomSize, offset: atomOffset)
                 atoms.append(mdat)
                 foundMdat = true
@@ -241,8 +247,12 @@ final class Mp4Restructure {
         return (moovAtom.storage, moovAtomSize)
     }
 
-    private func readUInt32FromData(data: Data, offset: Int) -> UInt32 {
-        let valueData = data.subdata(in: offset ..< offset + 4)
-        return UInt32(bigEndian: valueData.withUnsafeBytes { $0.load(as: UInt32.self) })
+    func getInteger<T: FixedWidthInteger>(data: Data, offset: Int) throws -> T {
+        let sizeOfInteger = MemoryLayout<T>.size
+        guard sizeOfInteger <= data.count else {
+            throw ByteBuffer.Error.eof
+        }
+        let _offset = offset + sizeOfInteger
+        return T(data: data[_offset - sizeOfInteger ..< _offset]).bigEndian
     }
 }
diff --git a/AudioStreaming/Streaming/Audio Source/Mp4/RemoteMp4Restructure.swift b/AudioStreaming/Streaming/Audio Source/Mp4/RemoteMp4Restructure.swift
@@ -74,35 +74,39 @@ final class RemoteMp4Restructure {
                         return
                     }
                     self.audioData.append(data)
-                    let value = self.mp4Restructure.checkIsOptimized(data: self.audioData)
-                    if let value {
-                        if let offset = value.offset, !value.optimized {
-                            guard response.response?.statusCode == 206 else {
-                                Logger.error("⛔️ mp4 error: no moov before mdat and the stream is not seekable", category: .networking)
-                                completion(.failure(Mp4RestructureError.nonOptimizedMp4AndServerCannotSeek))
-                                return
-                            }
-                            // stop request, fetch moov and restructure
-                            self.audioData = Data()
-                            self.task?.cancel()
-                            self.task = nil
-                            self.fetchAndRestructureMoovAtom(offset: offset) { result in
-                                switch result {
-                                case let .success(value):
-                                    let data = value.data
-                                    let offset = value.offset
-                                    self.dataOptimized = true
-                                    completion(.success(RestructuredData(initialData: data, mdatOffset: offset)))
-                                case let .failure(error):
-                                    completion(.failure(Mp4RestructureError.networkError(error)))
+                    do {
+                        let value = try self.mp4Restructure.checkIsOptimized(data: self.audioData)
+                        if let value {
+                            if let offset = value.offset, !value.optimized {
+                                guard response.response?.statusCode == 206 else {
+                                    Logger.error("⛔️ mp4 error: no moov before mdat and the stream is not seekable", category: .networking)
+                                    completion(.failure(Mp4RestructureError.nonOptimizedMp4AndServerCannotSeek))
+                                    return
+                                }
+                                // stop request, fetch moov and restructure
+                                self.audioData = Data()
+                                self.task?.cancel()
+                                self.task = nil
+                                self.fetchAndRestructureMoovAtom(offset: offset) { result in
+                                    switch result {
+                                    case let .success(value):
+                                        let data = value.data
+                                        let offset = value.offset
+                                        self.dataOptimized = true
+                                        completion(.success(RestructuredData(initialData: data, mdatOffset: offset)))
+                                    case let .failure(error):
+                                        completion(.failure(Mp4RestructureError.networkError(error)))
+                                    }
                                 }
+                            } else {
+                                self.audioData = Data()
+                                self.task?.cancel()
+                                self.task = nil
+                                completion(.success(nil))
                             }
-                        } else {
-                            self.audioData = Data()
-                            self.task?.cancel()
-                            self.task = nil
-                            completion(.success(nil))
                         }
+                    } catch {
+                        completion(.failure(Mp4RestructureError.invalidAtomSize))
                     }
                 case let .stream(.failure(error)):
                     completion(.failure(Mp4RestructureError.networkError(error)))