Merge branch 'charlesmchen/websocketEdgeCases'

src/Network/WebSockets/TSSocketManager.m

@@ -57,7 +57,8 @@ NSString *const SocketConnectingNotification = @"SocketConnectingNotification";
     OWSAssert([NSThread isMainThread]);
 
     _signalService = [OWSSignalService new];
-    
+    _status = kSocketStatusClosed;
+
     [self addObserver:self forKeyPath:@"status" options:0 context:kSocketStatusObservationContext];
 
     return self;
@@ -117,6 +118,13 @@ NSString *const SocketConnectingNotification = @"SocketConnectingNotification";
     DDLogWarn(@"Creating new websocket");
 
     // If socket is not already open or connecting, connect now.
+    //
+    // First we need to close the existing websocket, if any.
+    // The websocket delegate methods are invoked _after_ the websocket
+    // state changes, so we may be just learning about a socket failure
+    // or close event now.
+    self.status = kSocketStatusClosed;
+    // Now open a new socket.
     self.status = kSocketStatusConnecting;
 }
 
@@ -152,6 +160,10 @@ NSString *const SocketConnectingNotification = @"SocketConnectingNotification";
 
     // If this status update is redundant, verify that
     // class state and socket state are aligned.
+    //
+    // Note: it's not safe to check the socket's readyState here as
+    //       it may have been just updated on another thread. If so,
+    //       we'll learn of that state change soon.
     if (_status == status) {
         switch (status) {
             case kSocketStatusClosed:
@@ -159,11 +171,9 @@ NSString *const SocketConnectingNotification = @"SocketConnectingNotification";
                 break;
             case kSocketStatusOpen:
                 OWSAssert(self.websocket);
-                OWSAssert([self.websocket readyState] == SR_OPEN);
                 break;
             case kSocketStatusConnecting:
                 OWSAssert(self.websocket);
-                OWSAssert([self.websocket readyState] == SR_CONNECTING);
                 break;
         }
         return;
@@ -211,7 +221,13 @@ NSString *const SocketConnectingNotification = @"SocketConnectingNotification";
             socket.delegate = self;
             
             [self setWebsocket:socket];
+
+            // [SRWebSocket open] could hypothetically call a delegate method (e.g. if
+            // the socket failed immediately for some reason), so we update the status
+            // _before_ calling it, not after.
+            _status = status;
             [socket open];
+            return;
         }
     }
 
@@ -250,12 +266,22 @@ NSString *const SocketConnectingNotification = @"SocketConnectingNotification";
 
 - (void)webSocketDidOpen:(SRWebSocket *)webSocket {
     OWSAssert([NSThread isMainThread]);
-    
+    OWSAssert(webSocket);
+    if (webSocket != self.websocket) {
+        // Ignore events from obsolete web sockets.
+        return;
+    }
+
     self.status = kSocketStatusOpen;
 }
 
 - (void)webSocket:(SRWebSocket *)webSocket didFailWithError:(NSError *)error {
     OWSAssert([NSThread isMainThread]);
+    OWSAssert(webSocket);
+    if (webSocket != self.websocket) {
+        // Ignore events from obsolete web sockets.
+        return;
+    }
 
     DDLogError(@"Websocket did fail with error: %@", error);
 
@@ -266,6 +292,11 @@ NSString *const SocketConnectingNotification = @"SocketConnectingNotification";
 
 - (void)webSocket:(SRWebSocket *)webSocket didReceiveMessage:(NSData *)data {
     OWSAssert([NSThread isMainThread]);
+    OWSAssert(webSocket);
+    if (webSocket != self.websocket) {
+        // Ignore events from obsolete web sockets.
+        return;
+    }
 
     WebSocketMessage *wsMessage = [WebSocketMessage parseFromData:data];
 
@@ -354,6 +385,11 @@ NSString *const SocketConnectingNotification = @"SocketConnectingNotification";
            reason:(NSString *)reason
          wasClean:(BOOL)wasClean {
     OWSAssert([NSThread isMainThread]);
+    OWSAssert(webSocket);
+    if (webSocket != self.websocket) {
+        // Ignore events from obsolete web sockets.
+        return;
+    }
 
     DDLogWarn(@"Websocket did close with code: %ld", (long)code);