//

// Copyright 2012 Square Inc.

//

// 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 "SRWebSocket.h"

#if TARGET_OS_IPHONE

#define HAS_ICU

#endif

#ifdef HAS_ICU

#import <unicode/utf8.h>

#endif

#if TARGET_OS_IPHONE

#import <Endian.h>

#else

#import <CoreServices/CoreServices.h>

#endif

#import <CommonCrypto/CommonDigest.h>

#import <Security/SecRandom.h>

#import "base64.h"

#import "NSData+SRB64Additions.h"

typedef enum {

SROpCodeTextFrame = 0x1,

SROpCodeBinaryFrame = 0x2,

//3-7Reserved

SROpCodeConnectionClose = 0x8,

SROpCodePing = 0x9,

SROpCodePong = 0xA,

//B-F reserved

} SROpCode;

typedef enum {

SRStatusCodeNormal = 1000,

SRStatusCodeGoingAway = 1001,

SRStatusCodeProtocolError = 1002,

SRStatusCodeUnhandledType = 1003,

// 1004 reserved

SRStatusNoStatusReceived = 1005,

// 1004-1006 reserved

SRStatusCodeInvalidUTF8 = 1007,

SRStatusCodePolicyViolated = 1008,

SRStatusCodeMessageTooBig = 1009,

} SRStatusCode;

typedef struct {

BOOL fin;

// BOOL rsv1;

// BOOL rsv2;

// BOOL rsv3;

uint8_t opcode;

BOOL masked;

uint64_t payload_length;

} frame_header;

static NSString *const SRWebSocketAppendToSecKeyString = @"258EAFA5-E914-47DA-95CA-C5AB0DC85B11";

static inline int32_t validate_dispatch_data_partial_string(NSData *data);

static inline dispatch_queue_t log_queue();

static inline void SRFastLog(NSString *format, ...);

@interface NSData (SRWebSocket)

- (NSString *)stringBySHA1ThenBase64Encoding;

@end

@interface NSString (SRWebSocket)

- (NSString *)stringBySHA1ThenBase64Encoding;

@end

@interface NSURL (SRWebSocket)

// The origin isn't really applicable for a native application

// So instead, just map ws -> http and wss -> https

- (NSString *)SR_origin;

@end

@interface _SRRunLoopThread : NSThread

@property (nonatomic, readonly) NSRunLoop *runLoop;

@end

static NSString *newSHA1String(const char *bytes, size_t length) {

uint8_t md[CC_SHA1_DIGEST_LENGTH];

CC_SHA1(bytes, length, md);

size_t buffer_size = ((sizeof(md) * 3 + 2) / 2);

char *buffer = (char *)malloc(buffer_size);

int len = b64_ntop(md, CC_SHA1_DIGEST_LENGTH, buffer, buffer_size);

if (len == -1) {

free(buffer);

return nil;

} else{

return [[NSString alloc] initWithBytesNoCopy:buffer length:len encoding:NSASCIIStringEncoding freeWhenDone:YES];

}

}

@implementation NSData (SRWebSocket)

- (NSString *)stringBySHA1ThenBase64Encoding;

{

return newSHA1String(self.bytes, self.length);

}

@end

@implementation NSString (SRWebSocket)

- (NSString *)stringBySHA1ThenBase64Encoding;

{

return newSHA1String(self.UTF8String, self.length);

}

@end

NSString *const SRWebSocketErrorDomain = @"SRWebSocketErrorDomain";

// Returns number of bytes consumed. returning 0 means you didn't match.

// Sends bytes to callback handler;

typedef size_t (^stream_scanner)(NSData *collected_data);

typedef void (^data_callback)(SRWebSocket *webSocket, NSData *data);

@interface SRIOConsumer : NSObject {

stream_scanner _scanner;

data_callback _handler;

size_t _bytesNeeded;

BOOL _readToCurrentFrame;

BOOL _unmaskBytes;

}

@property (nonatomic, copy, readonly) stream_scanner consumer;

@property (nonatomic, copy, readonly) data_callback handler;

