fix(ios): websocket module support extra headers (#4069) (#4081)

wwwcg · wwwcg · commit ca3967f9988e · 2024-10-24T11:15:27.000+08:00
diff --git a/framework/ios/debug/websocket/HippySRWebSocket.h b/framework/ios/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/framework/ios/debug/websocket/HippySRWebSocket.m b/framework/ios/debug/websocket/HippySRWebSocket.m
@@ -276,8 +276,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.
@@ -293,11 +294,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/framework/ios/debug/websocket/HippyWebSocketManager.m b/framework/ios/debug/websocket/HippyWebSocketManager.m
@@ -28,6 +28,9 @@
 #import "HippyDefines.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);
