zwopple · eunikolsky · Nov 8, 2018
diff --git a/PSAutobahnClientTests/PSWebSocketServerTests.m b/PSAutobahnClientTests/PSWebSocketServerTests.m
@@ -0,0 +1,63 @@
+//  Copyright 2014-Present Zwopple Limited
+//
+//  Licensed under the Apache License, Version 2.0 (the "License");
+//  you may not use this file except in compliance with the License.
+//  You may obtain a copy of the License at
+//
+//  http://www.apache.org/licenses/LICENSE-2.0
+//
+//  Unless required by applicable law or agreed to in writing, software
+//  distributed under the License is distributed on an "AS IS" BASIS,
+//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+//  See the License for the specific language governing permissions and
+//  limitations under the License.
+
+#import <XCTest/XCTest.h>
+
+#import "PSWebSocket.h"
+#import "PSWebSocketServer.h"
+
+@interface FakeServerDelegate : NSObject <PSWebSocketServerDelegate> @end
+@implementation FakeServerDelegate
+
+- (void)server:(PSWebSocketServer *)server didFailWithError:(NSError *)error { }
+- (void)server:(PSWebSocketServer *)server webSocket:(PSWebSocket *)webSocket didCloseWithCode:(NSInteger)code reason:(NSString *)reason wasClean:(BOOL)wasClean { }
+- (void)server:(PSWebSocketServer *)server webSocket:(PSWebSocket *)webSocket didFailWithError:(NSError *)error { }
+- (void)server:(PSWebSocketServer *)server webSocket:(PSWebSocket *)webSocket didReceiveMessage:(id)message { }
+- (void)server:(PSWebSocketServer *)server webSocketDidOpen:(PSWebSocket *)webSocket { }
+- (void)serverDidStart:(PSWebSocketServer *)server { }
+- (void)serverDidStop:(PSWebSocketServer *)server { }
+
+@end
+
+@interface FakeServerDelegateWithFlushOutput : FakeServerDelegate @end
+@implementation FakeServerDelegateWithFlushOutput
+
+- (void)server:(PSWebSocketServer *)server webSocketDidFlushOutput:(PSWebSocket *)webSocket {}
+
+@end
+
+// opening the class for testing
+@interface PSWebSocketServer() <PSWebSocketDelegate> @end
+
+@interface PSWebSocketServerTests: XCTestCase @end
+@implementation PSWebSocketServerTests
+
+- (void)testChangingDelegateWhileFlushingOutputShouldNotCrashWithUnrecognizedSelector {
+    FakeServerDelegateWithFlushOutput *delegateWithFlushOutput = [FakeServerDelegateWithFlushOutput new];
+    FakeServerDelegate *delegate = [FakeServerDelegate new];
+
+    PSWebSocketServer *sut = [PSWebSocketServer serverWithHost:@"" port:9999];
+    sut.delegate = delegateWithFlushOutput;
+    sut.delegateQueue = dispatch_get_global_queue(QOS_CLASS_UTILITY, 0);
+
+    PSWebSocket *socket = [PSWebSocket clientSocketWithRequest:[NSURLRequest requestWithURL:[NSURL URLWithString:@"http://127.1/"]]];
+
+    for (uint16_t i = 0; i < 5000; ++i) {
+        sut.delegate = delegateWithFlushOutput;
+        [sut webSocketDidFlushOutput:socket];
+        sut.delegate = delegate;
+    }
+}
+
+@end
diff --git a/PocketSocket.xcodeproj/project.pbxproj b/PocketSocket.xcodeproj/project.pbxproj
@@ -7,6 +7,7 @@
 	objects = {
 
 /* Begin PBXBuildFile section */
+		07C803FF2194AEE40038A4AB /* libPocketSocket.a in Frameworks */ = {isa = PBXBuildFile; fileRef = EEE5E30A18B37DD500BAE47A /* libPocketSocket.a */; };
 		275D3ADB1B02811E0013B9A9 /* PSWebSocketBuffer.m in Sources */ = {isa = PBXBuildFile; fileRef = EEE5E33818B37DEC00BAE47A /* PSWebSocketBuffer.m */; };
 		275D3ADC1B02811E0013B9A9 /* PSWebSocket.m in Sources */ = {isa = PBXBuildFile; fileRef = EEE5E35218B37DF300BAE47A /* PSWebSocket.m */; };
 		275D3ADD1B02811E0013B9A9 /* PSWebSocketUTF8Decoder.m in Sources */ = {isa = PBXBuildFile; fileRef = EEE5E34218B37DEC00BAE47A /* PSWebSocketUTF8Decoder.m */; };
@@ -16,6 +17,7 @@
 		275D3AE11B02811E0013B9A9 /* PSWebSocketNetworkThread.m in Sources */ = {isa = PBXBuildFile; fileRef = EEE5E34018B37DEC00BAE47A /* PSWebSocketNetworkThread.m */; };
 		275D3AE21B02811E0013B9A9 /* PSWebSocketServer.m in Sources */ = {isa = PBXBuildFile; fileRef = EE2A05DA18B5BBEC0066EEA4 /* PSWebSocketServer.m */; };
 		275D3AE41B02811E0013B9A9 /* Foundation.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = EEE5E30D18B37DD500BAE47A /* Foundation.framework */; };
