Fix a race condition where a download size could exceed maxSize (#10412)

Author: paulb777

FirebaseStorage/CHANGELOG.md

@@ -1,3 +1,6 @@
+# 10.2.0
+- [fixed] Fix a race condition where a download size could exceed the value of the `maxSize` parameter. (#10358)
+
 # 10.1.0
 - [fixed] Fixed a 10.0.0 regression where metadata passed to `putFile` was not properly initialized. (#10353)
 - [fixed] Fixed a 10.0.0 regression handling an empty JSON metadata field from the emulator. (#10370)

FirebaseStorage/Sources/StorageReference.swift

@@ -275,10 +275,12 @@ import Foundation
     let callbackQueue = fetcherService.callbackQueue ?? DispatchQueue.main
 
     task.observe(.success) { snapshot in
+      let error = self.checkSizeOverflow(task: snapshot.task, maxSize: maxSize)
       callbackQueue.async {
         if !completed {
           completed = true
-          completion(task.downloadData, nil)
+          let data = error == nil ? task.downloadData : nil
+          completion(data, error)
         }
       }
     }
@@ -291,14 +293,8 @@ import Foundation
       }
     }
     task.observe(.progress) { snapshot in
-      let task = snapshot.task
-      if task.progress.totalUnitCount > maxSize || task.progress.completedUnitCount > maxSize {
-        let error = StorageErrorCode.error(withCode: .downloadSizeExceeded,
-                                           infoDictionary: [
-                                             "totalSize": task.progress.totalUnitCount,
-                                             "maxAllowedSize": maxSize,
-                                           ])
-        (task as? StorageDownloadTask)?.cancel(withError: error)
+      if let error = self.checkSizeOverflow(task: snapshot.task, maxSize: maxSize) {
+        task.cancel(withError: error)
       }
     }
     task.enqueue()
@@ -695,4 +691,18 @@ import Foundation
     self.storage = storage
     self.path = path
   }
+
+  /**
+   * For maxSize API, return an error if the size is exceeded.
+   */
+  private func checkSizeOverflow(task: StorageTask, maxSize: Int64) -> NSError? {
+    if task.progress.totalUnitCount > maxSize || task.progress.completedUnitCount > maxSize {
+      return StorageErrorCode.error(withCode: .downloadSizeExceeded,
+                                    infoDictionary: [
+                                      "totalSize": task.progress.totalUnitCount,
+                                      "maxAllowedSize": maxSize,
+                                    ])
+    }
+    return nil
+  }
 }