fix: send booru headers for image requests

1 files changed, 61 insertions, 0 deletions

diff --git a/Sora/Data/Booru/BooruNetworkImageLoader.swift b/Sora/Data/Booru/BooruNetworkImageLoader.swift
new file mode 100644
index 0000000..c2a0f8a
--- /dev/null
+++ b/Sora/Data/Booru/BooruNetworkImageLoader.swift
@@ -0,0 +1,61 @@
+import CoreGraphics
+import Foundation
+import ImageIO
+import NetworkImage
+
+actor BooruNetworkImageLoader: NetworkImageLoader {
+  private let domain: String
+  private let sendUserAgent: Bool
+  private let customUserAgent: String
+  private var ongoingTasks: [URL: Task<CGImage, Error>] = [:]
+
+  init(
+    domain: String,
+    sendUserAgent: Bool,
+    customUserAgent: String
+  ) {
+    self.domain = domain
+    self.sendUserAgent = sendUserAgent
+    self.customUserAgent = customUserAgent
+  }
+
+  func image(from url: URL) async throws -> CGImage {
+    if let task = ongoingTasks[url] {
+      return try await task.value
+    }
+
+    let domain = self.domain
+    let sendUserAgent = self.sendUserAgent
+    let customUserAgent = self.customUserAgent
+
+    let task = Task<CGImage, Error> {
+      guard let data = await ImageCacheManager.shared.loadImageData(
+        for: url,
+        domain: domain,
+        sendUserAgent: sendUserAgent,
+        customUserAgent: customUserAgent
+      ) else {
+        throw URLError(.badServerResponse)
+      }
+
+      guard
+        let source = CGImageSourceCreateWithData(data as CFData, nil),
+        let image = CGImageSourceCreateImageAtIndex(
+          source,
+          0,
+          [kCGImageSourceShouldCache: true] as CFDictionary
+        )
+      else {
+        throw URLError(.cannotDecodeContentData)
+      }
+
+      return image
+    }
+
+    ongoingTasks[url] = task
+
+    defer { ongoingTasks.removeValue(forKey: url) }
+
+    return try await task.value
+  }
+}