Merge pull request #1124 from DataDog/maxep/RUMM-2819/optimize-tlv-read

RUMM-2819 TLV Lenght Limit

CHANGELOG.md

@@ -1,8 +1,10 @@
 # Unreleased
 
+- [BUGFIX] Fix 'Could not allocate memory' after corrupted TLV. See [#1089][] (Thanks [@cltnschlosser][])
+
 # 1.14.0 / 20-12-2022
 
-- [IMPROVEMENT] Add a method for sending error attributes on logs as strings[#1051][].
+- [IMPROVEMENT] Add a method for sending error attributes on logs as strings. See [#1051][].
 - [IMPROVEMENT] Add manual Open Telemetry b3 headers injection. See [#1057][]
 - [IMPROVEMENT] Add automatic Open Telemetry b3 headers injection. See [#1061][]
 - [IMPROVEMENT] Add manual and automatic W3C traceparent header injection. See [#1071][]
@@ -423,6 +425,7 @@
 [#1057]: https://github.com/DataDog/dd-sdk-ios/pull/1057
 [#1061]: https://github.com/DataDog/dd-sdk-ios/pull/1061
 [#1071]: https://github.com/DataDog/dd-sdk-ios/pull/1071
+[#1089]: https://github.com/DataDog/dd-sdk-ios/pull/1089
 [@00fa9a]: https://github.com/00FA9A
 [@britton-earnin]: https://github.com/Britton-Earnin
 [@hengyu]: https://github.com/Hengyu
@@ -449,3 +452,4 @@
 [@sdejesusf]: https://github.com/sdejesusF
 [@avdlee]: https://github.com/AvdLee
 [@dfed]: https://github.com/dfed
+[@cltnschlosser]: https://github.com/cltnschlosser

Sources/Datadog/DatadogCore/DatadogCore.swift

@@ -24,6 +24,7 @@ internal final class DatadogCore {
     /// The storage r/w GDC queue.
     let readWriteQueue = DispatchQueue(
         label: "com.datadoghq.ios-sdk-read-write",
+        autoreleaseFrequency: .workItem,
         target: .global(qos: .utility)
     )
 

Sources/Datadog/DatadogCore/Storage/DataBlock.swift

@@ -9,6 +9,9 @@ import Foundation
 /// Block size binary type
 internal typealias BlockSize = UInt32
 
+/// Default max data lenght in block (safety check) - 10 MB
+private let MAX_DATA_LENGHT: UInt64 = 10 * 1_024 * 1_024
+
 /// Block type supported in data stream
 internal enum BlockType: UInt16 {
     case event = 0x00
@@ -17,8 +20,11 @@ internal enum BlockType: UInt16 {
 /// Reported errors while manipulating data blocks.
 internal enum DataBlockError: Error {
     case readOperationFailed(streamError: Error?)
+    case invalidDataType
     case invalidByteSequence
-    case dataLenghtExceedsLimit
+    case bytesLengthExceedsLimit(limit: UInt64)
+    case dataAllocationFailure
+    case endOfStream
 }
 
 /// A data block in defined by its type and a byte sequence.
@@ -38,15 +44,15 @@ internal struct DataBlock {
     ///     +-  2 bytes -+-   4 bytes   -+- n bytes -|
     ///     | block type | data size (n) |    data   |
     ///     +------------+---------------+-----------+
-    ///
+    /// - Parameter maxLenght: Maximum data lenght of a block.
     /// - Returns: a data block in TLV.
-    func serialize() throws -> Data {
+    func serialize(maxLenght: UInt64 = MAX_DATA_LENGHT) throws -> Data {
         var buffer = Data()
         // T
         withUnsafeBytes(of: type.rawValue) { buffer.append(contentsOf: $0) }
         // L
-        guard let length = BlockSize(exactly: data.count) else {
-            throw DataBlockError.dataLenghtExceedsLimit
+        guard let length = BlockSize(exactly: data.count), length <= maxLenght else {
+            throw DataBlockError.bytesLengthExceedsLimit(limit: maxLenght)
         }
         withUnsafeBytes(of: length) { buffer.append(contentsOf: $0) }
         // V
@@ -63,14 +69,25 @@ internal final class DataBlockReader {
     /// The input data stream.
     private let stream: InputStream
 
+    /// Maximum data lenght of a block.
+    private let maxBlockLenght: UInt64
+
     /// Reads block from data input.
     ///
     /// At initilization, the reader will open a stream targeting the input data.
     /// The stream will be closed when the reader instance is deallocated.
     ///
-    /// - Parameter data: The data input
-    convenience init(data: Data) {
-        self.init(input: InputStream(data: data))
+    /// - Parameters:
+    ///   - data: The data input
+    ///   - maxBlockLenght: Maximum data lenght of a block.
+    convenience init(
+        data: Data,
+        maxBlockLenght: UInt64 = MAX_DATA_LENGHT
+    ) {
+        self.init(
+            input: InputStream(data: data),
+            maxBlockLenght: maxBlockLenght
+        )
     }
 
     /// Reads block from an input stream.
@@ -79,7 +96,11 @@ internal final class DataBlockReader {
     /// when the reader instance is deallocated.
     ///
     /// - Parameter stream: The input stream
-    init(input stream: InputStream) {
+    init(
+        input stream: InputStream,
+        maxBlockLenght: UInt64 = MAX_DATA_LENGHT
+    ) {
+        self.maxBlockLenght = maxBlockLenght
         self.stream = stream
         stream.open()
     }
@@ -97,26 +118,25 @@ internal final class DataBlockReader {
     /// - Returns: The next block or nil if none could be found.
     func next() throws -> DataBlock? {
         // look for the next known block
-        while stream.hasBytesAvailable {
-            // read an entire block before inferring the data type
-            // to leave the stream in a usuable state if an unkown
-            // type was encountered.
-            let type = try readType()
-            let data = try readData()
-
-            if let type = BlockType(rawValue: type) {
-                return DataBlock(type: type, data: data)
+        while true {
+            do {
+                return try readBlock()
+            } catch DataBlockError.invalidDataType {
+                continue
+            } catch DataBlockError.endOfStream {
+                // Some streams won't return false for hasBytesAvailable until a read is attempted
+                return nil
+            } catch {
+                throw error
             }
         }
-
-        return nil
     }
 
     /// Reads all data blocks from current index in the stream.
     ///
     /// - Throws: `DataBlockError` while reading the input stream.
     /// - Returns: The block sequence found in the input
-    func all() throws -> [DataBlock] {
+    func all(maxDataLenght: UInt64 = MAX_DATA_LENGHT) throws -> [DataBlock] {
         var blocks: [DataBlock] = []
 
         while let block = try next() {
@@ -132,18 +152,47 @@ internal final class DataBlockReader {
     /// - Throws: `DataBlockError` while reading the input stream.
     /// - Returns: Data bytes from stream.
     private func read(length: Int) throws -> Data {
-        var bytes = [UInt8](repeating: 0, count: length)
-        let count = stream.read(&bytes, maxLength: length)
+        guard length > 0 else {
+            return Data()
+        }
+
+        // Load from stream directly to data without unnecessary copies
+        var data = Data(count: length)
+        let count = try data.withUnsafeMutableBytes {
+            guard let buffer = $0.assumingMemoryBound(to: UInt8.self).baseAddress else {
+                throw DataBlockError.dataAllocationFailure
+            }
+            return stream.read(buffer, maxLength: length)
+        }
 
         if count < 0 {
             throw DataBlockError.readOperationFailed(streamError: stream.streamError)
         }
 
+        if count == 0 {
+            throw DataBlockError.endOfStream
+        }
+
         guard count == length else {
             throw DataBlockError.invalidByteSequence
         }
 
-        return Data(bytes)
+        return data
+    }
+
+    /// Reads a block.
+    private func readBlock() throws -> DataBlock {
+        // read an entire block before inferring the data type
+        // to leave the stream in a usuable state if an unkown
+        // type was encountered.
+        let type = try readType()
+        let data = try readData()
+
+        guard let type = BlockType(rawValue: type) else {
+            throw DataBlockError.invalidDataType
+        }
+
+        return DataBlock(type: type, data: data)
     }
 
     /// Reads a block type.
@@ -160,8 +209,10 @@ internal final class DataBlockReader {
         // even if `Int` is able to represent all `BlockSize` on 64 bit
         // arch, we make sure to avoid overflow and get the exact data
         // length.
-        guard let length = Int(exactly: size) else {
-            throw DataBlockError.dataLenghtExceedsLimit
+        // Additionally check that length hasn't been corrupted and
+        // we don't try to generate a huge buffer.
+        guard let length = Int(exactly: size), length <= maxBlockLenght else {
+            throw DataBlockError.bytesLengthExceedsLimit(limit: maxBlockLenght)
         }
 
         return try read(length: length)

Sources/Datadog/DatadogCore/Storage/Files/File.swift

@@ -27,13 +27,17 @@ internal protocol ReadableFile {
     /// Name of this file.
     var name: String { get }
 
-    /// Reads the available data in this file.
-    func read() throws -> Data
+    /// Creates InputStream for reading the available data from this file.
+    func stream() throws -> InputStream
 
     /// Deletes this file.
     func delete() throws
 }
 
+private enum FileError: Error {
+    case unableToCreateInputStream
+}
+
 /// An immutable `struct` designed to provide optimized and thread safe interface for file manipulation.
 /// It doesn't own the file, which means the file presence is not guaranteed - the file can be deleted by OS at any time (e.g. due to memory pressure).
 internal struct File: WritableFile, ReadableFile {
@@ -90,43 +94,11 @@ internal struct File: WritableFile, ReadableFile {
         }
     }
 
-    func read() throws -> Data {
-        let fileHandle = try FileHandle(forReadingFrom: url)
-
-        // NOTE: RUMM-669
-        // https://github.com/DataDog/dd-sdk-ios/issues/214
-        // https://en.wikipedia.org/wiki/Xcode#11.x_series
-        // compiler version needs to have iOS 13.4+ as base SDK
-        #if compiler(>=5.2)
-        /**
-         Even though the `fileHandle.seekToEnd()` should be available since iOS 13.0:
-         ```
-         @available(OSX 10.15, iOS 13.0, watchOS 6.0, tvOS 13.0, *)
-         public func readToEnd() throws -> Data?
-         ```
-         it crashes on iOS Simulators prior to iOS 13.4:
-         ```
-         Symbol not found: _$sSo12NSFileHandleC10FoundationE9readToEndAC4DataVSgyKF
-         ```
-        This is fixed in iOS 14/Xcode 12
-        */
-        if #available(iOS 13.4, tvOS 13.4, *) {
-            defer { try? fileHandle.close() }
-            return try fileHandle.readToEnd() ?? Data()
-        } else {
-            return try legacyRead(from: fileHandle)
-        }
-        #else
-        return try legacyRead(from: fileHandle)
-        #endif
-    }
-
-    private func legacyRead(from fileHandle: FileHandle) throws -> Data {
-        let data = fileHandle.readDataToEndOfFile()
-        try? objcExceptionHandler.rethrowToSwift {
-            fileHandle.closeFile()
+    func stream() throws -> InputStream {
+        guard let stream = InputStream(url: url) else {
+            throw FileError.unableToCreateInputStream
         }
-        return data
+        return stream
     }
 
     func size() throws -> UInt64 {

Sources/Datadog/DatadogCore/Storage/FilesOrchestrator.swift

@@ -9,11 +9,12 @@ import Foundation
 /// Orchestrates files in a single directory.
 internal class FilesOrchestrator {
     /// Directory where files are stored.
-    private let directory: Directory
+    let directory: Directory
     /// Date provider.
-    private let dateProvider: DateProvider
+    let dateProvider: DateProvider
     /// Performance rules for writing and reading files.
-    private let performance: StoragePerformancePreset
+    let performance: StoragePerformancePreset
+
     /// Name of the last file returned by `getWritableFile()`.
     private var lastWritableFileName: String? = nil
     /// Tracks number of times the file at `lastWritableFileURL` was returned from `getWritableFile()`.
@@ -33,10 +34,6 @@ internal class FilesOrchestrator {
     // MARK: - `WritableFile` orchestration
 
     func getWritableFile(writeSize: UInt64) throws -> WritableFile {
-        if writeSize > performance.maxObjectSize {
-            throw InternalError(description: "data exceeds the maximum size of \(performance.maxObjectSize) bytes.")
-        }
-
         let lastWritableFileOrNil = reuseLastWritableFileIfPossible(writeSize: writeSize)
 
         if let lastWritableFile = lastWritableFileOrNil { // if last writable file can be reused

Sources/Datadog/DatadogCore/Storage/Reading/FileReader.swift

@@ -31,7 +31,7 @@ internal final class FileReader: Reader {
         }
 
         do {
-            let events = try decode(data: file.read())
+            let events = try decode(stream: file.stream())
             return Batch(events: events, file: file)
         } catch {
             DD.telemetry.error("Failed to read data from file", error: error)
@@ -45,10 +45,13 @@ internal final class FileReader: Reader {
     /// consumed and decrypted if encryption is available. Decrypted events are finally joined with
     /// data-format separator.
     ///
-    /// - Parameter data: The data to decode.
+    /// - Parameter stream: The InputStream that provides data to decode.
     /// - Returns: The decoded and formatted data.
-    private func decode(data: Data) throws -> [Data] {
-        let reader = DataBlockReader(data: data)
+    private func decode(stream: InputStream) throws -> [Data] {
+        let reader = DataBlockReader(
+            input: stream,
+            maxBlockLenght: orchestrator.performance.maxObjectSize
+        )
 
         var failure: String? = nil
         defer {

Sources/Datadog/DatadogCore/Storage/Writing/FileWriter.swift

@@ -55,7 +55,9 @@ internal final class FileWriter: Writer {
         return try DataBlock(
             type: .event,
             data: encrypt(data: data)
-        ).serialize()
+        ).serialize(
+            maxLenght: orchestrator.performance.maxObjectSize
+        )
     }
 
     /// Encrypts data if encryption is available.

Sources/Datadog/DatadogCore/Upload/FeatureUpload.swift

@@ -20,6 +20,7 @@ internal struct FeatureUpload {
     ) {
         let uploadQueue = DispatchQueue(
             label: "com.datadoghq.ios-sdk-\(featureName)-upload",
+            autoreleaseFrequency: .workItem,
             target: .global(qos: .utility)
         )