From 9f4304b4b53b029f3fa86f7b797b7fa51eb442d1 Mon Sep 17 00:00:00 2001
From: wwwcg <isarraypro@gmail.com>
Date: Tue, 15 Oct 2024 16:50:35 +0800
Subject: [PATCH] fix(ios): websocket module support extra headers

---
 ios/sdk/debug/websocket/HippySRWebSocket.h     |  3 +++
 ios/sdk/debug/websocket/HippySRWebSocket.m     | 18 +++++++++++++++---
 .../debug/websocket/HippyWebSocketManager.m    | 16 ++++++++++++----
 3 files changed, 30 insertions(+), 7 deletions(-)

diff --git a/ios/sdk/debug/websocket/HippySRWebSocket.h b/ios/sdk/debug/websocket/HippySRWebSocket.h
index 997b7096357..ae520a0252f 100644
--- a/ios/sdk/debug/websocket/HippySRWebSocket.h
+++ b/ios/sdk/debug/websocket/HippySRWebSocket.h
@@ -66,6 +66,9 @@ extern NSString *const HippySRHTTPResponseErrorKey;
 
 // Some helper constructors.
 - (instancetype)initWithURL:(NSURL *)url protocols:(NSArray<NSString *> *)protocols;
+- (instancetype)initWithURL:(NSURL *)url 
+               extraHeaders:(NSDictionary *)extraHeaders
+                  protocols:(NSArray<NSString *> *)protocols;
 - (instancetype)initWithURL:(NSURL *)url;
 
 // Delegate queue will be dispatch_main_queue by default.
diff --git a/ios/sdk/debug/websocket/HippySRWebSocket.m b/ios/sdk/debug/websocket/HippySRWebSocket.m
index 50291f4287c..6a74c9494b2 100644
--- a/ios/sdk/debug/websocket/HippySRWebSocket.m
+++ b/ios/sdk/debug/websocket/HippySRWebSocket.m
@@ -278,8 +278,9 @@ - (instancetype)initWithURLRequest:(NSURLRequest *)request;
 - (instancetype)initWithURL:(NSURL *)URL;
 { return [self initWithURL:URL protocols:nil]; }
 
-- (instancetype)initWithURL:(NSURL *)URL protocols:(NSArray<NSString *> *)protocols;
-{
+- (instancetype)initWithURL:(NSURL *)URL
+               extraHeaders:(NSDictionary *)extraHeaders
+                  protocols:(NSArray<NSString *> *)protocols {
     NSMutableURLRequest *request;
     if (URL) {
         // Build a mutable request so we can fill the cookie header.
@@ -295,11 +296,22 @@ - (instancetype)initWithURL:(NSURL *)URL protocols:(NSArray<NSString *> *)protoc
 
         // Load and set the cookie header.
         NSArray<NSHTTPCookie *> *cookies = [[NSHTTPCookieStorage sharedHTTPCookieStorage] cookiesForURL:components.URL];
-        [request setAllHTTPHeaderFields:[NSHTTPCookie requestHeaderFieldsWithCookies:cookies]];
+        NSDictionary *cookieHeaders = [NSHTTPCookie requestHeaderFieldsWithCookies:cookies];
+        
+        // Set extraHeaders passed from the js
+        NSMutableDictionary *allHeaders = cookieHeaders.mutableCopy ?: [NSMutableDictionary dictionary];
+        if (extraHeaders) {
+            [allHeaders addEntriesFromDictionary:extraHeaders];
+        }
+        [request setAllHTTPHeaderFields:allHeaders];
     }
     return [self initWithURLRequest:request protocols:protocols];
 }
 
+- (instancetype)initWithURL:(NSURL *)URL protocols:(NSArray<NSString *> *)protocols {
+    return [self initWithURL:URL extraHeaders:nil protocols:protocols];
+}
+
 - (void)_HippySR_commonInit;
 {
     NSString *scheme = _url.scheme.lowercaseString;
diff --git a/ios/sdk/debug/websocket/HippyWebSocketManager.m b/ios/sdk/debug/websocket/HippyWebSocketManager.m
index 6f3bcb87842..c80570bdf2c 100644
--- a/ios/sdk/debug/websocket/HippyWebSocketManager.m
+++ b/ios/sdk/debug/websocket/HippyWebSocketManager.m
@@ -28,6 +28,9 @@
 #import "HippySRWebSocket.h"
 
 static NSUInteger socketIndex = 0;
+static NSString *const kHippyWebSocketUrlParamKey = @"url";
+static NSString *const kHippyWebSocketHeaderParamKey = @"headers";
+static NSString *const kHippySecWebSocketProtocolKey = @"Sec-WebSocket-Protocol";
 
 #pragma mark - HippyWebSocketManager
 
@@ -62,11 +65,16 @@ - (void)invalidate {
 }
 
 HIPPY_EXPORT_METHOD(connect:(NSDictionary *)params resolver:(HippyPromiseResolveBlock)resolve rejecter:(HippyPromiseRejectBlock)reject) {
-    NSDictionary *headers = params[@"headers"];
-    NSString *url = params[@"url"];
-    NSString *protocols = headers[@"Sec-WebSocket-Protocol"];
+    NSDictionary *headers = params[kHippyWebSocketHeaderParamKey];
+    NSString *url = params[kHippyWebSocketUrlParamKey];
+    NSString *protocols = headers[kHippySecWebSocketProtocolKey];
+    // prepare extra headers
+    NSMutableDictionary *extraHeaders = headers.mutableCopy;
+    [extraHeaders removeObjectForKey:kHippySecWebSocketProtocolKey];
     NSArray<NSString *> *protocolArray = [protocols componentsSeparatedByString:@","];
-    HippySRWebSocket *socket = [[HippySRWebSocket alloc] initWithURL:[NSURL URLWithString:url] protocols:protocolArray];
+    HippySRWebSocket *socket = [[HippySRWebSocket alloc] initWithURL:[NSURL URLWithString:url]
+                                                        extraHeaders:extraHeaders
+                                                           protocols:protocolArray];
     socket.delegate = self;
     socket.socketID = socketIndex++;
     NSNumber *socketId = @(socket.socketID);