Fix web socket issue.

// FREEBIE

src/Network/WebSockets/TSSocketManager.m

@@ -1,9 +1,5 @@
 //
-//  TSSocketManager.m
+//  Copyright (c) 2017 Open Whisper Systems. All rights reserved.
-//  TextSecureiOS
-//
-//  Created by Frederic Jacobs on 17/05/14.
-//  Copyright (c) 2014 Open Whisper Systems. All rights reserved.
 //
 
 #import "SubProtocol.pb.h"
@@ -26,28 +22,31 @@ NSString *const SocketOpenedNotification     = @"SocketOpenedNotification";
 NSString *const SocketClosedNotification     = @"SocketClosedNotification";
 NSString *const SocketConnectingNotification = @"SocketConnectingNotification";
 
+// TSSocketManager's properties should only be accessed from the main thread.
 @interface TSSocketManager ()
 
-@property (nonatomic, readonly, strong) OWSSignalService *signalService;
+@property (nonatomic, readonly) OWSSignalService *signalService;
 
-@property (nonatomic, retain) NSTimer *pingTimer;
+@property (nonatomic) NSTimer *pingTimer;
-@property (nonatomic, retain) NSTimer *reconnectTimer;
+@property (nonatomic) NSTimer *reconnectTimer;
 
 
-@property (nonatomic, retain) SRWebSocket *websocket;
+@property (nonatomic) SRWebSocket *websocket;
 @property (nonatomic) SocketStatus status;
 
 @property (nonatomic) UIBackgroundTaskIdentifier fetchingTaskIdentifier;
 
-@property BOOL didConnectBg;
+@property (nonatomic) BOOL didConnectBg;
-@property BOOL didRetreiveMessageBg;
+@property (nonatomic) BOOL didRetreiveMessageBg;
-@property BOOL shouldDownloadMessage;
+@property (nonatomic) BOOL shouldDownloadMessage;
 
-@property (nonatomic, retain) NSTimer *backgroundKeepAliveTimer;
+@property (nonatomic) NSTimer *backgroundKeepAliveTimer;
-@property (nonatomic, retain) NSTimer *backgroundConnectTimer;
+@property (nonatomic) NSTimer *backgroundConnectTimer;
 
 @end
 
+#pragma mark -
+
 @implementation TSSocketManager
 
 - (instancetype)init
@@ -57,9 +56,11 @@ NSString *const SocketConnectingNotification = @"SocketConnectingNotification";
     if (!self) {
         return self;
     }
+    
+    OWSAssert([NSThread isMainThread]);
 
     _signalService = [OWSSignalService new];
-    _websocket = nil;
+    
     [self addObserver:self forKeyPath:@"status" options:0 context:kSocketStatusObservationContext];
 
     return self;