@property (nonatomic, assign) size_t bytesNeeded;

@property (nonatomic, assign, readonly) BOOL readToCurrentFrame;

@property (nonatomic, assign, readonly) BOOL unmaskBytes;

@end

// This class is not thread-safe, and is expected to always be run on the same queue

@interface SRIOConsumerPool : NSObject

- (id)initWithBufferCapacity:(NSUInteger)poolSize;

- (SRIOConsumer *)consumerWithScanner:(stream_scanner)scanner handler:(data_callback)handler bytesNeeded:(size_t)bytesNeeded readToCurrentFrame:(BOOL)readToCurrentFrame unmaskBytes:(BOOL)unmaskBytes;

- (void)returnConsumer:(SRIOConsumer *)consumer;

@end

@interface SRWebSocket () <NSStreamDelegate>

- (void)_writeData:(NSData *)data;

- (void)_closeWithProtocolError:(NSString *)message;

- (void)_failWithError:(NSError *)error;

- (void)_disconnect;

- (void)_readFrameNew;

- (void)_readFrameContinue;

- (void)_pumpScanner;

- (void)_pumpWriting;

- (void)_addConsumerWithScanner:(stream_scanner)consumer callback:(data_callback)callback;

- (void)_addConsumerWithDataLength:(size_t)dataLength callback:(data_callback)callback readToCurrentFrame:(BOOL)readToCurrentFrame unmaskBytes:(BOOL)unmaskBytes;

- (void)_addConsumerWithScanner:(stream_scanner)consumer callback:(data_callback)callback dataLength:(size_t)dataLength;

- (void)_readUntilBytes:(const void *)bytes length:(size_t)length callback:(data_callback)dataHandler;

- (void)_readUntilHeaderCompleteWithCallback:(data_callback)dataHandler;

- (void)_sendFrameWithOpcode:(SROpCode)opcode data:(id)data;

- (BOOL)_checkHandshake:(CFHTTPMessageRef)httpMessage;

- (void)_SR_commonInit;

- (void)_initializeStreams;

- (void)_connect;

@property (nonatomic) SRReadyState readyState;

@property (nonatomic) NSOperationQueue *delegateOperationQueue;

@property (nonatomic) dispatch_queue_t delegateDispatchQueue;

@end

@implementation SRWebSocket {

NSInteger _webSocketVersion;

NSOperationQueue *_delegateOperationQueue;

dispatch_queue_t _delegateDispatchQueue;

dispatch_queue_t _workQueue;

NSMutableArray *_consumers;

NSInputStream *_inputStream;

NSOutputStream *_outputStream;

NSMutableData *_readBuffer;

NSUInteger _readBufferOffset;

NSMutableData *_outputBuffer;

NSUInteger _outputBufferOffset;

uint8_t _currentFrameOpcode;

size_t _currentFrameCount;

size_t _readOpCount;

uint32_t _currentStringScanPosition;

NSMutableData *_currentFrameData;

NSString *_closeReason;

NSString *_secKey;

BOOL _pinnedCertFound;

uint8_t _currentReadMaskKey[4];

size_t _currentReadMaskOffset;

BOOL _consumerStopped;

BOOL _closeWhenFinishedWriting;

BOOL _failed;

BOOL _secure;

NSURLRequest *_urlRequest;

CFHTTPMessageRef _receivedHTTPHeaders;

BOOL _sentClose;

BOOL _didFail;

int _closeCode;

BOOL _isPumping;

BOOL _didSchedule;

// We use this to retain ourselves.

__strong SRWebSocket *_selfRetain;

NSArray *_requestedProtocols;

SRIOConsumerPool *_consumerPool;

}

@synthesize delegate = _delegate;

@synthesize url = _url;

@synthesize readyState = _readyState;

@synthesize protocol = _protocol;

static __strong NSData *CRLFCRLF;

+ (void)initialize;

{

CRLFCRLF = [[NSData alloc] initWithBytes:"\r\n\r\n" length:4];

}