+		AC576A417BC55056BF0FE522 /* PSWebSocketServerTests.m in Sources */ = {isa = PBXBuildFile; fileRef = AC57675323BB4AFDB2A920A8 /* PSWebSocketServerTests.m */; };
 		EE2A05DB18B5BBEC0066EEA4 /* PSWebSocketServer.m in Sources */ = {isa = PBXBuildFile; fileRef = EE2A05DA18B5BBEC0066EEA4 /* PSWebSocketServer.m */; };
 		EE337EC318B4AE1F003F95B9 /* libz.dylib in Frameworks */ = {isa = PBXBuildFile; fileRef = EEE5E38418B385DE00BAE47A /* libz.dylib */; };
 		EE337EC418B4AE23003F95B9 /* Foundation.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = EEE5E30D18B37DD500BAE47A /* Foundation.framework */; };
@@ -83,6 +85,7 @@
 
 /* Begin PBXFileReference section */
 		275D3AE91B02811E0013B9A9 /* libPocketSocket-Mac.a */ = {isa = PBXFileReference; explicitFileType = archive.ar; includeInIndex = 0; path = "libPocketSocket-Mac.a"; sourceTree = BUILT_PRODUCTS_DIR; };
+		AC57675323BB4AFDB2A920A8 /* PSWebSocketServerTests.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = PSWebSocketServerTests.m; sourceTree = "<group>"; };
 		EE2A05D918B5BBEC0066EEA4 /* PSWebSocketServer.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = PSWebSocketServer.h; sourceTree = "<group>"; };
 		EE2A05DA18B5BBEC0066EEA4 /* PSWebSocketServer.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = PSWebSocketServer.m; sourceTree = "<group>"; };
 		EE337EC518B4AE28003F95B9 /* CFNetwork.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = CFNetwork.framework; path = System/Library/Frameworks/CFNetwork.framework; sourceTree = SDKROOT; };
@@ -161,6 +164,7 @@
 			isa = PBXFrameworksBuildPhase;
 			buildActionMask = 2147483647;
 			files = (
+				07C803FF2194AEE40038A4AB /* libPocketSocket.a in Frameworks */,
 				EE337ECC18B4AE46003F95B9 /* libSystem.dylib in Frameworks */,
 				EE337ECA18B4AE34003F95B9 /* Security.framework in Frameworks */,
 				EE337EC618B4AE28003F95B9 /* CFNetwork.framework in Frameworks */,
@@ -282,6 +286,7 @@
 				EEE5E36A18B37F8700BAE47A /* PSAutobahnClientTests.m */,
 				EEE5E37218B37FC000BAE47A /* PSAutobahnClientWebSocketOperation.h */,
 				EEE5E37318B37FC000BAE47A /* PSAutobahnClientWebSocketOperation.m */,
+				AC57675323BB4AFDB2A920A8 /* PSWebSocketServerTests.m */,
 				EEE5E36518B37F8700BAE47A /* Supporting Files */,
 			);
 			path = PSAutobahnClientTests;
@@ -482,6 +487,7 @@
 				EEE5E37A18B380F200BAE47A /* PSWebSocketInflater.m in Sources */,
 				EEE5E37618B380EA00BAE47A /* PSWebSocket.m in Sources */,
 				EEE5E36B18B37F8700BAE47A /* PSAutobahnClientTests.m in Sources */,
+				AC576A417BC55056BF0FE522 /* PSWebSocketServerTests.m in Sources */,
 			);
 			runOnlyForDeploymentPostprocessing = 0;
 		};

