Merge pull request #169 from jleandroperez/master

Issue #156: Prevents crash in handleEvent method
facebookincubator · Feb 19, 2016 · 541f7b6 · 541f7b6
2 parents 4c0d934 + a770a5f
commit 541f7b6
Showing 1 changed file with 58 additions and 12 deletions.
diff --git a/SocketRocket/SRWebSocket.m b/SocketRocket/SRWebSocket.m
@@ -236,6 +236,7 @@ @implementation SRWebSocket {
 
     BOOL _sentClose;
     BOOL _didFail;
+    BOOL _cleanupScheduled;
     int _closeCode;
 
     BOOL _isPumping;
@@ -762,11 +763,11 @@ - (void)_failWithError:(NSError *)error;
             }];
 
             self.readyState = SR_CLOSED;
-            _selfRetain = nil;
 
             SRFastLog(@"Failing with error %@", error.localizedDescription);
 
             [self closeConnection];
+            [self _scheduleCleanup];
         }
     });
 }
@@ -1174,13 +1175,15 @@ - (void)_pumpWriting;
          _inputStream.streamStatus != NSStreamStatusClosed) &&
         !_sentClose) {
         _sentClose = YES;
-
-        [_outputStream close];
-        [_inputStream close];
 
-
-        for (NSArray *runLoop in [_scheduledRunloops copy]) {
-            [self unscheduleFromRunLoop:[runLoop objectAtIndex:0] forMode:[runLoop objectAtIndex:1]];
+        @synchronized(self) {
+            [_outputStream close];
+            [_inputStream close];
+
+
+            for (NSArray *runLoop in [_scheduledRunloops copy]) {
+                [self unscheduleFromRunLoop:[runLoop objectAtIndex:0] forMode:[runLoop objectAtIndex:1]];
+            }
         }
 
         if (!_failed) {
@@ -1191,7 +1194,7 @@ - (void)_pumpWriting;
             }];
         }
 
-        _selfRetain = nil;
+        [self _scheduleCleanup];
     }
 }
 
@@ -1218,6 +1221,41 @@ - (void)_addConsumerWithScanner:(stream_scanner)consumer callback:(data_callback
 }
 
 
+- (void)_scheduleCleanup
+{
+    @synchronized(self) {
+        if (_cleanupScheduled) {
+            return;
+        }
+
+        _cleanupScheduled = YES;
+
+        // Cleanup NSStream delegate's in the same RunLoop used by the streams themselves:
+        // This way we'll prevent race conditions between handleEvent and SRWebsocket's dealloc
+        NSTimer *timer = [NSTimer timerWithTimeInterval:(0.0f) target:self selector:@selector(_cleanupSelfReference:) userInfo:nil repeats:NO];
+        [[NSRunLoop SR_networkRunLoop] addTimer:timer forMode:NSDefaultRunLoopMode];
+    }
+}
+
+- (void)_cleanupSelfReference:(NSTimer *)timer
+{
+    @synchronized(self) {
+        // Nuke NSStream delegate's
+        _inputStream.delegate = nil;
+        _outputStream.delegate = nil;
+
+        // Remove the streams, right now, from the networkRunLoop
+        [_inputStream close];
+        [_outputStream close];
+    }
+
+    // Cleanup selfRetain in the same GCD queue as usual
+    dispatch_async(_workQueue, ^{
+        _selfRetain = nil;
+    });
+}
+
+
 static const char CRLFCRLFBytes[] = {'\r', '\n', '\r', '\n'};
 
 - (void)_readUntilHeaderCompleteWithCallback:(data_callback)dataHandler;
@@ -1459,6 +1497,8 @@ - (void)_sendFrameWithOpcode:(SROpCode)opcode data:(id)data;
 
 - (void)stream:(NSStream *)aStream handleEvent:(NSStreamEvent)eventCode;
 {
+    __weak typeof(self) weakSelf = self;
+
     if (_secure && !_pinnedCertFound && (eventCode == NSStreamEventHasBytesAvailable || eventCode == NSStreamEventHasSpaceAvailable)) {
 
         NSArray *sslCerts = [_urlRequest SR_SSLPinnedCertificates];
@@ -1484,7 +1524,8 @@ - (void)stream:(NSStream *)aStream handleEvent:(NSStreamEvent)eventCode;
 
             if (!_pinnedCertFound) {
                 dispatch_async(_workQueue, ^{
-                    [self _failWithError:[NSError errorWithDomain:SRWebSocketErrorDomain code:23556 userInfo:[NSDictionary dictionaryWithObject:[NSString stringWithFormat:@"Invalid server cert"] forKey:NSLocalizedDescriptionKey]]];
+                    NSDictionary *userInfo = @{ NSLocalizedDescriptionKey : @"Invalid server cert" };
+                    [weakSelf _failWithError:[NSError errorWithDomain:@"org.lolrus.SocketRocket" code:23556 userInfo:userInfo]];
                 });
                 return;
             } else if (aStream == _outputStream) {
@@ -1496,6 +1537,12 @@ - (void)stream:(NSStream *)aStream handleEvent:(NSStreamEvent)eventCode;
     }
 
     dispatch_async(_workQueue, ^{
+        [weakSelf safeHandleEvent:eventCode stream:aStream];
+    });
+}
+
+- (void)safeHandleEvent:(NSStreamEvent)eventCode stream:(NSStream *)aStream
+{
         switch (eventCode) {
             case NSStreamEventOpenCompleted: {
                 SRFastLog(@"NSStreamEventOpenCompleted %@", aStream);
@@ -1533,9 +1580,9 @@ - (void)stream:(NSStream *)aStream handleEvent:(NSStreamEvent)eventCode;
                     dispatch_async(_workQueue, ^{
                         if (self.readyState != SR_CLOSED) {
                             self.readyState = SR_CLOSED;
-                            _selfRetain = nil;
+                            [self _scheduleCleanup];
                         }
-
+                        
                         if (!_sentClose && !_failed) {
                             _sentClose = YES;
                             // If we get closed in this state it's probably not clean because we should be sending this when we send messages
@@ -1583,7 +1630,6 @@ - (void)stream:(NSStream *)aStream handleEvent:(NSStreamEvent)eventCode;
                 SRFastLog(@"(default)  %@", aStream);
                 break;
         }
-    });
 }
 
 @end