- (id)initWithURLRequest:(NSURLRequest *)request protocols:(NSArray *)protocols;

{

self = [super init];

if (self) {

assert(request.URL);

_url = request.URL;

_urlRequest = request;

_requestedProtocols = [protocols copy];

[self _SR_commonInit];

}

return self;

}

- (id)initWithURLRequest:(NSURLRequest *)request;

{

return [self initWithURLRequest:request protocols:nil];

}

- (id)initWithURL:(NSURL *)url;

{

return [self initWithURL:url protocols:nil];

}

- (id)initWithURL:(NSURL *)url protocols:(NSArray *)protocols;

{

NSMutableURLRequest *request = [[NSMutableURLRequest alloc] initWithURL:url];

return [self initWithURLRequest:request protocols:protocols];

}

- (void)_SR_commonInit;

{

NSString *scheme = _url.scheme.lowercaseString;

assert([scheme isEqualToString:@"ws"] || [scheme isEqualToString:@"http"] || [scheme isEqualToString:@"wss"] || [scheme isEqualToString:@"https"]);

if ([scheme isEqualToString:@"wss"] || [scheme isEqualToString:@"https"]) {

_secure = YES;

}

_readyState = SR_CONNECTING;

_consumerStopped = YES;

_webSocketVersion = 13;

_workQueue = dispatch_queue_create(NULL, DISPATCH_QUEUE_SERIAL);

_delegateDispatchQueue = dispatch_get_main_queue();

dispatch_retain(_delegateDispatchQueue);

_readBuffer = [[NSMutableData alloc] init];

_outputBuffer = [[NSMutableData alloc] init];

_currentFrameData = [[NSMutableData alloc] init];

_consumers = [[NSMutableArray alloc] init];

_consumerPool = [[SRIOConsumerPool alloc] init];

[self _initializeStreams];

// default handlers

}

- (void)dealloc

{

_inputStream.delegate = nil;

_outputStream.delegate = nil;

[_inputStream close];

[_outputStream close];

dispatch_release(_workQueue);

if (_receivedHTTPHeaders) {

CFRelease(_receivedHTTPHeaders);

_receivedHTTPHeaders = NULL;

}

if (_delegateDispatchQueue) {

dispatch_release(_delegateDispatchQueue);

_delegateDispatchQueue = NULL;

}

}

#ifndef NDEBUG

- (void)setReadyState:(SRReadyState)aReadyState;

{

[self willChangeValueForKey:@"readyState"];

assert(aReadyState > _readyState);

_readyState = aReadyState;

[self didChangeValueForKey:@"readyState"];

}

#endif

- (void)open;

{

assert(_url);

NSAssert(_readyState == SR_CONNECTING, @"Cannot call -(void)open on SRWebSocket more than once");

_selfRetain = self;

[self _connect];

}

// Calls block on delegate queue

- (void)_performDelegateBlock:(dispatch_block_t)block;

{

if (_delegateOperationQueue) {

[_delegateOperationQueue addOperationWithBlock:block];

} else {

assert(_delegateDispatchQueue);

dispatch_async(_delegateDispatchQueue, block);

}

}

- (void)setDelegateDispatchQueue:(dispatch_queue_t)queue;

{

if (queue) {

dispatch_retain(queue);

}

if (_delegateDispatchQueue) {

dispatch_release(_delegateDispatchQueue);

}

_delegateDispatchQueue = queue;

}

- (BOOL)_checkHandshake:(CFHTTPMessageRef)httpMessage;

{

NSString *acceptHeader = CFBridgingRelease(CFHTTPMessageCopyHeaderFieldValue(httpMessage, CFSTR("Sec-WebSocket-Accept")));

if (acceptHeader == nil) {

return NO;

}

NSString *concattedString = [_secKey stringByAppendingString:SRWebSocketAppendToSecKeyString];

NSString *expectedAccept = [concattedString stringBySHA1ThenBase64Encoding];

return [acceptHeader isEqualToString:expectedAccept];

}

- (void)_HTTPHeadersDidFinish;

