From 799d88275eba43e97392f48ca085f8a327e13448 Mon Sep 17 00:00:00 2001
From: Lorenzo Blasa <lblasa@fb.com>
Date: Sun, 3 Oct 2021 01:16:40 -0700
Subject: [PATCH] Remove WebSocket delegate before closing

Summary:
This change removes ourselves as a delegate before closing.

SocketRocket uses its own internal async queue to perform most operations.

After a disconnect, we don't expect to receive any more delegate calls as the handlers may contain references which may have become invalid.

So, removing ourselves as delegates will ensure that we don't get called after a disconnect.

For sanity, we are also taking a copy of the message handler instead of a reference to it.

Reviewed By: briantkelley

Differential Revision: D31360721

fbshipit-source-id: bae5a2423757cd9064ffac28afb8b78c28a20d87
---
 iOS/FlipperKit/FlipperPlatformWebSocket.mm | 4 +++-
 iOS/FlipperKit/FlipperWebSocket.mm         | 3 ++-
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/iOS/FlipperKit/FlipperPlatformWebSocket.mm b/iOS/FlipperKit/FlipperPlatformWebSocket.mm
index 505e99c8c..27010e68d 100644
--- a/iOS/FlipperKit/FlipperPlatformWebSocket.mm
+++ b/iOS/FlipperKit/FlipperPlatformWebSocket.mm
@@ -156,7 +156,9 @@ static constexpr int connectionKeepaliveSeconds = 10;
     // not notify the delegate. SocketRocket only triggers the close event
     // when the connection is closed from the server.
     _eventHandler(facebook::flipper::SocketEvent::CLOSE);
-
+    // 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;
   }
diff --git a/iOS/FlipperKit/FlipperWebSocket.mm b/iOS/FlipperKit/FlipperWebSocket.mm
index d73f429b1..d3bcec289 100644
--- a/iOS/FlipperKit/FlipperWebSocket.mm
+++ b/iOS/FlipperKit/FlipperWebSocket.mm
@@ -170,8 +170,9 @@ bool FlipperWebSocket::connect(FlipperConnectionManager* manager) {
     }
     eventHandler(event);
   };
+  auto messageHandler = messageHandler_;
   socket_.messageHandler = ^(const std::string& message) {
-    this->messageHandler_(message);
+    messageHandler(message);
   };
 
   if (endpoint_.secure) {