From 97207a7ad50afe6a67ba567cd08a7e84eb04ca0a Mon Sep 17 00:00:00 2001
From: EYHN <cneyhn@gmail.com>
Date: Tue, 21 Jan 2025 00:16:21 +0000
Subject: [PATCH] fix(ios): fix ios http and ws (#9805)

---
 .../apps/ios/App/App/AffineHttpHandler.swift  | 19 ++++++++++++++++---
 .../apps/ios/App/App/AffineWsHandler.swift    |  6 +++---
 packages/frontend/apps/ios/src/setup.ts       | 14 +++++++-------
 3 files changed, 26 insertions(+), 13 deletions(-)

diff --git a/packages/frontend/apps/ios/App/App/AffineHttpHandler.swift b/packages/frontend/apps/ios/App/App/AffineHttpHandler.swift
index 6426858bf2..4a0267de40 100644
--- a/packages/frontend/apps/ios/App/App/AffineHttpHandler.swift
+++ b/packages/frontend/apps/ios/App/App/AffineHttpHandler.swift
@@ -53,8 +53,8 @@ class AffineHttpHandler: NSObject, WKURLSchemeHandler {
       request.setValue(value, forHTTPHeaderField: key)
     }
     
-    URLSession.shared.dataTask(with: request) {
-      rawData, rawResponse, error in
+    let task = URLSession.shared.dataTask(with: request) {
+      (rawData, rawResponse, error) in
       urlSchemeTask.stopped?.withLock({
         if $0 {
           return
@@ -91,12 +91,16 @@ class AffineHttpHandler: NSObject, WKURLSchemeHandler {
         }
       })
     }
+    task.resume()
+    
+    urlSchemeTask.dataTask = task
   }
   
   func webView(_ webView: WKWebView, stop urlSchemeTask: WKURLSchemeTask) {
     urlSchemeTask.stopped?.withLock({
       $0 = true
     })
+    urlSchemeTask.dataTask?.cancel()
   }
 }
 
@@ -106,9 +110,18 @@ private extension WKURLSchemeTask {
       return objc_getAssociatedObject(self, &stoppedKey) as? Mutex<Bool> ?? nil
     }
     set {
-      objc_setAssociatedObject(self, &stoppedKey, newValue, .OBJC_ASSOCIATION_ASSIGN)
+      objc_setAssociatedObject(self, &stoppedKey, newValue, .OBJC_ASSOCIATION_RETAIN)
+    }
+  }
+  var dataTask: URLSessionDataTask? {
+    get {
+      return objc_getAssociatedObject(self, &dataTaskKey) as? URLSessionDataTask
+    }
+    set {
+      objc_setAssociatedObject(self, &dataTaskKey, newValue, .OBJC_ASSOCIATION_RETAIN)
     }
   }
 }
 
 private var stoppedKey = malloc(1)
+private var dataTaskKey = malloc(1)
diff --git a/packages/frontend/apps/ios/App/App/AffineWsHandler.swift b/packages/frontend/apps/ios/App/App/AffineWsHandler.swift
index be72ad0a1c..ae78327641 100644
--- a/packages/frontend/apps/ios/App/App/AffineWsHandler.swift
+++ b/packages/frontend/apps/ios/App/App/AffineWsHandler.swift
@@ -168,7 +168,7 @@ class AffineWsHandler: NSObject, WKURLSchemeHandler {
   
   func webView(_ webView: WKWebView, stop urlSchemeTask: WKURLSchemeTask) {
     urlSchemeTask.stopped?.withLock({
-      $0 = false
+      $0 = true
     })
     urlSchemeTask.wsTask?.cancel(with: .abnormalClosure, reason: "Closed".data(using: .utf8))
   }
@@ -180,7 +180,7 @@ private extension WKURLSchemeTask {
       return objc_getAssociatedObject(self, &stoppedKey) as? Mutex<Bool> ?? nil
     }
     set {
-      objc_setAssociatedObject(self, &stoppedKey, newValue, .OBJC_ASSOCIATION_ASSIGN)
+      objc_setAssociatedObject(self, &stoppedKey, newValue, .OBJC_ASSOCIATION_RETAIN)
     }
   }
   var wsTask: URLSessionWebSocketTask? {
@@ -188,7 +188,7 @@ private extension WKURLSchemeTask {
       return objc_getAssociatedObject(self, &wsTaskKey) as? URLSessionWebSocketTask
     }
     set {
-      objc_setAssociatedObject(self, &stoppedKey, newValue, .OBJC_ASSOCIATION_ASSIGN)
+      objc_setAssociatedObject(self, &wsTaskKey, newValue, .OBJC_ASSOCIATION_RETAIN)
     }
   }
 }
diff --git a/packages/frontend/apps/ios/src/setup.ts b/packages/frontend/apps/ios/src/setup.ts
index 68904eaf0b..f9fec4569b 100644
--- a/packages/frontend/apps/ios/src/setup.ts
+++ b/packages/frontend/apps/ios/src/setup.ts
@@ -14,25 +14,25 @@ import '@affine/core/bootstrap/browser';
  */
 const rawFetch = globalThis.fetch;
 globalThis.fetch = async (input: RequestInfo | URL, init?: RequestInit) => {
-  const url = new URL(
+  let url = new URL(
     typeof input === 'string'
       ? input
       : input instanceof URL
         ? input.toString()
         : input.url,
     globalThis.location.origin
-  );
+  ).href;
 
-  if (url.protocol === 'capacitor:') {
+  if (url.startsWith('capacitor:')) {
     return rawFetch(input, init);
   }
 
-  if (url.protocol === 'http:') {
-    url.protocol = 'affine-http:';
+  if (url.startsWith('http:')) {
+    url = 'affine-http:' + url.slice(5);
   }
 
-  if (url.protocol === 'https:') {
-    url.protocol = 'affine-https:';
+  if (url.startsWith('https:')) {
+    url = 'affine-https:' + url.slice(6);
   }
 
   return rawFetch(url, input instanceof Request ? input : init);