{

NSInteger responseCode = CFHTTPMessageGetResponseStatusCode(_receivedHTTPHeaders);

if (responseCode >= 400) {

SRFastLog(@"Request failed with response code %d", responseCode);

[self _failWithError:[NSError errorWithDomain:@"org.lolrus.SocketRocket" code:2132 userInfo:[NSDictionary dictionaryWithObject:[NSString stringWithFormat:@"received bad response code from server %d", responseCode] forKey:NSLocalizedDescriptionKey]]];

return;

}

if(![self _checkHandshake:_receivedHTTPHeaders]) {

[self _failWithError:[NSError errorWithDomain:SRWebSocketErrorDomain code:2133 userInfo:[NSDictionary dictionaryWithObject:[NSString stringWithFormat:@"Invalid Sec-WebSocket-Accept response"] forKey:NSLocalizedDescriptionKey]]];

return;

}

NSString *negotiatedProtocol = CFBridgingRelease(CFHTTPMessageCopyHeaderFieldValue(_receivedHTTPHeaders, CFSTR("Sec-WebSocket-Protocol")));

if (negotiatedProtocol) {

// Make sure we requested the protocol

if ([_requestedProtocols indexOfObject:negotiatedProtocol] == NSNotFound) {

[self _failWithError:[NSError errorWithDomain:SRWebSocketErrorDomain code:2133 userInfo:[NSDictionary dictionaryWithObject:[NSString stringWithFormat:@"Server specified Sec-WebSocket-Protocol that wasn't requested"] forKey:NSLocalizedDescriptionKey]]];

return;

}

_protocol = negotiatedProtocol;

}

self.readyState = SR_OPEN;

if (!_didFail) {

[self _readFrameNew];

}

[self _performDelegateBlock:^{

if ([self.delegate respondsToSelector:@selector(webSocketDidOpen:)]) {

[self.delegate webSocketDidOpen:self];

};

}];

}

- (void)_readHTTPHeader;

{

if (_receivedHTTPHeaders == NULL) {

_receivedHTTPHeaders = CFHTTPMessageCreateEmpty(NULL, NO);

}

[self _readUntilHeaderCompleteWithCallback:^(SRWebSocket *self, NSData *data) {

CFHTTPMessageAppendBytes(_receivedHTTPHeaders, (const UInt8 *)data.bytes, data.length);

if (CFHTTPMessageIsHeaderComplete(_receivedHTTPHeaders)) {

SRFastLog(@"Finished reading headers %@", CFBridgingRelease(CFHTTPMessageCopyAllHeaderFields(_receivedHTTPHeaders)));

[self _HTTPHeadersDidFinish];

} else {

[self _readHTTPHeader];

}

}];

}

- (void)didConnect

{

SRFastLog(@"Connected");

CFHTTPMessageRef request = CFHTTPMessageCreateRequest(NULL, CFSTR("GET"), (__bridge CFURLRef)_url, kCFHTTPVersion1_1);

// Set host first so it defaults

CFHTTPMessageSetHeaderFieldValue(request, CFSTR("Host"), (__bridge CFStringRef)(_url.port ? [NSString stringWithFormat:@"%@:%@", _url.host, _url.port] : _url.host));

NSMutableData *keyBytes = [[NSMutableData alloc] initWithLength:16];

SecRandomCopyBytes(kSecRandomDefault, keyBytes.length, keyBytes.mutableBytes);

_secKey = [keyBytes SR_stringByBase64Encoding];

assert([_secKey length] == 24);

CFHTTPMessageSetHeaderFieldValue(request, CFSTR("Upgrade"), CFSTR("websocket"));

CFHTTPMessageSetHeaderFieldValue(request, CFSTR("Connection"), CFSTR("Upgrade"));

CFHTTPMessageSetHeaderFieldValue(request, CFSTR("Sec-WebSocket-Key"), (__bridge CFStringRef)_secKey);

CFHTTPMessageSetHeaderFieldValue(request, CFSTR("Sec-WebSocket-Version"), (__bridge CFStringRef)[NSString stringWithFormat:@"%d", _webSocketVersion]);

CFHTTPMessageSetHeaderFieldValue(request, CFSTR("Origin"), (__bridge CFStringRef)_url.SR_origin);

if (_requestedProtocols) {

CFHTTPMessageSetHeaderFieldValue(request, CFSTR("Sec-WebSocket-Protocol"), (__bridge CFStringRef)[_requestedProtocols componentsJoinedByString:@", "]);

}

[_urlRequest.allHTTPHeaderFields enumerateKeysAndObjectsUsingBlock:^(id key, id obj, BOOL *stop) {

CFHTTPMessageSetHeaderFieldValue(request, (__bridge CFStringRef)key, (__bridge CFStringRef)obj);

}];

NSData *message = CFBridgingRelease(CFHTTPMessageCopySerializedMessage(request));

CFRelease(request);

[self _writeData:message];

[self _readHTTPHeader];

}

