From 2ba167f899064fc623af006a59a4331308284d93 Mon Sep 17 00:00:00 2001
From: Lorenzo Blasa <lblasa@meta.com>
Date: Fri, 27 Jan 2023 04:07:37 -0800
Subject: [PATCH] Back out "Clear handlers on disconnect"

Summary:
^

Revert as to validate this is not causing regressions: T143523262

Reviewed By: passy

Differential Revision: D42800560

fbshipit-source-id: 8db61454eabfdb259637bb97c2bb4754984ecf6f
---
 iOS/FlipperKit/FlipperPlatformWebSocket.h  | 10 +++---
 iOS/FlipperKit/FlipperPlatformWebSocket.mm | 38 +++++++++++-----------
 iOS/FlipperKit/FlipperWebSocket.mm         | 18 +++++-----
 3 files changed, 33 insertions(+), 33 deletions(-)

diff --git a/iOS/FlipperKit/FlipperPlatformWebSocket.h b/iOS/FlipperKit/FlipperPlatformWebSocket.h
index 07eb7d251..0044b134f 100644
--- a/iOS/FlipperKit/FlipperPlatformWebSocket.h
+++ b/iOS/FlipperKit/FlipperPlatformWebSocket.h
@@ -23,15 +23,15 @@ NS_ASSUME_NONNULL_BEGIN
 /// A message handler used to dispatch messages received from the server.
 @property(nonatomic) facebook::flipper::SocketMessageHandler messageHandler;
 
+/// A certificate provider used to obtain the client certificate used for
+/// authentication.
+@property(nonatomic)
+    facebook::flipper::SocketCertificateProvider certificateProvider;
+
 /// Initializes an instance of FliperWebSocketTransport with an endpoint URL.
 /// @param url Endpoint URL used to establish the connection.
 - (instancetype)initWithURL:(NSURL* _Nonnull)url;
 
-/// A certificate provider used to obtain the client certificate used for
-/// authentication.
-- (void)setCertificateProvider:
-    (facebook::flipper::SocketCertificateProvider)certificateProvider;
-
 /// Connect to the endpoint.
 - (void)connect;
 
diff --git a/iOS/FlipperKit/FlipperPlatformWebSocket.mm b/iOS/FlipperKit/FlipperPlatformWebSocket.mm
index 8116579ea..a51c5ac27 100644
--- a/iOS/FlipperKit/FlipperPlatformWebSocket.mm
+++ b/iOS/FlipperKit/FlipperPlatformWebSocket.mm
@@ -144,10 +144,6 @@ static constexpr int connectionKeepaliveSeconds = 10;
     return;
   }
 
-  _socket = [[SRWebSocket alloc] initWithURL:self->_url
-                              securityPolicy:self->_policy];
-  _socket.delegate = self;
-
   __weak auto weakSelf = self;
   [_dispatchQueue addOperationWithBlock:^{
     __strong auto strongSelf = weakSelf;
@@ -162,32 +158,33 @@ static constexpr int connectionKeepaliveSeconds = 10;
       return;
     }
 
+    strongSelf->_socket = [[SRWebSocket alloc] initWithURL:self->_url
+                                            securityPolicy:self->_policy];
+    strongSelf->_socket.delegate = self;
     [strongSelf->_socket open];
   }];
 }
 
 - (void)disconnect {
-  _socket.delegate = nil;
-
-  // Manually trigger a 'close' event as SocketRocket close method will
-  // not notify the delegate. SocketRocket only triggers the close event
-  // when the connection is closed from the server. Furthermore,
-  // we are clearing the delegate above.
-  _eventHandler(facebook::flipper::SocketEvent::CLOSE);
-
-  _eventHandler = [](facebook::flipper::SocketEvent) {};
-  _messageHandler = ^(const std::string&) {
-  };
-  _policy.certificateProvider = [](char* _Nonnull, size_t) { return ""; };
-
   [_dispatchQueue cancelAllOperations];
-  [_dispatchQueue waitUntilAllOperationsAreFinished];
 
   if ([_keepAlive isValid]) {
     [_keepAlive invalidate];
   }
   _keepAlive = nil;
-  _socket = nil;
+
+  // Manually trigger a 'close' event as SocketRocket close method will
+  // not notify the delegate. SocketRocket only triggers the close event
+  // when the connection is closed from the server.
+  _eventHandler(facebook::flipper::SocketEvent::CLOSE);
+
+  if (_socket) {
+    // Clear the socket delegate before close. Ensures that we won't get
+    // any messages after the disconnect takes place.
+    _socket.delegate = nil;
+    [_socket close];
+    _socket = nil;
+  };
 }
 
 - (void)send:(NSString*)message
@@ -205,6 +202,7 @@ static constexpr int connectionKeepaliveSeconds = 10;
 
 - (void)setCertificateProvider:
     (facebook::flipper::SocketCertificateProvider)certificateProvider {
+  _certificateProvider = certificateProvider;
   _policy.certificateProvider = certificateProvider;
 }
 
@@ -243,6 +241,7 @@ static constexpr int connectionKeepaliveSeconds = 10;
   } else {
     _eventHandler(facebook::flipper::SocketEvent::ERROR);
   }
+  _socket = nil;
 }
 
 - (void)_webSocketDidClose {
@@ -252,6 +251,7 @@ static constexpr int connectionKeepaliveSeconds = 10;
   _keepAlive = nil;
 
   _eventHandler(facebook::flipper::SocketEvent::CLOSE);
+  _socket = nil;
 }
 
 - (void)_webSocketDidReceiveMessage:(id)message {
diff --git a/iOS/FlipperKit/FlipperWebSocket.mm b/iOS/FlipperKit/FlipperWebSocket.mm
index 24480c7b1..3389d7791 100644
--- a/iOS/FlipperKit/FlipperWebSocket.mm
+++ b/iOS/FlipperKit/FlipperWebSocket.mm
@@ -112,15 +112,15 @@ bool FlipperWebSocket::connect(FlipperConnectionManager* manager) {
   };
 
   if (endpoint_.secure) {
-    [socket_
-        setCertificateProvider:[this](char* _Nonnull password, size_t length) {
-          auto pkcs12 = connectionContextStore_->getCertificate();
-          if (pkcs12.first.length() == 0) {
-            return std::string("");
-          }
-          strncpy(password, pkcs12.second.c_str(), length);
-          return pkcs12.first;
-        }];
+    socket_.certificateProvider = [this](
+                                      char* _Nonnull password, size_t length) {
+      auto pkcs12 = connectionContextStore_->getCertificate();
+      if (pkcs12.first.length() == 0) {
+        return std::string("");
+      }
+      strncpy(password, pkcs12.second.c_str(), length);
+      return pkcs12.first;
+    };
   }
 
   [socket_ connect];