diff --git a/SocketRocket/SRWebSocket.h b/SocketRocket/SRWebSocket.h index 34b7d4360..a85e89bd0 100644 --- a/SocketRocket/SRWebSocket.h +++ b/SocketRocket/SRWebSocket.h @@ -59,6 +59,8 @@ extern NSString *const SRHTTPResponseErrorKey; @protocol SRWebSocketDelegate; +typedef void (^SRSendCompletionBlock)(NSError * _Nullable error); + ///-------------------------------------- #pragma mark - SRWebSocket ///-------------------------------------- @@ -278,6 +280,17 @@ extern NSString *const SRHTTPResponseErrorKey; */ - (BOOL)sendString:(NSString *)string error:(NSError **)error NS_SWIFT_NAME(send(string:)); +/** + Send a UTF-8 String to the server. + + @param string String to send. + @param completion The call back of send result. + If an error occurs, this block will invoked with an `NSError` object containing information about the error, otherwise this block will be invoked with `nil`. + + @return `YES` if the string was scheduled to send, otherwise - `NO`. + */ +- (BOOL)sendString:(NSString *)string completion:(nullable SRSendCompletionBlock)completion NS_SWIFT_NAME(send(string:completion:)); + /** Send binary data to the server. @@ -302,6 +315,17 @@ extern NSString *const SRHTTPResponseErrorKey; */ - (BOOL)sendDataNoCopy:(nullable NSData *)data error:(NSError **)error NS_SWIFT_NAME(send(dataNoCopy:)); +/** + Send binary data to the server, without making a defensive copy of it first. + + @param data Data to send. + @param completion The call back of send result. + If an error occurs, this block will invoked with an `NSError` object containing information about the error, otherwise this block will be invoked with `nil`. + + @return `YES` if the string was scheduled to send, otherwise - `NO`. + */ +- (BOOL)sendDataNoCopy:(nullable NSData *)data completion:(nullable SRSendCompletionBlock)completion NS_SWIFT_NAME(send(dataNoCopy:completion:)); + /** Send Ping message to the server with optional data. diff --git a/SocketRocket/SRWebSocket.m b/SocketRocket/SRWebSocket.m index 83f3e128f..6bd274b21 100644 --- a/SocketRocket/SRWebSocket.m +++ b/SocketRocket/SRWebSocket.m @@ -69,6 +69,31 @@ NSString *const SRWebSocketErrorDomain = @"SRWebSocketErrorDomain"; NSString *const SRHTTPResponseErrorKey = @"HTTPResponseStatusCode"; +@interface SRDataCallback : NSObject +@property (nonatomic, assign) NSRange range; +@property (nonatomic, copy, readonly) SRSendCompletionBlock completion; + ++ (instancetype)new NS_UNAVAILABLE; +- (instancetype)init NS_UNAVAILABLE; +- (instancetype)initWithRange:(NSRange)range + completion:(SRSendCompletionBlock)completion NS_DESIGNATED_INITIALIZER; + +@end + +@implementation SRDataCallback + +- (instancetype)initWithRange:(NSRange)range completion:(SRSendCompletionBlock)completion +{ + self = [super init]; + + _range = range; + _completion = [completion copy]; + + return self; +} + +@end + @interface SRWebSocket () @property (atomic, assign, readwrite) SRReadyState readyState; @@ -138,6 +163,8 @@ @implementation SRWebSocket { // proxy support SRProxyConnect *_proxyConnect; + + NSMutableDictionary *_sendCallbacks; } @synthesize readyState = _readyState; @@ -179,6 +206,8 @@ - (instancetype)initWithURLRequest:(NSURLRequest *)request protocols:(NSArray _Nullable delegate, SRDelegateAvailableMethods availableMethods) { @@ -566,14 +600,30 @@ - (void)_failWithError:(NSError *)error; }); } -- (void)_writeData:(NSData *)data; +- (void)_writeData:(NSData *)data { - [self assertOnWorkQueue]; + [self _writeData:data completion:nil]; +} +- (void)_writeData:(NSData *)data completion:(nullable SRSendCompletionBlock)completion +{ + [self assertOnWorkQueue]; + if (_closeWhenFinishedWriting) { + if (completion) { + completion(SRErrorWithCodeDescription(2134, @"socket is closed")); + } return; } + + if (completion) { + NSUInteger location = dispatch_data_get_size(_outputBuffer); + NSRange dataRange = NSMakeRange(location, data.length); + SRDataCallback *record = [[SRDataCallback alloc] initWithRange:dataRange completion:completion]; + _sendCallbacks[[NSValue valueWithRange:dataRange]] = record; + } + __block NSData *strongData = data; dispatch_data_t newData = dispatch_data_create(data.bytes, data.length, nil, ^{ strongData = nil; @@ -613,6 +663,24 @@ - (BOOL)sendString:(NSString *)string error:(NSError **)error return YES; } +- (BOOL)sendString:(NSString *)string completion:(nullable SRSendCompletionBlock)completion +{ + if (self.readyState != SR_OPEN) { + NSString *message = @"Invalid State: Cannot call `sendString:completion:` until connection is open."; + if (completion) { + completion(SRErrorWithCodeDescription(2134, message)); + } + SRDebugLog(message); + return NO; + } + + string = [string copy]; + dispatch_async(_workQueue, ^{ + [self _sendFrameWithOpcode:SROpCodeTextFrame data:[string dataUsingEncoding:NSUTF8StringEncoding] completion:completion]; + }); + return YES; +} + - (BOOL)sendData:(nullable NSData *)data error:(NSError **)error { data = [data copy]; @@ -640,6 +708,27 @@ - (BOOL)sendDataNoCopy:(nullable NSData *)data error:(NSError **)error return YES; } +- (BOOL)sendDataNoCopy:(nullable NSData *)data completion:(nullable SRSendCompletionBlock)completion +{ + if (self.readyState != SR_OPEN) { + NSString *message = @"Invalid State: Cannot call `sendDataNoCopy:completion:` until connection is open."; + if (completion) { + completion(SRErrorWithCodeDescription(2134, message)); + } + SRDebugLog(message); + return NO; + } + + dispatch_async(_workQueue, ^{ + if (data) { + [self _sendFrameWithOpcode:SROpCodeBinaryFrame data:data completion:completion]; + } else { + [self _sendFrameWithOpcode:SROpCodeTextFrame data:nil completion:completion]; + } + }); + return YES; +} + - (BOOL)sendPing:(nullable NSData *)data error:(NSError **)error { if (self.readyState != SR_OPEN) { @@ -1060,7 +1149,27 @@ - (void)_pumpWriting; _outputBufferOffset += bytesWritten; + NSMutableArray *removeKeys = [NSMutableArray array]; + [_sendCallbacks enumerateKeysAndObjectsUsingBlock:^(NSValue * _Nonnull key, SRDataCallback * _Nonnull obj, BOOL * _Nonnull stop) { + if (NSMaxRange(obj.range) <= _outputBufferOffset) { + [removeKeys addObject:key]; + obj.completion(nil); + } + }]; + [_sendCallbacks removeObjectsForKeys:removeKeys]; + if (_outputBufferOffset > SRDefaultBufferSize() && _outputBufferOffset > dataLength / 2) { + + NSArray *callbacks = _sendCallbacks.allValues; + [_sendCallbacks removeAllObjects]; + [callbacks enumerateObjectsUsingBlock:^(SRDataCallback * _Nonnull obj, NSUInteger idx, BOOL * _Nonnull stop) { + NSRange dataRange = obj.range; + dataRange.location -= _outputBufferOffset; + obj.range = dataRange; + _sendCallbacks[[NSValue valueWithRange:dataRange]] = obj; + }]; + + _outputBuffer = dispatch_data_create_subrange(_outputBuffer, _outputBufferOffset, dataLength - _outputBufferOffset); _outputBufferOffset = 0; } @@ -1320,11 +1429,14 @@ -(void)_pumpScanner; static const size_t SRFrameHeaderOverhead = 32; -- (void)_sendFrameWithOpcode:(SROpCode)opCode data:(NSData *)data +- (void)_sendFrameWithOpcode:(SROpCode)opCode data:(NSData *)data completion:(nullable SRSendCompletionBlock)completion { [self assertOnWorkQueue]; if (!data) { + if (completion) { + completion(nil); + } return; } @@ -1332,7 +1444,11 @@ - (void)_sendFrameWithOpcode:(SROpCode)opCode data:(NSData *)data NSMutableData *frameData = [[NSMutableData alloc] initWithLength:payloadLength + SRFrameHeaderOverhead]; if (!frameData) { - [self closeWithCode:SRStatusCodeMessageTooBig reason:@"Message too big"]; + NSString *reason = @"Message too big"; + [self closeWithCode:SRStatusCodeMessageTooBig reason:reason]; + if (completion) { + completion(SRErrorWithCodeDescription(SRStatusCodeMessageTooBig, reason)); + } return; } uint8_t *frameBuffer = (uint8_t *)frameData.mutableBytes; @@ -1387,7 +1503,12 @@ - (void)_sendFrameWithOpcode:(SROpCode)opCode data:(NSData *)data assert(frameBufferSize <= frameData.length); frameData.length = frameBufferSize; - [self _writeData:frameData]; + [self _writeData:frameData completion:completion]; +} + +- (void)_sendFrameWithOpcode:(SROpCode)opCode data:(NSData *)data +{ + [self _sendFrameWithOpcode:opCode data:data completion:nil]; } - (void)stream:(NSStream *)aStream handleEvent:(NSStreamEvent)eventCode