- (void)_initializeStreams;

{

NSInteger port = _url.port.integerValue;

if (port == 0) {

if (!_secure) {

port = 80;

} else {

port = 443;

}

}

NSString *host = _url.host;

CFReadStreamRef readStream = NULL;

CFWriteStreamRef writeStream = NULL;

CFStreamCreatePairWithSocketToHost(NULL, (__bridge CFStringRef)host, port, &readStream, &writeStream);

_outputStream = CFBridgingRelease(writeStream);

_inputStream = CFBridgingRelease(readStream);

if (_secure) {

NSMutableDictionary *SSLOptions = [[NSMutableDictionary alloc] init];

[_outputStream setProperty:(__bridge id)kCFStreamSocketSecurityLevelNegotiatedSSL forKey:(__bridge id)kCFStreamPropertySocketSecurityLevel];

// If we're using pinned certs, don't validate the certificate chain

if ([_urlRequest SR_SSLPinnedCertificates].count) {

[SSLOptions setValue:[NSNumber numberWithBool:NO] forKey:(__bridge id)kCFStreamSSLValidatesCertificateChain];

}

#if DEBUG

[SSLOptions setValue:[NSNumber numberWithBool:NO] forKey:(__bridge id)kCFStreamSSLValidatesCertificateChain];

NSLog(@"SocketRocket: In debug mode. Allowing connection to any root cert");

#endif

[_outputStream setProperty:SSLOptions

forKey:(__bridge id)kCFStreamPropertySSLSettings];

}

_inputStream.delegate = self;

_outputStream.delegate = self;

}

- (void)_connect;

{

if (!_didSchedule) {

[self scheduleInRunLoop:[NSRunLoop SR_networkRunLoop] forMode:NSDefaultRunLoopMode];

}

[_outputStream open];

[_inputStream open];

}

- (void)scheduleInRunLoop:(NSRunLoop *)aRunLoop forMode:(NSString *)mode;

{

_didSchedule = YES;

[_outputStream scheduleInRunLoop:aRunLoop forMode:mode];

[_inputStream scheduleInRunLoop:aRunLoop forMode:mode];

}

- (void)unscheduleFromRunLoop:(NSRunLoop *)aRunLoop forMode:(NSString *)mode;

{

[_outputStream removeFromRunLoop:aRunLoop forMode:mode];

[_inputStream removeFromRunLoop:aRunLoop forMode:mode];

}

- (void)close;

{

[self closeWithCode:-1 reason:nil];

}

- (void)closeWithCode:(NSInteger)code reason:(NSString *)reason;

