From 131d757cf8e6e368ad338728379174f7cfff9326 Mon Sep 17 00:00:00 2001
From: Nghia Tran <vinhnghiatran@gmail.com>
Date: Thu, 9 Nov 2023 08:59:27 +0700
Subject: [PATCH] Option to opt-out the Method Swizzling on WS/WSS Connection
 (#138)

* Able to opt-out the WS/WSS

* Bump 1.23.0
---
 Sources/Atlantis.swift        |  8 +++++---
 Sources/NetworkInjector.swift | 24 +++++++++++++++++++-----
 atlantis-proxyman.podspec     |  2 +-
 3 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/Sources/Atlantis.swift b/Sources/Atlantis.swift
index 343a3e1..fd6d4f3 100644
--- a/Sources/Atlantis.swift
+++ b/Sources/Atlantis.swift
@@ -77,14 +77,16 @@ public final class Atlantis: NSObject {
     /// Build version of Atlantis
     /// It's essential for Proxyman to known if it's compatible with this version
     /// Instead of receving the number from the info.plist, we should hardcode here because the info file doesn't exist in SPM
-    public static let buildVersion: String = "1.22.0"
+    public static let buildVersion: String = "1.23.0"
 
     /// Start Swizzle all network functions and monitoring the traffic
     /// It also starts looking Bonjour network from Proxyman app.
     /// If hostName is nil, Atlantis will find all Proxyman apps in the network. It's useful if we have only one machine for personal usage.
     /// If hostName is not nil, Atlantis will try to connect to particular mac machine. It's useful if you have multiple Proxyman.
     /// - Parameter hostName: Host name of Mac machine. You can find your current Host Name in Proxyman -> Certificate -> Install on iOS -> By Atlantis -> Show Start Atlantis
-    @objc public class func start(hostName: String? = nil) {
+    /// - Parameter shouldCaptureWebSocketTraffic: Determine if Atlantis should perform the Method Swizzling on WS/WSS connection. Default is true.
+    @objc public class func start(hostName: String? = nil, shouldCaptureWebSocketTraffic: Bool = true) {
+        // save config
         let configuration = Configuration.default(hostName: hostName)
 
         //
@@ -105,7 +107,7 @@ public final class Atlantis: NSObject {
 
         // Enable the injector
         Atlantis.shared.configuration = configuration
-        Atlantis.shared.injector.injectAllNetworkClasses()
+        Atlantis.shared.injector.injectAllNetworkClasses(config: NetworkConfiguration(shouldCaptureWebSocketTraffic: shouldCaptureWebSocketTraffic))
 
         // Start transport layer if need
         if Atlantis.shared.isEnabledTransportLayer {
diff --git a/Sources/NetworkInjector.swift b/Sources/NetworkInjector.swift
index f89a1a7..69426c1 100644
--- a/Sources/NetworkInjector.swift
+++ b/Sources/NetworkInjector.swift
@@ -8,10 +8,21 @@
 
 import Foundation
 
+struct NetworkConfiguration {
+
+    /// Whether or not Atlantis should perform the Method Swizzling on WS/WSS connection.
+    let shouldCaptureWebSocketTraffic: Bool
+
+    // init
+    init(shouldCaptureWebSocketTraffic: Bool = true) {
+        self.shouldCaptureWebSocketTraffic = shouldCaptureWebSocketTraffic
+    }
+}
+
 protocol Injector {
 
     var delegate: InjectorDelegate? { get set }
-    func injectAllNetworkClasses()
+    func injectAllNetworkClasses(config: NetworkConfiguration)
 }
 
 protocol InjectorDelegate: AnyObject {
@@ -38,17 +49,17 @@ final class NetworkInjector: Injector {
 
     // MARK: - Internal
 
-    func injectAllNetworkClasses() {
+    func injectAllNetworkClasses(config: NetworkConfiguration = NetworkConfiguration()) {
         // Make sure we swizzle *ONCE*
         DispatchQueue.once {
-            injectAllURLSession()
+            injectAllURLSession(config)
         }
     }
 }
 
 extension NetworkInjector {
 
-    private func injectAllURLSession() {
+    private func injectAllURLSession(_ config: NetworkConfiguration) {
 
         // iOS 8: __NSCFURLSessionConnection
         // iOS 9, 10, 11, 12, 13, 14: __NSCFURLLocalSessionConnection
@@ -66,7 +77,10 @@ extension NetworkInjector {
         injectURLSessionUploadTasks()
 
         // Websocket
-        injectURLSessionWebsocketTasks()
+        // Able to opt-out the WS/WSS if needed
+        if config.shouldCaptureWebSocketTraffic {
+            injectURLSessionWebsocketTasks()
+        }
     }
 
     private func injectIntoURLSessionDelegate(anyClass: AnyClass) {
diff --git a/atlantis-proxyman.podspec b/atlantis-proxyman.podspec
index 7885232..bffd507 100644
--- a/atlantis-proxyman.podspec
+++ b/atlantis-proxyman.podspec
@@ -1,6 +1,6 @@
 Pod::Spec.new do |spec|
   spec.name         = "atlantis-proxyman"
-  spec.version      = "1.22.0"
+  spec.version      = "1.23.0"
   spec.summary      = "A lightweight and powerful iOS framework for intercepting HTTP/HTTPS Traffic"
   spec.description  = <<-DESC
   A lightweight and powerful iOS framework for intercepting HTTP/HTTPS Traffic from your app. No more messing around with proxy, certificate config.