diff --git a/PocketSocket/PSWebSocket.m b/PocketSocket/PSWebSocket.m
@@ -313,8 +313,8 @@ - (void)connect {
     if(_secure && _mode == PSWebSocketModeClient) {
 
         __block BOOL customTrustEvaluation = NO;
-        [self executeDelegateAndWait:^{
-            customTrustEvaluation = [_delegate respondsToSelector:@selector(webSocket:evaluateServerTrust:)];
+        [self executeDelegateAndWait:^(id <PSWebSocketDelegate> delegate){
+            customTrustEvaluation = [delegate respondsToSelector:@selector(webSocket:evaluateServerTrust:)];
         }];
 
         NSMutableDictionary *ssl = [NSMutableDictionary dictionary];
@@ -502,7 +502,7 @@ - (void)failWithCode:(NSInteger)code reason:(NSString *)reason {
 }
 - (void)failWithError:(NSError *)error {
     if(error.code == PSWebSocketStatusCodeProtocolError && [error.domain isEqualToString:PSWebSocketErrorDomain]) {
-        [self executeDelegate:^{
+        [self executeDelegate:^(id <PSWebSocketDelegate> delegate){
             _closeCode = error.code;
             _closeReason = error.localizedDescription;
             [self closeWithCode:_closeCode reason:_closeReason];
@@ -551,7 +551,7 @@ - (void)driver:(PSWebSocketDriver *)driver didReceiveMessage:(id)message {
     [self notifyDelegateDidReceiveMessage:message];
 }
 - (void)driver:(PSWebSocketDriver *)driver didReceivePing:(NSData *)ping {
-    [self executeDelegate:^{
+    [self executeDelegate:^(id <PSWebSocketDelegate> delegate){
         [self executeWork:^{
             [driver sendPong:ping];
         }];
@@ -560,7 +560,7 @@ - (void)driver:(PSWebSocketDriver *)driver didReceivePing:(NSData *)ping {
 - (void)driver:(PSWebSocketDriver *)driver didReceivePong:(NSData *)pong {
     void (^handler)(NSData *pong) = [_pingHandlers firstObject];
     if(handler) {
-        [self executeDelegate:^{
+        [self executeDelegate:^(id <PSWebSocketDelegate> delegate){
             handler(pong);
         }];
         [_pingHandlers removeObjectAtIndex:0];
@@ -636,44 +636,44 @@ - (void)stream:(NSStream *)stream handleEvent:(NSStreamEvent)event {
 #pragma mark - Delegation
 
 - (void)notifyDelegateDidOpen {
-    [self executeDelegate:^{
-        [_delegate webSocketDidOpen:self];
+    [self executeDelegate:^(id <PSWebSocketDelegate> delegate){
+        [delegate webSocketDidOpen:self];
     }];
 }
 - (void)notifyDelegateDidReceiveMessage:(id)message {
-    [self executeDelegate:^{
-        [_delegate webSocket:self didReceiveMessage:message];
+    [self executeDelegate:^(id <PSWebSocketDelegate> delegate){
+        [delegate webSocket:self didReceiveMessage:message];
     }];
 }
 - (void)notifyDelegateDidFailWithError:(NSError *)error {
-    [self executeDelegate:^{
-        [_delegate webSocket:self didFailWithError:error];
+    [self executeDelegate:^(id <PSWebSocketDelegate> delegate){
+        [delegate webSocket:self didFailWithError:error];
     }];
 }
 - (void)notifyDelegateDidCloseWithCode:(NSInteger)code reason:(NSString *)reason wasClean:(BOOL)wasClean {
-    [self executeDelegate:^{
-        [_delegate webSocket:self didCloseWithCode:code reason:reason wasClean:wasClean];
+    [self executeDelegate:^(id <PSWebSocketDelegate> delegate){
+        [delegate webSocket:self didCloseWithCode:code reason:reason wasClean:wasClean];
     }];
 }
 - (void)notifyDelegateDidFlushInput {
-    [self executeDelegate:^{
-        if ([_delegate respondsToSelector:@selector(webSocketDidFlushInput:)]) {
-            [_delegate webSocketDidFlushInput:self];
+    [self executeDelegate:^(id <PSWebSocketDelegate> delegate){
+        if ([delegate respondsToSelector:@selector(webSocketDidFlushInput:)]) {
+            [delegate webSocketDidFlushInput:self];
         }
     }];
 }
 - (void)notifyDelegateDidFlushOutput {
-    [self executeDelegate:^{
-        if ([_delegate respondsToSelector:@selector(webSocketDidFlushOutput:)]) {
-            [_delegate webSocketDidFlushOutput:self];
+    [self executeDelegate:^(id <PSWebSocketDelegate> delegate){
+        if ([delegate respondsToSelector:@selector(webSocketDidFlushOutput:)]) {
+            [delegate webSocketDidFlushOutput:self];
         }
     }];
 }
 - (BOOL)askDelegateToEvaluateServerTrust:(SecTrustRef)trust {
     __block BOOL result = NO;
-    [self executeDelegateAndWait:^{
-        if ([_delegate respondsToSelector:@selector(webSocket:evaluateServerTrust:)]) {
-            result = [_delegate webSocket:self evaluateServerTrust:trust];
+    [self executeDelegateAndWait:^(id <PSWebSocketDelegate> delegate){
+        if ([delegate respondsToSelector:@selector(webSocket:evaluateServerTrust:)]) {
+            result = [delegate webSocket:self evaluateServerTrust:trust];
         }
     }];
     return result;
@@ -689,13 +689,21 @@ - (void)executeWorkAndWait:(void (^)(void))work {
     NSParameterAssert(work);
     dispatch_sync(_workQueue, work);
 }
-- (void)executeDelegate:(void (^)(void))work {
+- (void)executeDelegate:(void (^)(id <PSWebSocketDelegate>))work {
     NSParameterAssert(work);
-    dispatch_async((_delegateQueue) ? _delegateQueue : dispatch_get_main_queue(), work);
+    id <PSWebSocketDelegate> delegate = _delegate;
+    void (^workRetainingDelegate)(void) = ^void(void) {
+        work(delegate);
+    };
+    dispatch_async((_delegateQueue) ? _delegateQueue : dispatch_get_main_queue(), workRetainingDelegate);
 }
-- (void)executeDelegateAndWait:(void (^)(void))work {
+- (void)executeDelegateAndWait:(void (^)(id <PSWebSocketDelegate>))work {
     NSParameterAssert(work);
-    dispatch_sync((_delegateQueue) ? _delegateQueue : dispatch_get_main_queue(), work);
+    id <PSWebSocketDelegate> delegate = _delegate;
+    void (^workRetainingDelegate)(void) = ^void(void) {
+        work(delegate);
+    };
+    dispatch_sync((_delegateQueue) ? _delegateQueue : dispatch_get_main_queue(), workRetainingDelegate);
 }
 
 #pragma mark - Dealloc

diff --git a/PocketSocket/PSWebSocketServer.m b/PocketSocket/PSWebSocketServer.m
@@ -580,53 +580,53 @@ - (void)stream:(NSStream *)stream handleEvent:(NSStreamEvent)event {
 #pragma mark - Delegation
 
 - (void)notifyDelegateDidStart {
-    [self executeDelegate:^{
-        [_delegate serverDidStart:self];
+    [self executeDelegate:^(id <PSWebSocketServerDelegate> delegate){
+        [delegate serverDidStart:self];
     }];
 }
 - (void)notifyDelegateFailedToStart:(NSError *)error {
-    [self executeDelegate:^{
-        [_delegate server:self didFailWithError:error];
+    [self executeDelegate:^(id <PSWebSocketServerDelegate> delegate){
+        [delegate server:self didFailWithError:error];
     }];
 }
 - (void)notifyDelegateDidStop {
-    [self executeDelegate:^{
-        [_delegate serverDidStop:self];
+    [self executeDelegate:^(id <PSWebSocketServerDelegate> delegate){
+        [delegate serverDidStop:self];
     }];
 }
 
 - (void)notifyDelegateWebSocketDidOpen:(PSWebSocket *)webSocket {
-    [self executeDelegate:^{
-        [_delegate server:self webSocketDidOpen:webSocket];
+    [self executeDelegate:^(id <PSWebSocketServerDelegate> delegate){
+        [delegate server:self webSocketDidOpen:webSocket];
     }];
 }
 - (void)notifyDelegateWebSocket:(PSWebSocket *)webSocket didReceiveMessage:(id)message {
-    [self executeDelegate:^{
-        [_delegate server:self webSocket:webSocket didReceiveMessage:message];
+    [self executeDelegate:^(id <PSWebSocketServerDelegate> delegate){
+        [delegate server:self webSocket:webSocket didReceiveMessage:message];
     }];
 }
 
 - (void)notifyDelegateWebSocket:(PSWebSocket *)webSocket didFailWithError:(NSError *)error {
-    [self executeDelegate:^{
-        [_delegate server:self webSocket:webSocket didFailWithError:error];
+    [self executeDelegate:^(id <PSWebSocketServerDelegate> delegate){
+        [delegate server:self webSocket:webSocket didFailWithError:error];
     }];
 }
 - (void)notifyDelegateWebSocket:(PSWebSocket *)webSocket didCloseWithCode:(NSInteger)code reason:(NSString *)reason wasClean:(BOOL)wasClean {
-    [self executeDelegate:^{
-        [_delegate server:self webSocket:webSocket didCloseWithCode:code reason:reason wasClean:wasClean];
+    [self executeDelegate:^(id <PSWebSocketServerDelegate> delegate){
+        [delegate server:self webSocket:webSocket didCloseWithCode:code reason:reason wasClean:wasClean];
     }];
 }
 - (void)notifyDelegateWebSocketDidFlushInput:(PSWebSocket *)webSocket {
-    [self executeDelegate:^{
-        if ([_delegate respondsToSelector: @selector(server:webSocketDidFlushInput:)]) {
-            [_delegate server:self webSocketDidFlushInput:webSocket];
+    [self executeDelegate:^(id <PSWebSocketServerDelegate> delegate){
+        if ([delegate respondsToSelector: @selector(server:webSocketDidFlushInput:)]) {
+            [delegate server:self webSocketDidFlushInput:webSocket];
         };
     }];
 }
 - (void)notifyDelegateWebSocketDidFlushOutput:(PSWebSocket *)webSocket {
-    [self executeDelegate:^{
-        if ([_delegate respondsToSelector: @selector(server:webSocketDidFlushOutput:)]) {
-            [_delegate server:self webSocketDidFlushOutput:webSocket];
+    [self executeDelegate:^(id <PSWebSocketServerDelegate> delegate){
+        if ([delegate respondsToSelector: @selector(server:webSocketDidFlushOutput:)]) {
+            [delegate server:self webSocketDidFlushOutput:webSocket];
         }
     }];
 }
@@ -635,22 +635,22 @@ - (BOOL)askDelegateShouldAcceptConnection:(PSWebSocketServerConnection *)connect
                                  response:(NSHTTPURLResponse **)outResponse {
     __block BOOL accept;
     __block NSHTTPURLResponse* response = nil;
-    [self executeDelegateAndWait:^{
-        if([_delegate respondsToSelector:@selector(server:acceptWebSocketWithRequest:address:trust:response:)]) {
+    [self executeDelegateAndWait:^(id <PSWebSocketServerDelegate> delegate){
+        if([delegate respondsToSelector:@selector(server:acceptWebSocketWithRequest:address:trust:response:)]) {
             NSData* address = PSPeerAddressOfInputStream(connection.inputStream);
             SecTrustRef trust = (SecTrustRef)CFReadStreamCopyProperty(
                                                   (__bridge CFReadStreamRef)connection.inputStream,
                                                   kCFStreamPropertySSLPeerTrust);
-            accept = [_delegate server:self
+            accept = [delegate server:self
             acceptWebSocketWithRequest:request
                                address:address
                                  trust:trust
                               response:&response];
             if(trust) {
                 CFRelease(trust);
             }
-        } else if([_delegate respondsToSelector:@selector(server:acceptWebSocketWithRequest:)]) {
-            accept = [_delegate server:self acceptWebSocketWithRequest:request];
+        } else if([delegate respondsToSelector:@selector(server:acceptWebSocketWithRequest:)]) {
+            accept = [delegate server:self acceptWebSocketWithRequest:request];
         } else {
             accept = YES;
         }
@@ -669,13 +669,24 @@ - (void)executeWorkAndWait:(void (^)(void))work {
     NSParameterAssert(work);
     dispatch_sync(_workQueue, work);
 }
-- (void)executeDelegate:(void (^)(void))work {
+- (void)executeDelegate:(void (^)(id <PSWebSocketServerDelegate>))work {
     NSParameterAssert(work);
-    dispatch_async((_delegateQueue) ? _delegateQueue : dispatch_get_main_queue(), work);
-}
-- (void)executeDelegateAndWait:(void (^)(void))work {
+    /// The block below captures the `delegate` strongly, so even if the weak `_delegate`
+    /// changes while the block is executing, it won't crash with
+    /// "unrecognized selector sent to instance"
+    id <PSWebSocketServerDelegate> delegate = _delegate;
+    void (^workRetainingDelegate)(void) = ^void(void) {
+        work(delegate);
+    };
+    dispatch_async((_delegateQueue) ? _delegateQueue : dispatch_get_main_queue(), workRetainingDelegate);
+}
+- (void)executeDelegateAndWait:(void (^)(id <PSWebSocketServerDelegate>))work {
     NSParameterAssert(work);
-    dispatch_sync((_delegateQueue) ? _delegateQueue : dispatch_get_main_queue(), work);
+    id <PSWebSocketServerDelegate> delegate = _delegate;
+    void (^workRetainingDelegate)(void) = ^void(void) {
+        work(delegate);
+    };
+    dispatch_sync((_delegateQueue) ? _delegateQueue : dispatch_get_main_queue(), workRetainingDelegate);
 }
 
 #pragma mark - Dealloc