{

assert(code);

if (self.readyState == SR_CLOSING || self.readyState == SR_CLOSED) {

return;

}

BOOL wasConnecting = self.readyState == SR_CONNECTING;

self.readyState = SR_CLOSING;

SRFastLog(@"Closing with code %d reason %@", code, reason);

dispatch_async(_workQueue, ^{

if (wasConnecting) {

[self _disconnect];

return;

}

size_t maxMsgSize = [reason maximumLengthOfBytesUsingEncoding:NSUTF8StringEncoding];

NSMutableData *mutablePayload = [[NSMutableData alloc] initWithLength:sizeof(uint16_t) + maxMsgSize];

NSData *payload = mutablePayload;

((uint16_t *)mutablePayload.mutableBytes)[0] = EndianU16_BtoN(code);

if (reason) {

NSRange remainingRange = {0};

NSUInteger usedLength = 0;

BOOL success = [reason getBytes:(char *)mutablePayload.mutableBytes + sizeof(uint16_t) maxLength:payload.length - sizeof(uint16_t) usedLength:&usedLength encoding:NSUTF8StringEncoding options:NSStringEncodingConversionExternalRepresentation range:NSMakeRange(0, reason.length) remainingRange:&remainingRange];

assert(success);

assert(remainingRange.length == 0);

if (usedLength != maxMsgSize) {

payload = [payload subdataWithRange:NSMakeRange(0, usedLength + sizeof(uint16_t))];

}

}

[self _sendFrameWithOpcode:SROpCodeConnectionClose data:payload];

});

}

- (void)_closeWithProtocolError:(NSString *)message;

{

// Need to shunt this on the _callbackQueue first to see if they received any messages

[self _performDelegateBlock:^{

[self closeWithCode:SRStatusCodeProtocolError reason:message];

dispatch_async(_workQueue, ^{

[self _disconnect];

});

}];

}

- (void)_failWithError:(NSError *)error;

{

dispatch_async(_workQueue, ^{

if (self.readyState != SR_CLOSED) {

_failed = YES;

[self _performDelegateBlock:^{

if ([self.delegate respondsToSelector:@selector(webSocket:didFailWithError:)]) {

[self.delegate webSocket:self didFailWithError:error];

}

}];

self.readyState = SR_CLOSED;

_selfRetain = nil;

SRFastLog(@"Failing with error %@", error.localizedDescription);

[self _disconnect];

}

});

}

- (void)_writeData:(NSData *)data;

{

assert(dispatch_get_current_queue() == _workQueue);

if (_closeWhenFinishedWriting) {

return;

}

[_outputBuffer appendData:data];

[self _pumpWriting];

}

- (void)send:(id)data;

{

NSAssert(self.readyState != SR_CONNECTING, @"Invalid State: Cannot call send: until connection is open");

// TODO: maybe not copy this for performance

data = [data copy];

dispatch_async(_workQueue, ^{

if ([data isKindOfClass:[NSString class]]) {

[self _sendFrameWithOpcode:SROpCodeTextFrame data:[(NSString *)data dataUsingEncoding:NSUTF8StringEncoding]];

} else if ([data isKindOfClass:[NSData class]]) {

[self _sendFrameWithOpcode:SROpCodeBinaryFrame data:data];

} else if (data == nil) {

[self _sendFrameWithOpcode:SROpCodeTextFrame data:data];

} else {

assert(NO);

}

});

}

- (void)handlePing:(NSData *)pingData;

{

// Need to pingpong this off _callbackQueue first to make sure messages happen in order

[self _performDelegateBlock:^{

dispatch_async(_workQueue, ^{

[self _sendFrameWithOpcode:SROpCodePong data:pingData];

});

}];

}

- (void)handlePong;

{

// NOOP

}

- (void)_handleMessage:(id)message

{

SRFastLog(@"Received message");

[self _performDelegateBlock:^{

[self.delegate webSocket:self didReceiveMessage:message];

}];

}

static inline BOOL closeCodeIsValid(int closeCode) {

if (closeCode < 1000) {

return NO;

}

if (closeCode >= 1000 && closeCode <= 1011) {

if (closeCode == 1004 ||

closeCode == 1005 ||

closeCode == 1006) {

return NO;

}

return YES;

}

if (closeCode >= 3000 && closeCode <= 3999) {

return YES;

}

if (closeCode >= 4000 && closeCode <= 4999) {

return YES;

}

return NO;

}

// Note from RFC:

//

// If there is a body, the first two

// bytes of the body MUST be a 2-byte unsigned integer (in network byte

// order) representing a status code with value /code/ defined in