@@ -87,6 +88,21 @@ NSString *const SocketConnectingNotification = @"SocketConnectingNotification";
 
 - (void)becomeActive
 {
+    OWSAssert([NSThread isMainThread]);
+    
+    if ([NSThread isMainThread]) {
+        [self ensureWebsocket];
+    } else {
+        dispatch_async(dispatch_get_main_queue(), ^{
+            [self ensureWebsocket];
+        });
+    }
+}
+
+- (void)ensureWebsocket
+{
+    OWSAssert([NSThread isMainThread]);
+    
     if (self.signalService.isCensored) {
         DDLogWarn(@"%@ Refusing to start websocket in `becomeActive`.", self.tag);
         return;
@@ -98,10 +114,12 @@ NSString *const SocketConnectingNotification = @"SocketConnectingNotification";
         switch ([socket readyState]) {
             case SR_OPEN:
                 DDLogVerbose(@"WebSocket already open on connection request");
+                OWSAssert(self.status == kSocketStatusOpen);
                 self.status = kSocketStatusOpen;
                 return;
             case SR_CONNECTING:
                 DDLogVerbose(@"WebSocket is already connecting");
+                OWSAssert(self.status == kSocketStatusConnecting);
                 self.status = kSocketStatusConnecting;
                 return;
             default:
@@ -110,9 +128,10 @@ NSString *const SocketConnectingNotification = @"SocketConnectingNotification";
     }
 
     // Discard the old socket which is already closed or is closing.
-    [socket close];
+    [self closeWebSocket];
-    self.status = kSocketStatusClosed;
+
-    socket.delegate = nil;
+    OWSAssert(self.status == kSocketStatusClosed);
+    self.status = kSocketStatusConnecting;
 
     // Create a new web socket.
     NSString *webSocketConnect = [textSecureWebSocketAPI stringByAppendingString:[self webSocketAuthenticationString]];
@@ -127,13 +146,32 @@ NSString *const SocketConnectingNotification = @"SocketConnectingNotification";
 }
 
 + (void)resignActivity {
-    SRWebSocket *socket = [[self sharedManager] websocket];
+    OWSAssert([NSThread isMainThread]);
-    [socket close];
+
+    [[self sharedManager] closeWebSocket];
+}
+
+- (void)closeWebSocket
+{
+    OWSAssert([NSThread isMainThread]);
+
+    if (self.websocket) {
+        DDLogWarn(@"%@ closeWebSocket.", self.tag);
+    }
+
+    [self.websocket close];
+    self.websocket.delegate = nil;
+    self.websocket = nil;
+    [self.pingTimer invalidate];
+    self.pingTimer = nil;
+    self.status = kSocketStatusClosed;
 }
 
 #pragma mark - Delegate methods
 
 - (void)webSocketDidOpen:(SRWebSocket *)webSocket {
+    OWSAssert([NSThread isMainThread]);
+    
     self.pingTimer = [NSTimer timerWithTimeInterval:kWebSocketHeartBeat
                                              target:self
                                            selector:@selector(webSocketHeartBeat)
@@ -143,6 +181,7 @@ NSString *const SocketConnectingNotification = @"SocketConnectingNotification";
     // Additionally, we want the ping timer to work in the background too.
     [[NSRunLoop mainRunLoop] addTimer:self.pingTimer forMode:NSDefaultRunLoopMode];
 
+    OWSAssert(self.status == kSocketStatusConnecting);
     self.status = kSocketStatusOpen;
     [self.reconnectTimer invalidate];
     self.reconnectTimer = nil;
@@ -150,10 +189,12 @@ NSString *const SocketConnectingNotification = @"SocketConnectingNotification";
 }
 
 - (void)webSocket:(SRWebSocket *)webSocket didFailWithError:(NSError *)error {
+    OWSAssert([NSThread isMainThread]);
+    
     DDLogError(@"Error connecting to socket %@", error);
 
-    [self.pingTimer invalidate];
+    [self closeWebSocket];
-    self.status = kSocketStatusClosed;
+
     [self scheduleRetry];
 }
 
@@ -194,6 +235,8 @@ NSString *const SocketConnectingNotification = @"SocketConnectingNotification";
 }
 
 - (void)keepAliveBackground {
+    OWSAssert([NSThread isMainThread]);
+    
     [self.backgroundConnectTimer invalidate];
 
     if (self.fetchingTaskIdentifier) {
@@ -214,6 +257,8 @@ NSString *const SocketConnectingNotification = @"SocketConnectingNotification";
 }
 
 - (void)sendWebSocketMessageAcknowledgement:(WebSocketRequestMessage *)request {
+    OWSAssert([NSThread isMainThread]);
+
     WebSocketResponseMessageBuilder *response = [WebSocketResponseMessage builder];
     [response setStatus:200];
     [response setMessage:@"OK"];
@@ -227,6 +272,7 @@ NSString *const SocketConnectingNotification = @"SocketConnectingNotification";
     [self.websocket sendDataNoCopy:message.build.data error:&error];
     if (error) {
         DDLogWarn(@"Error while trying to write on websocket %@", error);
+        [self scheduleRetry];
     }
 }
 
@@ -234,20 +280,27 @@ NSString *const SocketConnectingNotification = @"SocketConnectingNotification";
  didCloseWithCode:(NSInteger)code
            reason:(NSString *)reason
          wasClean:(BOOL)wasClean {
-    [self.pingTimer invalidate];
+    OWSAssert([NSThread isMainThread]);
-    self.status = kSocketStatusClosed;
+
+    [self closeWebSocket];
 
-    if (!wasClean && [UIApplication sharedApplication].applicationState == UIApplicationStateActive) {
+    if (!wasClean && [self shouldKeepWebSocketAlive]) {
         [self scheduleRetry];
     }
 }
 
 - (void)webSocketHeartBeat {
-    NSError *error;
+    OWSAssert([NSThread isMainThread]);
-
+
-    [self.websocket sendPing:nil error:&error];
+    if ([self shouldKeepWebSocketAlive]) {
-    if (error) {
+        NSError *error;
-        DDLogWarn(@"Error in websocket heartbeat: %@", error.localizedDescription);
+        [self.websocket sendPing:nil error:&error];
+        if (error) {
+            DDLogWarn(@"Error in websocket heartbeat: %@", error.localizedDescription);
+            [self scheduleRetry];
+        }
+    } else {
+        [self closeWebSocket];
     }
 }
 
@@ -258,8 +311,29 @@ NSString *const SocketConnectingNotification = @"SocketConnectingNotification";
                          [TSStorageManager serverAuthToken]];
 }
 
+- (BOOL)shouldKeepWebSocketAlive {
+    if ([UIApplication sharedApplication].applicationState == UIApplicationStateActive) {
+        // If app is active, keep web socket alive.
+        return YES;
+    } else if (self.backgroundKeepAliveTimer ||
+               self.backgroundConnectTimer ||
+               self.fetchingTaskIdentifier != UIBackgroundTaskInvalid) {
+        // If app is doing any work in the background, keep web socket alive.
+        return YES;
+    } else {
+        return NO;
+    }
+}
+
 - (void)scheduleRetry {
-    if (![self.reconnectTimer isValid]) {
+    OWSAssert([NSThread isMainThread]);
+   
+    if (![self shouldKeepWebSocketAlive]) {
+        // Don't bother retrying if app is inactive and not doing any background activity.
+        [self.reconnectTimer invalidate];
+        self.reconnectTimer = nil;
+    } else if (![self.reconnectTimer isValid]) {
+        // TODO: It'd be nice to do exponential backoff.
         self.reconnectTimer = [NSTimer timerWithTimeInterval:kWebSocketReconnectTry
                                                       target:[self class]
                                                     selector:@selector(becomeActive)
@@ -275,37 +349,53 @@ NSString *const SocketConnectingNotification = @"SocketConnectingNotification";
 #pragma mark - Background Connect
 
 + (void)becomeActiveFromForeground {
-    TSSocketManager *sharedInstance = [self sharedManager];
+    dispatch_async(dispatch_get_main_queue(), ^{
+        [[self sharedManager] becomeActiveFromForeground];
+    });
+}
 
-    if (sharedInstance.fetchingTaskIdentifier != UIBackgroundTaskInvalid) {
+- (void)becomeActiveFromForeground {
-        [sharedInstance closeBackgroundTask];
+    OWSAssert([NSThread isMainThread]);
+    
+    if (self.fetchingTaskIdentifier != UIBackgroundTaskInvalid) {
+        [self closeBackgroundTask];
     }
-
+    
     [self becomeActive];
 }
 
 + (void)becomeActiveFromBackgroundExpectMessage:(BOOL)expected {
-    TSSocketManager *sharedInstance = [TSSocketManager sharedManager];
+    dispatch_async(dispatch_get_main_queue(), ^{
-
+        [[TSSocketManager sharedManager] becomeActiveFromBackgroundExpectMessage:expected];
-    if (sharedInstance.fetchingTaskIdentifier == UIBackgroundTaskInvalid) {
+    });
-        sharedInstance.backgroundConnectTimer = [NSTimer timerWithTimeInterval:kBackgroundConnectTimer
+}
-                                                                        target:sharedInstance
+
-                                                                      selector:@selector(backgroundConnectTimerExpired)
+- (void)becomeActiveFromBackgroundExpectMessage:(BOOL)expected {
-                                                                      userInfo:nil
+    OWSAssert([NSThread isMainThread]);
-                                                                       repeats:NO];
+    
+    if (self.fetchingTaskIdentifier == UIBackgroundTaskInvalid) {
+        [self.backgroundConnectTimer invalidate];
+        self.backgroundConnectTimer = [NSTimer timerWithTimeInterval:kBackgroundConnectTimer
+                                                              target:self
+                                                            selector:@selector(backgroundConnectTimerExpired)
+                                                            userInfo:nil
+                                                             repeats:NO];
         NSRunLoop *loop = [NSRunLoop mainRunLoop];
         [loop addTimer:[TSSocketManager sharedManager].backgroundConnectTimer forMode:NSDefaultRunLoopMode];
-
+        
-        [sharedInstance.backgroundKeepAliveTimer invalidate];
+        [self.backgroundKeepAliveTimer invalidate];
-        sharedInstance.didConnectBg          = NO;
+        self.backgroundKeepAliveTimer = nil;
-        sharedInstance.didRetreiveMessageBg  = NO;
+        self.didConnectBg          = NO;
-        sharedInstance.shouldDownloadMessage = expected;
+        self.didRetreiveMessageBg  = NO;
-        sharedInstance.fetchingTaskIdentifier =
+        self.shouldDownloadMessage = expected;
-            [[UIApplication sharedApplication] beginBackgroundTaskWithExpirationHandler:^{
+        self.fetchingTaskIdentifier =
-              [TSSocketManager resignActivity];
+        [[UIApplication sharedApplication] beginBackgroundTaskWithExpirationHandler:^{
-              [[TSSocketManager sharedManager] closeBackgroundTask];
+            OWSAssert([NSThread isMainThread]);
-            }];
+            
-
+            [TSSocketManager resignActivity];
+            [[TSSocketManager sharedManager] closeBackgroundTask];
+        }];
+        
         [self becomeActive];
     } else {
         DDLogWarn(@"Got called to become active in the background but there was already a background task running.");
@@ -313,17 +403,28 @@ NSString *const SocketConnectingNotification = @"SocketConnectingNotification";
 }
 
 - (void)backgroundConnectTimerExpired {
+    [self.backgroundConnectTimer invalidate];
+    self.backgroundConnectTimer = nil;
+    
     [self backgroundTimeExpired];
 }
 
 - (void)backgroundTimeExpired {
-    [[self class] resignActivity];
+    OWSAssert([NSThread isMainThread]);
+    
+    if (![self shouldKeepWebSocketAlive]) {
+        [[self class] resignActivity];
+    }
     [self closeBackgroundTask];
 }
 
 - (void)closeBackgroundTask {
+    OWSAssert([NSThread isMainThread]);
+
     [self.backgroundKeepAliveTimer invalidate];
+    self.backgroundKeepAliveTimer = nil;
     [self.backgroundConnectTimer invalidate];
+    self.backgroundConnectTimer = nil;
 
     /*
      If VOIP Push worked, we should just have to check if message was retreived and if not, alert the user.