Fix synchronous download manager

CHANGELOG.md

@@ -2,6 +2,7 @@
 ### PaymentSheet
 * [Fixed] Fixed an animation glitch when dismissing PaymentSheet in React Native.
 * [Fixed] Fixed an issue with FlowController in vertical layout where the payment method could incorrectly be preserved across a call to `update` when it's no longer valid.
+* [Fixed] Fixed a potential deadlock when `paymentOption` is accessed from Swift concurrency.
 
 
 ## 23.32.0 2024-10-21

StripePaymentSheet/StripePaymentSheet/Source/Helpers/DownloadManager.swift

@@ -22,6 +22,8 @@ import UIKit
 
     private let session: URLSession
     private let analyticsClient: STPAnalyticsClient
+    private let imageCacheLock = NSLock()
+    private var imageCache: [URL: UIImage] = [:]
 
     public init(
         urlSessionConfiguration: URLSessionConfiguration = .default,
@@ -52,7 +54,7 @@ import UIKit
 extension DownloadManager {
 
     /// Downloads an image from a provided URL, using either a synchronous method or an asynchronous method.
-    /// If no `updateHandler` is provided, this function will block the current thread until the image is downloaded. If an `updateHandler` is provided, the function does not wait for the download to finish and returns a placeholder image immediately instead. When the image finishes downloading, the `updateHandler` will be called with the downloaded image.
+    /// If no `updateHandler` is provided, this function will block the current thread until the image is downloaded. If an `updateHandler` is provided, the function does not wait for the download to finish and returns the image if it was cached or a placeholder image instead. When the image finishes downloading, the `updateHandler` will be called with the downloaded image.
     /// - Parameters:
     ///   - url: The URL from which to download the image.
     ///   - placeholder: An optional parameter indicating a placeholder image to display while the download is in progress. If not provided, a default placeholder image will be used instead.
@@ -61,23 +63,30 @@ extension DownloadManager {
     /// - Returns: A `UIImage` instance. If `updateHandler` is `nil`, this would be the downloaded image, otherwise, this would be the placeholder image.
     public func downloadImage(url: URL, placeholder: UIImage?, updateHandler: UpdateImageHandler?) -> UIImage {
         let placeholder = placeholder ?? imagePlaceHolder()
-        // If no `updateHandler` is provided use a blocking method to fetch the image
-        guard let updateHandler else {
-            return downloadImageBlocking(url: url, placeholder: placeholder)
+        imageCacheLock.lock()
+        let cachedImage = imageCache[url]
+        imageCacheLock.unlock()
+
+        if let updateHandler {
+            Task {
+                await downloadImageAsync(url: url, placeholder: placeholder, updateHandler: updateHandler)
+            }
         }
-
-        Task {
-           await downloadImageAsync(url: url, placeholder: placeholder, updateHandler: updateHandler)
-        }
-        // Immediately return a placeholder, when the download operation completes `updateHandler` will be called with the downloaded image
-        return placeholder
+        // Immediately return the cached image or a placeholder. When the download operation completes `updateHandler` will be called with the downloaded image.
+        return cachedImage ?? placeholder
     }
 
     // Common download function
     private func downloadImage(url: URL, placeholder: UIImage) async -> UIImage {
         do {
             let (data, _) = try await session.data(from: url)
             let image = try UIImage.from(imageData: data) // Throws a Error.failedToMakeImageFromData
+            DispatchQueue.global(qos: .userInteractive).async {
+                // Cache the image in memory
+                self.imageCacheLock.lock()
+                self.imageCache[url] = image
+                self.imageCacheLock.unlock()
+            }
             return image
         } catch {
             let errorAnalytic = ErrorAnalytic(event: .stripePaymentSheetDownloadManagerError,
@@ -96,14 +105,11 @@ extension DownloadManager {
         }
     }
 
-    private func downloadImageBlocking(url: URL, placeholder: UIImage) -> UIImage {
-        return _unsafeWait({
-            return await self.downloadImage(url: url, placeholder: placeholder)
-        })
-    }
-
-    func resetDiskCache() {
+    func resetCache() {
         session.configuration.urlCache?.removeAllCachedResponses()
+        imageCacheLock.lock()
+        imageCache = [:]
+        imageCacheLock.unlock()
     }
 }
 
@@ -143,23 +149,3 @@ private extension UIImage {
         return image
     }
 }
-
-// MARK: Workarounds for using Swift async from a sync context
-// https://forums.swift.org/t/using-async-functions-from-synchronous-functions-and-breaking-all-the-rules/59782/4
-private class ImageBox {
-    var image: UIImage = DownloadManager.sharedManager.imagePlaceHolder()
-}
-
-private extension DownloadManager {
-    /// Unsafely awaits an async function from a synchronous context.
-    func _unsafeWait(_ downloadImage: @escaping () async -> UIImage) -> UIImage {
-        let imageBox = ImageBox()
-        let sema = DispatchSemaphore(value: 0)
-        Task {
-            imageBox.image = await downloadImage()
-            sema.signal()
-        }
-        sema.wait()
-        return imageBox.image
-    }
-}

StripePaymentSheet/StripePaymentSheetTests/PaymentSheet/DownloadManagerTest.swift

@@ -28,7 +28,7 @@ class DownloadManagerTest: APIStubbedTestCase {
         self.urlSessionConfig = APIStubbedTestCase.stubbedURLSessionConfig()
         self.analyticsClient = STPAnalyticsClient()
         self.rm = DownloadManager(urlSessionConfiguration: urlSessionConfig, analyticsClient: analyticsClient)
-        self.rm.resetDiskCache()
+        self.rm.resetCache()
     }
 
     func testURLCacheConfiguration() {
@@ -39,27 +39,54 @@ class DownloadManagerTest: APIStubbedTestCase {
         XCTAssertEqual(configurationUrlCache?.diskCapacity, 30_000_000)
     }
 
-    func testSynchronous_validImage() {
+    func testDownloadImageWithoutUpdateHandler_validImage() {
+        let imageData = validImageData()
         stub(condition: { request in
             return request.url?.path.contains("/validImage.png") ?? false
         }) { _ in
-            return HTTPStubsResponse(data: self.validImageData(), statusCode: 200, headers: nil)
+            return HTTPStubsResponse(data: imageData, statusCode: 200, headers: nil)
         }
-
-        let image = rm.downloadImage(url: validURL, placeholder: nil, updateHandler: nil)
-        XCTAssertEqual(image.size, validImageSize)
+        // Downloading an image the first time...
+        let firstDownloadExpectation = expectation(description: "First download completed")
+        var image1_final: UIImage!
+        let image1_initial = rm.downloadImage(url: validURL, placeholder: nil) { _image in
+            image1_final = _image
+            firstDownloadExpectation.fulfill()
+        }
+        // ...should return a placeholder...
+        XCTAssertEqual(image1_initial.size, placeholderImageSize)
+        // ...and call the completion with the image.
+        waitForExpectations(timeout: 1)
+        XCTAssertEqual(image1_final.size, validImageSize)
+
+        // Calling `downloadImage` a second time...
+        let image2 = rm.downloadImage(url: validURL, placeholder: nil, updateHandler: nil)
+        // ...should return the cached image
+        XCTAssertEqual(image2, image1_final)
     }
 
-    func testSynchronous_invalidImage() {
+    func testDownloadImageWithoutUpdateHandler_invalidImage() {
         stub(condition: { request in
             return request.url?.path.contains("/invalidImage.png") ?? false
         }) { _ in
             return HTTPStubsResponse(error: NotFoundError())
         }
 
-        let image = rm.downloadImage(url: invalidURL, placeholder: nil, updateHandler: nil)
-
-        XCTAssertEqual(image.size, placeholderImageSize)
+        // Downloading an invalid image the first time...
+        let firstDownloadExpectation = expectation(description: "First download completed")
+        firstDownloadExpectation.isInverted = true
+        let image1_initial = rm.downloadImage(url: invalidURL, placeholder: nil) { _ in
+            firstDownloadExpectation.fulfill()
+        }
+        // ...should return a placeholder...
+        XCTAssertEqual(image1_initial.size, placeholderImageSize)
+        // ...without calling the update handler.
+        waitForExpectations(timeout: 1)
+
+        // Calling `downloadImage` a second time...
+        let image2 = rm.downloadImage(url: validURL, placeholder: nil, updateHandler: nil)
+        // ...should continue to return the placeholder
+        XCTAssertEqual(image2.size, placeholderImageSize)
     }
 
     func testAsync_validImage() {
@@ -177,10 +204,11 @@ class DownloadManagerTest: APIStubbedTestCase {
         }
 
         let placeholder = rm.imagePlaceHolder()
-        let image = rm.downloadImage(url: validURL, placeholder: placeholder, updateHandler: nil)
-
+        let image = rm.downloadImage(url: validURL, placeholder: placeholder, updateHandler: { _ in })
         XCTAssertEqual(image, placeholder)
 
+        // Wait a beat for the error analytic to get sent.
+        wait(seconds: 0.1)
         // Validate analytic
         let firstAnalytic = try XCTUnwrap(analyticsClient._testLogHistory.first)
         XCTAssertEqual("stripepaymentsheet.downloadmanager.error", firstAnalytic["event"] as? String)
@@ -193,13 +221,15 @@ class DownloadManagerTest: APIStubbedTestCase {
         stub(condition: { request in
             return request.url == self.validURL
         }) { _ in
-            return HTTPStubsResponse(data: "invalid image data".data(using: .utf8)!, statusCode: 200, headers: nil)
+            return HTTPStubsResponse(data: Data("invalid image data".utf8), statusCode: 200, headers: nil)
         }
 
         let placeholder = rm.imagePlaceHolder()
-        let image = rm.downloadImage(url: validURL, placeholder: placeholder, updateHandler: nil)
+        let image = rm.downloadImage(url: validURL, placeholder: placeholder, updateHandler: { _ in })
 
         XCTAssertEqual(image, placeholder)
+        // Wait a beat for the error analytic to get sent.
+        wait(seconds: 0.1)
 
         // Validate analytic
         let firstAnalytic = try XCTUnwrap(analyticsClient._testLogHistory.first)

StripePaymentSheet/StripePaymentSheetTests/PaymentSheet/PaymentSheetPaymentMethodTypeTest.swift

@@ -29,7 +29,7 @@ class PaymentSheetPaymentMethodTypeTest: XCTestCase {
         FormSpecProvider.shared.load { _ in
             e.fulfill()
         }
-        DownloadManager.sharedManager.resetDiskCache()
+        DownloadManager.sharedManager.resetCache()
         waitForExpectations(timeout: 10)
         // A Payment methods with a client-side asset and a form spec image URL...
         let loadExpectation = expectation(description: "Load form spec image")
@@ -60,7 +60,7 @@ class PaymentSheetPaymentMethodTypeTest: XCTestCase {
     }
 
     func testMakeImage_without_client_asset() {
-        DownloadManager.sharedManager.resetDiskCache()
+        DownloadManager.sharedManager.resetCache()
         let e = expectation(description: "Load specs")
         FormSpecProvider.shared.load { _ in
             e.fulfill()