// Section 7.4. Following the 2-byte integer the body MAY contain UTF-8

// encoded data with value /reason/, the interpretation of which is not

// defined by this specification.

- (void)handleCloseWithData:(NSData *)data;

{

size_t dataSize = data.length;

__block uint16_t closeCode = 0;

SRFastLog(@"Received close frame");

if (dataSize == 1) {

// TODO handle error

[self _closeWithProtocolError:@"Payload for close must be larger than 2 bytes"];

return;

} else if (dataSize >= 2) {

[data getBytes:&closeCode length:sizeof(closeCode)];

_closeCode = EndianU16_BtoN(closeCode);

if (!closeCodeIsValid(_closeCode)) {

[self _closeWithProtocolError:[NSString stringWithFormat:@"Cannot have close code of %d", _closeCode]];

return;

}

if (dataSize > 2) {

_closeReason = [[NSString alloc] initWithData:[data subdataWithRange:NSMakeRange(2, dataSize - 2)] encoding:NSUTF8StringEncoding];

if (!_closeReason) {

[self _closeWithProtocolError:@"Close reason MUST be valid UTF-8"];

return;

}

}

} else {

_closeCode = SRStatusNoStatusReceived;

}

assert(dispatch_get_current_queue() == _workQueue);

if (self.readyState == SR_OPEN) {

[self closeWithCode:1000 reason:nil];

}

dispatch_async(_workQueue, ^{

[self _disconnect];

});

}

- (void)_disconnect;

{

assert(dispatch_get_current_queue() == _workQueue);

SRFastLog(@"Trying to disconnect");

_closeWhenFinishedWriting = YES;

[self _pumpWriting];

}

- (void)_handleFrameWithData:(NSData *)frameData opCode:(NSInteger)opcode;

{

// Check that the current data is valid UTF8

BOOL isControlFrame = (opcode == SROpCodePing || opcode == SROpCodePong || opcode == SROpCodeConnectionClose);

if (!isControlFrame) {

[self _readFrameNew];

} else {

dispatch_async(_workQueue, ^{

[self _readFrameContinue];

});

}

switch (opcode) {

case SROpCodeTextFrame: {

NSString *str = [[NSString alloc] initWithData:frameData encoding:NSUTF8StringEncoding];

if (str == nil && frameData) {

[self closeWithCode:SRStatusCodeInvalidUTF8 reason:@"Text frames must be valid UTF-8"];

dispatch_async(_workQueue, ^{

[self _disconnect];

});

return;

}

[self _handleMessage:str];

break;

}

case SROpCodeBinaryFrame:

[self _handleMessage:[frameData copy]];

break;

case SROpCodeConnectionClose:

[self handleCloseWithData:frameData];

break;

case SROpCodePing:

[self handlePing:frameData];

break;

case SROpCodePong:

[self handlePong];

break;

default:

[self _closeWithProtocolError:[NSString stringWithFormat:@"Unknown opcode %d", opcode]];

// TODO: Handle invalid opcode

break;

}

}

- (void)_handleFrameHeader:(frame_header)frame_header curData:(NSData *)curData;

{

assert(frame_header.opcode != 0);

if (self.readyState != SR_OPEN) {

return;

}

BOOL isControlFrame = (frame_header.opcode == SROpCodePing || frame_header.opcode == SROpCodePong || frame_header.opcode == SROpCodeConnectionClose);

if (isControlFrame && !frame_header.fin) {

[self _closeWithProtocolError:@"Fragmented control frames not allowed"];

return;

}

if (isControlFrame && frame_header.payload_length >= 126) {

[self _closeWithProtocolError:@"Control frames cannot have payloads larger than 126 bytes"];

return;

}

if (!isControlFrame) {

_currentFrameOpcode = frame_header.opcode;

_currentFrameCount += 1;

}

if (frame_header.payload_length == 0) {

if (isControlFrame) {

[self _handleFrameWithData:curData opCode:frame_header.opcode];

} else {

if (frame_header.fin) {

[self _handleFrameWithData:_currentFrameData opCode:frame_header.opcode];

} else {

// TODO add assert that opcode is not a control;

[self _readFrameContinue];

}

}

} else {

[self _addConsumerWithDataLength:frame_header.payload_length callback:^(SRWebSocket *self, NSData *newData) {

if (isControlFrame) {

[self _handleFrameWithData:newData opCode:frame_header.opcode];

} else {

if (frame_header.fin) {

[self _handleFrameWithData:self->_currentFrameData opCode:frame_header.opcode];

} else {

// TODO add assert that opcode is not a control;

[self _readFrameContinue];

}

}

} readToCurrentFrame:!isControlFrame unmaskBytes:frame_header.masked];

}

}

