Merge pull request facebookincubator#414 from facebook/nlutsenko.prox…

…y.crash Fix crash on SRProxyConnect deallocation, when timeout was encountered.
wulongwang · Jun 17, 2016 · aaab483 · aaab483
2 parents e17d8ce + 14d35ac
commit aaab483
Showing 1 changed file with 22 additions and 4 deletions.
diff --git a/SocketRocket/Internal/Proxy/SRProxyConnect.m b/SocketRocket/Internal/Proxy/SRProxyConnect.m
@@ -39,7 +39,7 @@ @implementation SRProxyConnect
     BOOL _connectionRequiresSSL;
 
     NSMutableArray<NSData *> *_inputQueue;
-    NSOperationQueue * _writeQueue;
+    dispatch_queue_t _writeQueue;
 }
 
 ///--------------------------------------
@@ -54,12 +54,26 @@ -(instancetype)initWithURL:(NSURL *)url
     _url = url;
     _connectionRequiresSSL = SRURLRequiresSSL(url);
 
-    _writeQueue =  [[NSOperationQueue alloc] init];
+    _writeQueue = dispatch_queue_create("com.facebook.socketrocket.proxyconnect.write", DISPATCH_QUEUE_SERIAL);
     _inputQueue = [NSMutableArray arrayWithCapacity:2];
 
     return self;
 }
 
+- (void)dealloc
+{
+    // If we get deallocated before the socket open finishes - we need to cleanup everything.
+
+    [self.inputStream removeFromRunLoop:[NSRunLoop SR_networkRunLoop] forMode:NSDefaultRunLoopMode];
+    self.inputStream.delegate = nil;
+    [self.inputStream close];
+    self.inputStream = nil;
+
+    self.outputStream.delegate = nil;
+    [self.outputStream close];
+    self.outputStream = nil;
+}
+
 ///--------------------------------------
 #pragma mark - Open
 ///--------------------------------------
@@ -113,6 +127,10 @@ - (void)_failWithError:(NSError *)error
         CFRelease(_receivedHTTPHeaders);
         _receivedHTTPHeaders = NULL;
     }
+
+    self.inputStream.delegate = nil;
+    self.outputStream.delegate = nil;
+
     [self.inputStream removeFromRunLoop:[NSRunLoop SR_networkRunLoop]
                                 forMode:NSDefaultRunLoopMode];
     [self.inputStream close];
@@ -433,7 +451,7 @@ - (void)_writeData:(NSData *)data
     const uint8_t * bytes = data.bytes;
     __block NSInteger timeout = SRProxyConnectWriteTimeout * 1000000; // wait timeout before giving up
     __weak typeof(self) wself = self;
-    [_writeQueue addOperationWithBlock:^() {
+    dispatch_async(_writeQueue, ^{
         if (!wself) {
             return;
         }
@@ -452,7 +470,7 @@ - (void)_writeData:(NSData *)data
             }
         }
         [outStream write:bytes maxLength:data.length];
-    }];
+    });
 }
 @end