/* From RFC:

0 1 2 3

0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1

+-+-+-+-+-------+-+-------------+-------------------------------+

|F|R|R|R| opcode|M| Payload len | Extended payload length |

|I|S|S|S| (4) |A| (7) | (16/64) |

|N|V|V|V| |S| | (if payload len==126/127) |

| |1|2|3| |K| | |

+-+-+-+-+-------+-+-------------+ - - - - - - - - - - - - - - - +

| Extended payload length continued, if payload len == 127 |

+ - - - - - - - - - - - - - - - +-------------------------------+

| |Masking-key, if MASK set to 1 |

+-------------------------------+-------------------------------+

| Masking-key (continued) | Payload Data |

+-------------------------------- - - - - - - - - - - - - - - - +

: Payload Data continued ... :

+ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - +

| Payload Data continued ... |

+---------------------------------------------------------------+

*/

static const uint8_t SRFinMask = 0x80;

static const uint8_t SROpCodeMask = 0x0F;

static const uint8_t SRRsvMask = 0x70;

static const uint8_t SRMaskMask = 0x80;

static const uint8_t SRPayloadLenMask = 0x7F;

- (void)_readFrameContinue;

{

assert((_currentFrameCount == 0 && _currentFrameOpcode == 0) || (_currentFrameCount > 0 && _currentFrameOpcode > 0));

[self _addConsumerWithDataLength:2 callback:^(SRWebSocket *self, NSData *data) {

__block frame_header header = {0};

const uint8_t *headerBuffer = data.bytes;

assert(data.length >= 2);

if (headerBuffer[0] & SRRsvMask) {

[self _closeWithProtocolError:@"Server used RSV bits"];

return;

}

uint8_t receivedOpcode = (SROpCodeMask & headerBuffer[0]);

BOOL isControlFrame = (receivedOpcode == SROpCodePing || receivedOpcode == SROpCodePong || receivedOpcode == SROpCodeConnectionClose);

if (!isControlFrame && receivedOpcode != 0 && self->_currentFrameCount > 0) {

[self _closeWithProtocolError:@"all data frames after the initial data frame must have opcode 0"];

return;

}

if (receivedOpcode == 0 && self->_currentFrameCount == 0) {

[self _closeWithProtocolError:@"cannot continue a message"];

return;

}

header.opcode = receivedOpcode == 0 ? self->_currentFrameOpcode : receivedOpcode;

header.fin = !!(SRFinMask & headerBuffer[0]);

header.masked = !!(SRMaskMask & headerBuffer[1]);

header.payload_length = SRPayloadLenMask & headerBuffer[1];

headerBuffer = NULL;

if (header.masked) {

[self _closeWithProtocolError:@"Client must receive unmasked data"];

}

size_t extra_bytes_needed = header.masked ? sizeof(_currentReadMaskKey) : 0;

if (header.payload_length == 126) {

extra_bytes_needed += sizeof(uint16_t);

} else if (header.payload_length == 127) {

extra_bytes_needed += sizeof(uint64_t);

}

if (extra_bytes_needed == 0) {

[self _handleFrameHeader:header curData:self->_currentFrameData];

} else {

[self _addConsumerWithDataLength:extra_bytes_needed callback:^(SRWebSocket *self, NSData *data) {

size_t mapped_size = data.length;

const void *mapped_buffer = data.bytes;

size_t offset = 0;