1//
2//   Copyright 2012 Square Inc.
3//
4//   Licensed under the Apache License, Version 2.0 (the "License");
5//   you may not use this file except in compliance with the License.
6//   You may obtain a copy of the License at
7//
8//       http://www.apache.org/licenses/LICENSE-2.0
9//
10//   Unless required by applicable law or agreed to in writing, software
11//   distributed under the License is distributed on an "AS IS" BASIS,
12//   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13//   See the License for the specific language governing permissions and
14//   limitations under the License.
15//
16
17
18#import "SRWebSocket.h"
19
20#if TARGET_OS_IPHONE
21#define HAS_ICU
22#endif
23
24#ifdef HAS_ICU
25#import <unicode/utf8.h>
26#endif
27
28#if TARGET_OS_IPHONE
29#import <Endian.h>
30#else
31#import <CoreServices/CoreServices.h>
32#endif
33
34#import <CommonCrypto/CommonDigest.h>
35#import <Security/SecRandom.h>
36
37#if OS_OBJECT_USE_OBJC_RETAIN_RELEASE
38#define sr_dispatch_retain(x)
39#define sr_dispatch_release(x)
40#define maybe_bridge(x) ((__bridge void *) x)
41#else
42#define sr_dispatch_retain(x) dispatch_retain(x)
43#define sr_dispatch_release(x) dispatch_release(x)
44#define maybe_bridge(x) (x)
45#endif
46
47#if !__has_feature(objc_arc)
48#error SocketRocket must be compiled with ARC enabled
49#endif
50
51
52typedef enum  {
53    SROpCodeTextFrame = 0x1,
54    SROpCodeBinaryFrame = 0x2,
55    // 3-7 reserved.
56    SROpCodeConnectionClose = 0x8,
57    SROpCodePing = 0x9,
58    SROpCodePong = 0xA,
59    // B-F reserved.
60} SROpCode;
61
62typedef struct {
63    BOOL fin;
64//  BOOL rsv1;
65//  BOOL rsv2;
66//  BOOL rsv3;
67    uint8_t opcode;
68    BOOL masked;
69    uint64_t payload_length;
70} frame_header;
71
72static NSString *const SRWebSocketAppendToSecKeyString = @"258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
73
74static inline int32_t validate_dispatch_data_partial_string(NSData *data);
75static inline void SRFastLog(NSString *format, ...);
76
77@interface NSData (SRWebSocket)
78
79- (NSString *)stringBySHA1ThenBase64Encoding;
80
81@end
82
83
84@interface NSString (SRWebSocket)
85
86- (NSString *)stringBySHA1ThenBase64Encoding;
87
88@end
89
90
91@interface NSURL (SRWebSocket)
92
93// The origin isn't really applicable for a native application.
94// So instead, just map ws -> http and wss -> https.
95- (NSString *)SR_origin;
96
97@end
98
99
100@interface _SRRunLoopThread : NSThread
101
102@property (nonatomic, readonly) NSRunLoop *runLoop;
103
104@end
105
106
107static NSString *newSHA1String(const char *bytes, size_t length) {
108    uint8_t md[CC_SHA1_DIGEST_LENGTH];
109
110    assert(length >= 0);
111    assert(length <= UINT32_MAX);
112    CC_SHA1(bytes, (CC_LONG)length, md);
113
114    NSData *data = [NSData dataWithBytes:md length:CC_SHA1_DIGEST_LENGTH];
115
116    if ([data respondsToSelector:@selector(base64EncodedStringWithOptions:)]) {
117        return [data base64EncodedStringWithOptions:0];
118    }
119
120    return [data base64Encoding];
121}
122
123@implementation NSData (SRWebSocket)
124
125- (NSString *)stringBySHA1ThenBase64Encoding;
126{
127    return newSHA1String(self.bytes, self.length);
128}
129
130@end
131
132
133@implementation NSString (SRWebSocket)
134
135- (NSString *)stringBySHA1ThenBase64Encoding;
136{
137    return newSHA1String(self.UTF8String, self.length);
138}
139
140@end
141
142NSString *const SRWebSocketErrorDomain = @"SRWebSocketErrorDomain";
143NSString *const SRHTTPResponseErrorKey = @"HTTPResponseStatusCode";
144
145// Returns number of bytes consumed. Returning 0 means you didn't match.
146// Sends bytes to callback handler;
147typedef size_t (^stream_scanner)(NSData *collected_data);
148
149typedef void (^data_callback)(SRWebSocket *webSocket,  NSData *data);
150
151@interface SRIOConsumer : NSObject {
152    stream_scanner _scanner;
153    data_callback _handler;
154    size_t _bytesNeeded;
155    BOOL _readToCurrentFrame;
156    BOOL _unmaskBytes;
157}
158@property (nonatomic, copy, readonly) stream_scanner consumer;
159@property (nonatomic, copy, readonly) data_callback handler;
160@property (nonatomic, assign) size_t bytesNeeded;
161@property (nonatomic, assign, readonly) BOOL readToCurrentFrame;
162@property (nonatomic, assign, readonly) BOOL unmaskBytes;
163
164@end
165
166// This class is not thread-safe, and is expected to always be run on the same queue.
167@interface SRIOConsumerPool : NSObject
168
169- (id)initWithBufferCapacity:(NSUInteger)poolSize;
170
171- (SRIOConsumer *)consumerWithScanner:(stream_scanner)scanner handler:(data_callback)handler bytesNeeded:(size_t)bytesNeeded readToCurrentFrame:(BOOL)readToCurrentFrame unmaskBytes:(BOOL)unmaskBytes;
172- (void)returnConsumer:(SRIOConsumer *)consumer;
173
174@end
175
176@interface SRWebSocket ()  <NSStreamDelegate>
177
178- (void)_writeData:(NSData *)data;
179- (void)_closeWithProtocolError:(NSString *)message;
180- (void)_failWithError:(NSError *)error;
181
182- (void)_disconnect;
183
184- (void)_readFrameNew;
185- (void)_readFrameContinue;
186
187- (void)_pumpScanner;
188
189- (void)_pumpWriting;
190
191- (void)_addConsumerWithScanner:(stream_scanner)consumer callback:(data_callback)callback;
192- (void)_addConsumerWithDataLength:(size_t)dataLength callback:(data_callback)callback readToCurrentFrame:(BOOL)readToCurrentFrame unmaskBytes:(BOOL)unmaskBytes;
193- (void)_addConsumerWithScanner:(stream_scanner)consumer callback:(data_callback)callback dataLength:(size_t)dataLength;
194- (void)_readUntilBytes:(const void *)bytes length:(size_t)length callback:(data_callback)dataHandler;
195- (void)_readUntilHeaderCompleteWithCallback:(data_callback)dataHandler;
196
197- (void)_sendFrameWithOpcode:(SROpCode)opcode data:(id)data;
198
199- (BOOL)_checkHandshake:(CFHTTPMessageRef)httpMessage;
200- (void)_SR_commonInit;
201
202- (void)_initializeStreams;
203- (void)_connect;
204
205@property (nonatomic) SRReadyState readyState;
206
207@property (nonatomic) NSOperationQueue *delegateOperationQueue;
208@property (nonatomic) dispatch_queue_t delegateDispatchQueue;
209
210@end
211
212
213@implementation SRWebSocket {
214    NSInteger _webSocketVersion;
215
216    NSOperationQueue *_delegateOperationQueue;
217    dispatch_queue_t _delegateDispatchQueue;
218
219    dispatch_queue_t _workQueue;
220    NSMutableArray *_consumers;
221
222    NSInputStream *_inputStream;
223    NSOutputStream *_outputStream;
224
225    NSMutableData *_readBuffer;
226    NSUInteger _readBufferOffset;
227
228    NSMutableData *_outputBuffer;
229    NSUInteger _outputBufferOffset;
230
231    uint8_t _currentFrameOpcode;
232    size_t _currentFrameCount;
233    size_t _readOpCount;
234    uint32_t _currentStringScanPosition;
235    NSMutableData *_currentFrameData;
236
237    NSString *_closeReason;
238
239    NSString *_secKey;
240
241    BOOL _pinnedCertFound;
242
243    uint8_t _currentReadMaskKey[4];
244    size_t _currentReadMaskOffset;
245
246    BOOL _consumerStopped;
247
248    BOOL _closeWhenFinishedWriting;
249    BOOL _failed;
250
251    BOOL _secure;
252    NSURLRequest *_urlRequest;
253
254    CFHTTPMessageRef _receivedHTTPHeaders;
255
256    BOOL _sentClose;
257    BOOL _didFail;
258    int _closeCode;
259
260    BOOL _isPumping;
261
262    NSMutableSet *_scheduledRunloops;
263
264    // We use this to retain ourselves.
265    __strong SRWebSocket *_selfRetain;
266
267    NSArray *_requestedProtocols;
268    SRIOConsumerPool *_consumerPool;
269}
270
271@synthesize delegate = _delegate;
272@synthesize url = _url;
273@synthesize readyState = _readyState;
274@synthesize protocol = _protocol;
275
276static __strong NSData *CRLFCRLF;
277
278+ (void)initialize;
279{
280    CRLFCRLF = [[NSData alloc] initWithBytes:"\r\n\r\n" length:4];
281}
282
283- (id)initWithURLRequest:(NSURLRequest *)request protocols:(NSArray *)protocols;
284{
285    self = [super init];
286    if (self) {
287        assert(request.URL);
288        _url = request.URL;
289        _urlRequest = request;
290
291        _requestedProtocols = [protocols copy];
292
293        [self _SR_commonInit];
294    }
295
296    return self;
297}
298
299- (id)initWithURLRequest:(NSURLRequest *)request;
300{
301    return [self initWithURLRequest:request protocols:nil];
302}
303
304- (id)initWithURL:(NSURL *)url;
305{
306    return [self initWithURL:url protocols:nil];
307}
308
309- (id)initWithURL:(NSURL *)url protocols:(NSArray *)protocols;
310{
311    NSMutableURLRequest *request = [[NSMutableURLRequest alloc] initWithURL:url];
312    return [self initWithURLRequest:request protocols:protocols];
313}
314
315- (void)_SR_commonInit;
316{
317
318    NSString *scheme = _url.scheme.lowercaseString;
319    assert([scheme isEqualToString:@"ws"] || [scheme isEqualToString:@"http"] || [scheme isEqualToString:@"wss"] || [scheme isEqualToString:@"https"]);
320
321    if ([scheme isEqualToString:@"wss"] || [scheme isEqualToString:@"https"]) {
322        _secure = YES;
323    }
324
325    _readyState = SR_CONNECTING;
326    _consumerStopped = YES;
327    _webSocketVersion = 13;
328
329    _workQueue = dispatch_queue_create(NULL, DISPATCH_QUEUE_SERIAL);
330
331    // Going to set a specific on the queue so we can validate we're on the work queue
332    dispatch_queue_set_specific(_workQueue, (__bridge void *)self, maybe_bridge(_workQueue), NULL);
333
334    _delegateDispatchQueue = dispatch_get_main_queue();
335    sr_dispatch_retain(_delegateDispatchQueue);
336
337    _readBuffer = [[NSMutableData alloc] init];
338    _outputBuffer = [[NSMutableData alloc] init];
339
340    _currentFrameData = [[NSMutableData alloc] init];
341
342    _consumers = [[NSMutableArray alloc] init];
343
344    _consumerPool = [[SRIOConsumerPool alloc] init];
345
346    _scheduledRunloops = [[NSMutableSet alloc] init];
347
348    [self _initializeStreams];
349
350    // default handlers
351}
352
353- (void)assertOnWorkQueue;
354{
355    assert(dispatch_get_specific((__bridge void *)self) == maybe_bridge(_workQueue));
356}
357
358- (void)dealloc
359{
360    _inputStream.delegate = nil;
361    _outputStream.delegate = nil;
362
363    [_inputStream close];
364    [_outputStream close];
365
366    sr_dispatch_release(_workQueue);
367    _workQueue = NULL;
368
369    if (_receivedHTTPHeaders) {
370        CFRelease(_receivedHTTPHeaders);
371        _receivedHTTPHeaders = NULL;
372    }
373
374    if (_delegateDispatchQueue) {
375        sr_dispatch_release(_delegateDispatchQueue);
376        _delegateDispatchQueue = NULL;
377    }
378}
379
380#ifndef NDEBUG
381
382- (void)setReadyState:(SRReadyState)aReadyState;
383{
384    [self willChangeValueForKey:@"readyState"];
385    assert(aReadyState > _readyState);
386    _readyState = aReadyState;
387    [self didChangeValueForKey:@"readyState"];
388}
389
390#endif
391
392- (void)open;
393{
394    assert(_url);
395    NSAssert(_readyState == SR_CONNECTING, @"Cannot call -(void)open on SRWebSocket more than once");
396
397    _selfRetain = self;
398
399    [self _connect];
400}
401
402// Calls block on delegate queue
403- (void)_performDelegateBlock:(dispatch_block_t)block;
404{
405    if (_delegateOperationQueue) {
406        [_delegateOperationQueue addOperationWithBlock:block];
407    } else {
408        assert(_delegateDispatchQueue);
409        dispatch_async(_delegateDispatchQueue, block);
410    }
411}
412
413- (void)setDelegateDispatchQueue:(dispatch_queue_t)queue;
414{
415    if (queue) {
416        sr_dispatch_retain(queue);
417    }
418
419    if (_delegateDispatchQueue) {
420        sr_dispatch_release(_delegateDispatchQueue);
421    }
422
423    _delegateDispatchQueue = queue;
424}
425
426- (BOOL)_checkHandshake:(CFHTTPMessageRef)httpMessage;
427{
428    NSString *acceptHeader = CFBridgingRelease(CFHTTPMessageCopyHeaderFieldValue(httpMessage, CFSTR("Sec-WebSocket-Accept")));
429
430    if (acceptHeader == nil) {
431        return NO;
432    }
433
434    NSString *concattedString = [_secKey stringByAppendingString:SRWebSocketAppendToSecKeyString];
435    NSString *expectedAccept = [concattedString stringBySHA1ThenBase64Encoding];
436
437    return [acceptHeader isEqualToString:expectedAccept];
438}
439
440- (void)_HTTPHeadersDidFinish;
441{
442    NSInteger responseCode = CFHTTPMessageGetResponseStatusCode(_receivedHTTPHeaders);
443
444    if (responseCode >= 400) {
445        SRFastLog(@"Request failed with response code %d", responseCode);
446        [self _failWithError:[NSError errorWithDomain:SRWebSocketErrorDomain code:2132 userInfo:@{NSLocalizedDescriptionKey:[NSString stringWithFormat:@"received bad response code from server %ld", (long)responseCode], SRHTTPResponseErrorKey:@(responseCode)}]];
447        return;
448    }
449
450    if(![self _checkHandshake:_receivedHTTPHeaders]) {
451        [self _failWithError:[NSError errorWithDomain:SRWebSocketErrorDomain code:2133 userInfo:[NSDictionary dictionaryWithObject:[NSString stringWithFormat:@"Invalid Sec-WebSocket-Accept response"] forKey:NSLocalizedDescriptionKey]]];
452        return;
453    }
454
455    NSString *negotiatedProtocol = CFBridgingRelease(CFHTTPMessageCopyHeaderFieldValue(_receivedHTTPHeaders, CFSTR("Sec-WebSocket-Protocol")));
456    if (negotiatedProtocol) {
457        // Make sure we requested the protocol
458        if ([_requestedProtocols indexOfObject:negotiatedProtocol] == NSNotFound) {
459            [self _failWithError:[NSError errorWithDomain:SRWebSocketErrorDomain code:2133 userInfo:[NSDictionary dictionaryWithObject:[NSString stringWithFormat:@"Server specified Sec-WebSocket-Protocol that wasn't requested"] forKey:NSLocalizedDescriptionKey]]];
460            return;
461        }
462
463        _protocol = negotiatedProtocol;
464    }
465
466    self.readyState = SR_OPEN;
467
468    if (!_didFail) {
469        [self _readFrameNew];
470    }
471
472    [self _performDelegateBlock:^{
473        if ([self.delegate respondsToSelector:@selector(webSocketDidOpen:)]) {
474            [self.delegate webSocketDidOpen:self];
475        };
476    }];
477}
478
479
480- (void)_readHTTPHeader;
481{
482    if (_receivedHTTPHeaders == NULL) {
483        _receivedHTTPHeaders = CFHTTPMessageCreateEmpty(NULL, NO);
484    }
485
486    [self _readUntilHeaderCompleteWithCallback:^(SRWebSocket *self,  NSData *data) {
487        CFHTTPMessageAppendBytes(_receivedHTTPHeaders, (const UInt8 *)data.bytes, data.length);
488
489        if (CFHTTPMessageIsHeaderComplete(_receivedHTTPHeaders)) {
490            SRFastLog(@"Finished reading headers %@", CFBridgingRelease(CFHTTPMessageCopyAllHeaderFields(_receivedHTTPHeaders)));
491            [self _HTTPHeadersDidFinish];
492        } else {
493            [self _readHTTPHeader];
494        }
495    }];
496}
497
498- (void)didConnect
499{
500    SRFastLog(@"Connected");
501    CFHTTPMessageRef request = CFHTTPMessageCreateRequest(NULL, CFSTR("GET"), (__bridge CFURLRef)_url, kCFHTTPVersion1_1);
502
503    // Set host first so it defaults
504    CFHTTPMessageSetHeaderFieldValue(request, CFSTR("Host"), (__bridge CFStringRef)(_url.port ? [NSString stringWithFormat:@"%@:%@", _url.host, _url.port] : _url.host));
505
506    NSMutableData *keyBytes = [[NSMutableData alloc] initWithLength:16];
507    SecRandomCopyBytes(kSecRandomDefault, keyBytes.length, keyBytes.mutableBytes);
508
509    if ([keyBytes respondsToSelector:@selector(base64EncodedStringWithOptions:)]) {
510        _secKey = [keyBytes base64EncodedStringWithOptions:0];
511    } else {
512        _secKey = [keyBytes base64Encoding];
513    }
514
515    assert([_secKey length] == 24);
516
517    CFHTTPMessageSetHeaderFieldValue(request, CFSTR("Upgrade"), CFSTR("websocket"));
518    CFHTTPMessageSetHeaderFieldValue(request, CFSTR("Connection"), CFSTR("Upgrade"));
519    CFHTTPMessageSetHeaderFieldValue(request, CFSTR("Sec-WebSocket-Key"), (__bridge CFStringRef)_secKey);
520    CFHTTPMessageSetHeaderFieldValue(request, CFSTR("Sec-WebSocket-Version"), (__bridge CFStringRef)[NSString stringWithFormat:@"%ld", (long)_webSocketVersion]);
521
522    CFHTTPMessageSetHeaderFieldValue(request, CFSTR("Origin"), (__bridge CFStringRef)_url.SR_origin);
523
524    if (_requestedProtocols) {
525        CFHTTPMessageSetHeaderFieldValue(request, CFSTR("Sec-WebSocket-Protocol"), (__bridge CFStringRef)[_requestedProtocols componentsJoinedByString:@", "]);
526    }
527
528    [_urlRequest.allHTTPHeaderFields enumerateKeysAndObjectsUsingBlock:^(id key, id obj, BOOL *stop) {
529        CFHTTPMessageSetHeaderFieldValue(request, (__bridge CFStringRef)key, (__bridge CFStringRef)obj);
530    }];
531
532    NSData *message = CFBridgingRelease(CFHTTPMessageCopySerializedMessage(request));
533
534    CFRelease(request);
535
536    [self _writeData:message];
537    [self _readHTTPHeader];
538}
539
540- (void)_initializeStreams;
541{
542    assert(_url.port.unsignedIntValue <= UINT32_MAX);
543    uint32_t port = _url.port.unsignedIntValue;
544    if (port == 0) {
545        if (!_secure) {
546            port = 80;
547        } else {
548            port = 443;
549        }
550    }
551    NSString *host = _url.host;
552
553    CFReadStreamRef readStream = NULL;
554    CFWriteStreamRef writeStream = NULL;
555
556    CFStreamCreatePairWithSocketToHost(NULL, (__bridge CFStringRef)host, port, &readStream, &writeStream);
557
558    _outputStream = CFBridgingRelease(writeStream);
559    _inputStream = CFBridgingRelease(readStream);
560
561
562    if (_secure) {
563        NSMutableDictionary *SSLOptions = [[NSMutableDictionary alloc] init];
564
565        [_outputStream setProperty:(__bridge id)kCFStreamSocketSecurityLevelNegotiatedSSL forKey:(__bridge id)kCFStreamPropertySocketSecurityLevel];
566
567        // If we're using pinned certs, don't validate the certificate chain
568        if ([_urlRequest SR_SSLPinnedCertificates].count) {
569            [SSLOptions setValue:[NSNumber numberWithBool:NO] forKey:(__bridge id)kCFStreamSSLValidatesCertificateChain];
570        }
571
572#if DEBUG
573        [SSLOptions setValue:[NSNumber numberWithBool:NO] forKey:(__bridge id)kCFStreamSSLValidatesCertificateChain];
574        NSLog(@"SocketRocket: In debug mode.  Allowing connection to any root cert");
575#endif
576
577        [_outputStream setProperty:SSLOptions
578                            forKey:(__bridge id)kCFStreamPropertySSLSettings];
579    }
580
581    _inputStream.delegate = self;
582    _outputStream.delegate = self;
583}
584
585- (void)_connect;
586{
587    if (!_scheduledRunloops.count) {
588        [self scheduleInRunLoop:[NSRunLoop SR_networkRunLoop] forMode:NSDefaultRunLoopMode];
589    }
590
591
592    [_outputStream open];
593    [_inputStream open];
594}
595
596- (void)scheduleInRunLoop:(NSRunLoop *)aRunLoop forMode:(NSString *)mode;
597{
598    [_outputStream scheduleInRunLoop:aRunLoop forMode:mode];
599    [_inputStream scheduleInRunLoop:aRunLoop forMode:mode];
600
601    [_scheduledRunloops addObject:@[aRunLoop, mode]];
602}
603
604- (void)unscheduleFromRunLoop:(NSRunLoop *)aRunLoop forMode:(NSString *)mode;
605{
606    [_outputStream removeFromRunLoop:aRunLoop forMode:mode];
607    [_inputStream removeFromRunLoop:aRunLoop forMode:mode];
608
609    [_scheduledRunloops removeObject:@[aRunLoop, mode]];
610}
611
612- (void)close;
613{
614    [self closeWithCode:SRStatusCodeNormal reason:nil];
615}
616
617- (void)closeWithCode:(NSInteger)code reason:(NSString *)reason;
618{
619    assert(code);
620    dispatch_async(_workQueue, ^{
621        if (self.readyState == SR_CLOSING || self.readyState == SR_CLOSED) {
622            return;
623        }
624
625        BOOL wasConnecting = self.readyState == SR_CONNECTING;
626
627        self.readyState = SR_CLOSING;
628
629        SRFastLog(@"Closing with code %d reason %@", code, reason);
630
631        if (wasConnecting) {
632            [self _disconnect];
633            return;
634        }
635
636        size_t maxMsgSize = [reason maximumLengthOfBytesUsingEncoding:NSUTF8StringEncoding];
637        NSMutableData *mutablePayload = [[NSMutableData alloc] initWithLength:sizeof(uint16_t) + maxMsgSize];
638        NSData *payload = mutablePayload;
639
640        ((uint16_t *)mutablePayload.mutableBytes)[0] = EndianU16_BtoN(code);
641
642        if (reason) {
643            NSRange remainingRange = {0};
644
645            NSUInteger usedLength = 0;
646
647            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];
648
649            assert(success);
650            assert(remainingRange.length == 0);
651
652            if (usedLength != maxMsgSize) {
653                payload = [payload subdataWithRange:NSMakeRange(0, usedLength + sizeof(uint16_t))];
654            }
655        }
656
657
658        [self _sendFrameWithOpcode:SROpCodeConnectionClose data:payload];
659    });
660}
661
662- (void)_closeWithProtocolError:(NSString *)message;
663{
664    // Need to shunt this on the _callbackQueue first to see if they received any messages
665    [self _performDelegateBlock:^{
666        [self closeWithCode:SRStatusCodeProtocolError reason:message];
667        dispatch_async(_workQueue, ^{
668            [self _disconnect];
669        });
670    }];
671}
672
673- (void)_failWithError:(NSError *)error;
674{
675    dispatch_async(_workQueue, ^{
676        if (self.readyState != SR_CLOSED) {
677            _failed = YES;
678            [self _performDelegateBlock:^{
679                if ([self.delegate respondsToSelector:@selector(webSocket:didFailWithError:)]) {
680                    [self.delegate webSocket:self didFailWithError:error];
681                }
682            }];
683
684            self.readyState = SR_CLOSED;
685            _selfRetain = nil;
686
687            SRFastLog(@"Failing with error %@", error.localizedDescription);
688
689            [self _disconnect];
690        }
691    });
692}
693
694- (void)_writeData:(NSData *)data;
695{
696    [self assertOnWorkQueue];
697
698    if (_closeWhenFinishedWriting) {
699            return;
700    }
701    [_outputBuffer appendData:data];
702    [self _pumpWriting];
703}
704
705- (void)send:(id)data;
706{
707    NSAssert(self.readyState != SR_CONNECTING, @"Invalid State: Cannot call send: until connection is open");
708    // TODO: maybe not copy this for performance
709    data = [data copy];
710    dispatch_async(_workQueue, ^{
711        if ([data isKindOfClass:[NSString class]]) {
712            [self _sendFrameWithOpcode:SROpCodeTextFrame data:[(NSString *)data dataUsingEncoding:NSUTF8StringEncoding]];
713        } else if ([data isKindOfClass:[NSData class]]) {
714            [self _sendFrameWithOpcode:SROpCodeBinaryFrame data:data];
715        } else if (data == nil) {
716            [self _sendFrameWithOpcode:SROpCodeTextFrame data:data];
717        } else {
718            assert(NO);
719        }
720    });
721}
722
723- (void)sendPing:(NSData *)data;
724{
725    NSAssert(self.readyState == SR_OPEN, @"Invalid State: Cannot call send: until connection is open");
726    // TODO: maybe not copy this for performance
727    data = [data copy] ?: [NSData data]; // It's okay for a ping to be empty
728    dispatch_async(_workQueue, ^{
729        [self _sendFrameWithOpcode:SROpCodePing data:data];
730    });
731}
732
733- (void)handlePing:(NSData *)pingData;
734{
735    // Need to pingpong this off _callbackQueue first to make sure messages happen in order
736    [self _performDelegateBlock:^{
737        dispatch_async(_workQueue, ^{
738            [self _sendFrameWithOpcode:SROpCodePong data:pingData];
739        });
740    }];
741}
742
743- (void)handlePong:(NSData *)pongData;
744{
745    SRFastLog(@"Received pong");
746    [self _performDelegateBlock:^{
747        if ([self.delegate respondsToSelector:@selector(webSocket:didReceivePong:)]) {
748            [self.delegate webSocket:self didReceivePong:pongData];
749        }
750    }];
751}
752
753- (void)_handleMessage:(id)message
754{
755    SRFastLog(@"Received message");
756    [self _performDelegateBlock:^{
757        [self.delegate webSocket:self didReceiveMessage:message];
758    }];
759}
760
761
762static inline BOOL closeCodeIsValid(int closeCode) {
763    if (closeCode < 1000) {
764        return NO;
765    }
766
767    if (closeCode >= 1000 && closeCode <= 1011) {
768        if (closeCode == 1004 ||
769            closeCode == 1005 ||
770            closeCode == 1006) {
771            return NO;
772        }
773        return YES;
774    }
775
776    if (closeCode >= 3000 && closeCode <= 3999) {
777        return YES;
778    }
779
780    if (closeCode >= 4000 && closeCode <= 4999) {
781        return YES;
782    }
783
784    return NO;
785}
786
787//  Note from RFC:
788//
789//  If there is a body, the first two
790//  bytes of the body MUST be a 2-byte unsigned integer (in network byte
791//  order) representing a status code with value /code/ defined in
792//  Section 7.4.  Following the 2-byte integer the body MAY contain UTF-8
793//  encoded data with value /reason/, the interpretation of which is not
794//  defined by this specification.
795
796- (void)handleCloseWithData:(NSData *)data;
797{
798    size_t dataSize = data.length;
799    __block uint16_t closeCode = 0;
800
801    SRFastLog(@"Received close frame");
802
803    if (dataSize == 1) {
804        // TODO handle error
805        [self _closeWithProtocolError:@"Payload for close must be larger than 2 bytes"];
806        return;
807    } else if (dataSize >= 2) {
808        [data getBytes:&closeCode length:sizeof(closeCode)];
809        _closeCode = EndianU16_BtoN(closeCode);
810        if (!closeCodeIsValid(_closeCode)) {
811            [self _closeWithProtocolError:[NSString stringWithFormat:@"Cannot have close code of %d", _closeCode]];
812            return;
813        }
814        if (dataSize > 2) {
815            _closeReason = [[NSString alloc] initWithData:[data subdataWithRange:NSMakeRange(2, dataSize - 2)] encoding:NSUTF8StringEncoding];
816            if (!_closeReason) {
817                [self _closeWithProtocolError:@"Close reason MUST be valid UTF-8"];
818                return;
819            }
820        }
821    } else {
822        _closeCode = SRStatusNoStatusReceived;
823    }
824
825    [self assertOnWorkQueue];
826
827    if (self.readyState == SR_OPEN) {
828        [self closeWithCode:1000 reason:nil];
829    }
830    dispatch_async(_workQueue, ^{
831        [self _disconnect];
832    });
833}
834
835- (void)_disconnect;
836{
837    [self assertOnWorkQueue];
838    SRFastLog(@"Trying to disconnect");
839    _closeWhenFinishedWriting = YES;
840    [self _pumpWriting];
841}
842
843- (void)_handleFrameWithData:(NSData *)frameData opCode:(NSInteger)opcode;
844{
845    // Check that the current data is valid UTF8
846
847    BOOL isControlFrame = (opcode == SROpCodePing || opcode == SROpCodePong || opcode == SROpCodeConnectionClose);
848    if (!isControlFrame) {
849        [self _readFrameNew];
850    } else {
851        dispatch_async(_workQueue, ^{
852            [self _readFrameContinue];
853        });
854    }
855
856    switch (opcode) {
857        case SROpCodeTextFrame: {
858            NSString *str = [[NSString alloc] initWithData:frameData encoding:NSUTF8StringEncoding];
859            if (str == nil && frameData) {
860                [self closeWithCode:SRStatusCodeInvalidUTF8 reason:@"Text frames must be valid UTF-8"];
861                dispatch_async(_workQueue, ^{
862                    [self _disconnect];
863                });
864
865                return;
866            }
867            [self _handleMessage:str];
868            break;
869        }
870        case SROpCodeBinaryFrame:
871            [self _handleMessage:[frameData copy]];
872            break;
873        case SROpCodeConnectionClose:
874            [self handleCloseWithData:frameData];
875            break;
876        case SROpCodePing:
877            [self handlePing:frameData];
878            break;
879        case SROpCodePong:
880            [self handlePong:frameData];
881            break;
882        default:
883            [self _closeWithProtocolError:[NSString stringWithFormat:@"Unknown opcode %ld", (long)opcode]];
884            // TODO: Handle invalid opcode
885            break;
886    }
887}
888
889- (void)_handleFrameHeader:(frame_header)frame_header curData:(NSData *)curData;
890{
891    assert(frame_header.opcode != 0);
892
893    if (self.readyState != SR_OPEN) {
894        return;
895    }
896
897
898    BOOL isControlFrame = (frame_header.opcode == SROpCodePing || frame_header.opcode == SROpCodePong || frame_header.opcode == SROpCodeConnectionClose);
899
900    if (isControlFrame && !frame_header.fin) {
901        [self _closeWithProtocolError:@"Fragmented control frames not allowed"];
902        return;
903    }
904
905    if (isControlFrame && frame_header.payload_length >= 126) {
906        [self _closeWithProtocolError:@"Control frames cannot have payloads larger than 126 bytes"];
907        return;
908    }
909
910    if (!isControlFrame) {
911        _currentFrameOpcode = frame_header.opcode;
912        _currentFrameCount += 1;
913    }
914
915    if (frame_header.payload_length == 0) {
916        if (isControlFrame) {
917            [self _handleFrameWithData:curData opCode:frame_header.opcode];
918        } else {
919            if (frame_header.fin) {
920                [self _handleFrameWithData:_currentFrameData opCode:frame_header.opcode];
921            } else {
922                // TODO add assert that opcode is not a control;
923                [self _readFrameContinue];
924            }
925        }
926    } else {
927        assert(frame_header.payload_length <= SIZE_T_MAX);
928        [self _addConsumerWithDataLength:(size_t)frame_header.payload_length callback:^(SRWebSocket *self, NSData *newData) {
929            if (isControlFrame) {
930                [self _handleFrameWithData:newData opCode:frame_header.opcode];
931            } else {
932                if (frame_header.fin) {
933                    [self _handleFrameWithData:self->_currentFrameData opCode:frame_header.opcode];
934                } else {
935                    // TODO add assert that opcode is not a control;
936                    [self _readFrameContinue];
937                }
938
939            }
940        } readToCurrentFrame:!isControlFrame unmaskBytes:frame_header.masked];
941    }
942}
943
944/* From RFC:
945
946 0                   1                   2                   3
947 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
948 +-+-+-+-+-------+-+-------------+-------------------------------+
949 |F|R|R|R| opcode|M| Payload len |    Extended payload length    |
950 |I|S|S|S|  (4)  |A|     (7)     |             (16/64)           |
951 |N|V|V|V|       |S|             |   (if payload len==126/127)   |
952 | |1|2|3|       |K|             |                               |
953 +-+-+-+-+-------+-+-------------+ - - - - - - - - - - - - - - - +
954 |     Extended payload length continued, if payload len == 127  |
955 + - - - - - - - - - - - - - - - +-------------------------------+
956 |                               |Masking-key, if MASK set to 1  |
957 +-------------------------------+-------------------------------+
958 | Masking-key (continued)       |          Payload Data         |
959 +-------------------------------- - - - - - - - - - - - - - - - +
960 :                     Payload Data continued ...                :
961 + - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - +
962 |                     Payload Data continued ...                |
963 +---------------------------------------------------------------+
964 */
965
966static const uint8_t SRFinMask          = 0x80;
967static const uint8_t SROpCodeMask       = 0x0F;
968static const uint8_t SRRsvMask          = 0x70;
969static const uint8_t SRMaskMask         = 0x80;
970static const uint8_t SRPayloadLenMask   = 0x7F;
971
972
973- (void)_readFrameContinue;
974{
975    assert((_currentFrameCount == 0 && _currentFrameOpcode == 0) || (_currentFrameCount > 0 && _currentFrameOpcode > 0));
976
977    [self _addConsumerWithDataLength:2 callback:^(SRWebSocket *self, NSData *data) {
978        __block frame_header header = {0};
979
980        const uint8_t *headerBuffer = data.bytes;
981        assert(data.length >= 2);
982
983        if (headerBuffer[0] & SRRsvMask) {
984            [self _closeWithProtocolError:@"Server used RSV bits"];
985            return;
986        }
987
988        uint8_t receivedOpcode = (SROpCodeMask & headerBuffer[0]);
989
990        BOOL isControlFrame = (receivedOpcode == SROpCodePing || receivedOpcode == SROpCodePong || receivedOpcode == SROpCodeConnectionClose);
991
992        if (!isControlFrame && receivedOpcode != 0 && self->_currentFrameCount > 0) {
993            [self _closeWithProtocolError:@"all data frames after the initial data frame must have opcode 0"];
994            return;
995        }
996
997        if (receivedOpcode == 0 && self->_currentFrameCount == 0) {
998            [self _closeWithProtocolError:@"cannot continue a message"];
999            return;
1000        }
1001
1002        header.opcode = receivedOpcode == 0 ? self->_currentFrameOpcode : receivedOpcode;
1003
1004        header.fin = !!(SRFinMask & headerBuffer[0]);
1005
1006
1007        header.masked = !!(SRMaskMask & headerBuffer[1]);
1008        header.payload_length = SRPayloadLenMask & headerBuffer[1];
1009
1010        headerBuffer = NULL;
1011
1012        if (header.masked) {
1013            [self _closeWithProtocolError:@"Client must receive unmasked data"];
1014        }
1015
1016        size_t extra_bytes_needed = header.masked ? sizeof(_currentReadMaskKey) : 0;
1017
1018        if (header.payload_length == 126) {
1019            extra_bytes_needed += sizeof(uint16_t);
1020        } else if (header.payload_length == 127) {
1021            extra_bytes_needed += sizeof(uint64_t);
1022        }
1023
1024        if (extra_bytes_needed == 0) {
1025            [self _handleFrameHeader:header curData:self->_currentFrameData];
1026        } else {
1027            [self _addConsumerWithDataLength:extra_bytes_needed callback:^(SRWebSocket *self, NSData *data) {
1028                size_t mapped_size = data.length;
1029                const void *mapped_buffer = data.bytes;
1030                size_t offset = 0;
1031
1032                if (header.payload_length == 126) {
1033                    assert(mapped_size >= sizeof(uint16_t));
1034                    uint16_t newLen = EndianU16_BtoN(*(uint16_t *)(mapped_buffer));
1035                    header.payload_length = newLen;
1036                    offset += sizeof(uint16_t);
1037                } else if (header.payload_length == 127) {
1038                    assert(mapped_size >= sizeof(uint64_t));
1039                    header.payload_length = EndianU64_BtoN(*(uint64_t *)(mapped_buffer));
1040                    offset += sizeof(uint64_t);
1041                } else {
1042                    assert(header.payload_length < 126 && header.payload_length >= 0);
1043                }
1044
1045
1046                if (header.masked) {
1047                    assert(mapped_size >= sizeof(_currentReadMaskOffset) + offset);
1048                    memcpy(self->_currentReadMaskKey, ((uint8_t *)mapped_buffer) + offset, sizeof(self->_currentReadMaskKey));
1049                }
1050
1051                [self _handleFrameHeader:header curData:self->_currentFrameData];
1052            } readToCurrentFrame:NO unmaskBytes:NO];
1053        }
1054    } readToCurrentFrame:NO unmaskBytes:NO];
1055}
1056
1057- (void)_readFrameNew;
1058{
1059    dispatch_async(_workQueue, ^{
1060        [_currentFrameData setLength:0];
1061
1062        _currentFrameOpcode = 0;
1063        _currentFrameCount = 0;
1064        _readOpCount = 0;
1065        _currentStringScanPosition = 0;
1066
1067        [self _readFrameContinue];
1068    });
1069}
1070
1071- (void)_pumpWriting;
1072{
1073    [self assertOnWorkQueue];
1074
1075    NSUInteger dataLength = _outputBuffer.length;
1076    if (dataLength - _outputBufferOffset > 0 && _outputStream.hasSpaceAvailable) {
1077        NSInteger bytesWritten = [_outputStream write:_outputBuffer.bytes + _outputBufferOffset maxLength:dataLength - _outputBufferOffset];
1078        if (bytesWritten == -1) {
1079            [self _failWithError:[NSError errorWithDomain:SRWebSocketErrorDomain code:2145 userInfo:[NSDictionary dictionaryWithObject:@"Error writing to stream" forKey:NSLocalizedDescriptionKey]]];
1080             return;
1081        }
1082
1083        _outputBufferOffset += bytesWritten;
1084
1085        if (_outputBufferOffset > 4096 && _outputBufferOffset > (_outputBuffer.length >> 1)) {
1086            _outputBuffer = [[NSMutableData alloc] initWithBytes:(char *)_outputBuffer.bytes + _outputBufferOffset length:_outputBuffer.length - _outputBufferOffset];
1087            _outputBufferOffset = 0;
1088        }
1089    }
1090
1091    if (_closeWhenFinishedWriting &&
1092        _outputBuffer.length - _outputBufferOffset == 0 &&
1093        (_inputStream.streamStatus != NSStreamStatusNotOpen &&
1094         _inputStream.streamStatus != NSStreamStatusClosed) &&
1095        !_sentClose) {
1096        _sentClose = YES;
1097
1098        [_outputStream close];
1099        [_inputStream close];
1100
1101
1102        for (NSArray *runLoop in [_scheduledRunloops copy]) {
1103            [self unscheduleFromRunLoop:[runLoop objectAtIndex:0] forMode:[runLoop objectAtIndex:1]];
1104        }
1105
1106        if (!_failed) {
1107            [self _performDelegateBlock:^{
1108                if ([self.delegate respondsToSelector:@selector(webSocket:didCloseWithCode:reason:wasClean:)]) {
1109                    [self.delegate webSocket:self didCloseWithCode:_closeCode reason:_closeReason wasClean:YES];
1110                }
1111            }];
1112        }
1113
1114        _selfRetain = nil;
1115    }
1116}
1117
1118- (void)_addConsumerWithScanner:(stream_scanner)consumer callback:(data_callback)callback;
1119{
1120    [self assertOnWorkQueue];
1121    [self _addConsumerWithScanner:consumer callback:callback dataLength:0];
1122}
1123
1124- (void)_addConsumerWithDataLength:(size_t)dataLength callback:(data_callback)callback readToCurrentFrame:(BOOL)readToCurrentFrame unmaskBytes:(BOOL)unmaskBytes;
1125{
1126    [self assertOnWorkQueue];
1127    assert(dataLength);
1128
1129    [_consumers addObject:[_consumerPool consumerWithScanner:nil handler:callback bytesNeeded:dataLength readToCurrentFrame:readToCurrentFrame unmaskBytes:unmaskBytes]];
1130    [self _pumpScanner];
1131}
1132
1133- (void)_addConsumerWithScanner:(stream_scanner)consumer callback:(data_callback)callback dataLength:(size_t)dataLength;
1134{
1135    [self assertOnWorkQueue];
1136    [_consumers addObject:[_consumerPool consumerWithScanner:consumer handler:callback bytesNeeded:dataLength readToCurrentFrame:NO unmaskBytes:NO]];
1137    [self _pumpScanner];
1138}
1139
1140
1141static const char CRLFCRLFBytes[] = {'\r', '\n', '\r', '\n'};
1142
1143- (void)_readUntilHeaderCompleteWithCallback:(data_callback)dataHandler;
1144{
1145    [self _readUntilBytes:CRLFCRLFBytes length:sizeof(CRLFCRLFBytes) callback:dataHandler];
1146}
1147
1148- (void)_readUntilBytes:(const void *)bytes length:(size_t)length callback:(data_callback)dataHandler;
1149{
1150    // TODO optimize so this can continue from where we last searched
1151    stream_scanner consumer = ^size_t(NSData *data) {
1152        __block size_t found_size = 0;
1153        __block size_t match_count = 0;
1154
1155        size_t size = data.length;
1156        const unsigned char *buffer = data.bytes;
1157        for (size_t i = 0; i < size; i++ ) {
1158            if (((const unsigned char *)buffer)[i] == ((const unsigned char *)bytes)[match_count]) {
1159                match_count += 1;
1160                if (match_count == length) {
1161                    found_size = i + 1;
1162                    break;
1163                }
1164            } else {
1165                match_count = 0;
1166            }
1167        }
1168        return found_size;
1169    };
1170    [self _addConsumerWithScanner:consumer callback:dataHandler];
1171}
1172
1173
1174// Returns true if did work
1175- (BOOL)_innerPumpScanner {
1176
1177    BOOL didWork = NO;
1178
1179    if (self.readyState >= SR_CLOSING) {
1180        return didWork;
1181    }
1182
1183    if (!_consumers.count) {
1184        return didWork;
1185    }
1186
1187    size_t curSize = _readBuffer.length - _readBufferOffset;
1188    if (!curSize) {
1189        return didWork;
1190    }
1191
1192    SRIOConsumer *consumer = [_consumers objectAtIndex:0];
1193
1194    size_t bytesNeeded = consumer.bytesNeeded;
1195
1196    size_t foundSize = 0;
1197    if (consumer.consumer) {
1198        NSData *tempView = [NSData dataWithBytesNoCopy:(char *)_readBuffer.bytes + _readBufferOffset length:_readBuffer.length - _readBufferOffset freeWhenDone:NO];
1199        foundSize = consumer.consumer(tempView);
1200    } else {
1201        assert(consumer.bytesNeeded);
1202        if (curSize >= bytesNeeded) {
1203            foundSize = bytesNeeded;
1204        } else if (consumer.readToCurrentFrame) {
1205            foundSize = curSize;
1206        }
1207    }
1208
1209    NSData *slice = nil;
1210    if (consumer.readToCurrentFrame || foundSize) {
1211        NSRange sliceRange = NSMakeRange(_readBufferOffset, foundSize);
1212        slice = [_readBuffer subdataWithRange:sliceRange];
1213
1214        _readBufferOffset += foundSize;
1215
1216        if (_readBufferOffset > 4096 && _readBufferOffset > (_readBuffer.length >> 1)) {
1217            _readBuffer = [[NSMutableData alloc] initWithBytes:(char *)_readBuffer.bytes + _readBufferOffset length:_readBuffer.length - _readBufferOffset];            _readBufferOffset = 0;
1218        }
1219
1220        if (consumer.unmaskBytes) {
1221            NSMutableData *mutableSlice = [slice mutableCopy];
1222
1223            NSUInteger len = mutableSlice.length;
1224            uint8_t *bytes = mutableSlice.mutableBytes;
1225
1226            for (NSUInteger i = 0; i < len; i++) {
1227                bytes[i] = bytes[i] ^ _currentReadMaskKey[_currentReadMaskOffset % sizeof(_currentReadMaskKey)];
1228                _currentReadMaskOffset += 1;
1229            }
1230
1231            slice = mutableSlice;
1232        }
1233
1234        if (consumer.readToCurrentFrame) {
1235            [_currentFrameData appendData:slice];
1236
1237            _readOpCount += 1;
1238
1239            if (_currentFrameOpcode == SROpCodeTextFrame) {
1240                // Validate UTF8 stuff.
1241                size_t currentDataSize = _currentFrameData.length;
1242                if (_currentFrameOpcode == SROpCodeTextFrame && currentDataSize > 0) {
1243                    // TODO: Optimize the crap out of this.  Don't really have to copy all the data each time
1244
1245                    size_t scanSize = currentDataSize - _currentStringScanPosition;
1246
1247                    NSData *scan_data = [_currentFrameData subdataWithRange:NSMakeRange(_currentStringScanPosition, scanSize)];
1248                    int32_t valid_utf8_size = validate_dispatch_data_partial_string(scan_data);
1249
1250                    if (valid_utf8_size == -1) {
1251                        [self closeWithCode:SRStatusCodeInvalidUTF8 reason:@"Text frames must be valid UTF-8"];
1252                        dispatch_async(_workQueue, ^{
1253                            [self _disconnect];
1254                        });
1255                        return didWork;
1256                    } else {
1257                        _currentStringScanPosition += valid_utf8_size;
1258                    }
1259                }
1260
1261            }
1262
1263            consumer.bytesNeeded -= foundSize;
1264
1265            if (consumer.bytesNeeded == 0) {
1266                [_consumers removeObjectAtIndex:0];
1267                consumer.handler(self, nil);
1268                [_consumerPool returnConsumer:consumer];
1269                didWork = YES;
1270            }
1271        } else if (foundSize) {
1272            [_consumers removeObjectAtIndex:0];
1273            consumer.handler(self, slice);
1274            [_consumerPool returnConsumer:consumer];
1275            didWork = YES;
1276        }
1277    }
1278    return didWork;
1279}
1280
1281-(void)_pumpScanner;
1282{
1283    [self assertOnWorkQueue];
1284
1285    if (!_isPumping) {
1286        _isPumping = YES;
1287    } else {
1288        return;
1289    }
1290
1291    while ([self _innerPumpScanner]) {
1292
1293    }
1294
1295    _isPumping = NO;
1296}
1297
1298//#define NOMASK
1299
1300static const size_t SRFrameHeaderOverhead = 32;
1301
1302- (void)_sendFrameWithOpcode:(SROpCode)opcode data:(id)data;
1303{
1304    [self assertOnWorkQueue];
1305
1306    if (nil == data) {
1307        return;
1308    }
1309
1310    NSAssert([data isKindOfClass:[NSData class]] || [data isKindOfClass:[NSString class]], @"NSString or NSData");
1311
1312    size_t payloadLength = [data isKindOfClass:[NSString class]] ? [(NSString *)data lengthOfBytesUsingEncoding:NSUTF8StringEncoding] : [data length];
1313
1314    NSMutableData *frame = [[NSMutableData alloc] initWithLength:payloadLength + SRFrameHeaderOverhead];
1315    if (!frame) {
1316        [self closeWithCode:SRStatusCodeMessageTooBig reason:@"Message too big"];
1317        return;
1318    }
1319    uint8_t *frame_buffer = (uint8_t *)[frame mutableBytes];
1320
1321    // set fin
1322    frame_buffer[0] = SRFinMask | opcode;
1323
1324    BOOL useMask = YES;
1325#ifdef NOMASK
1326    useMask = NO;
1327#endif
1328
1329    if (useMask) {
1330    // set the mask and header
1331        frame_buffer[1] |= SRMaskMask;
1332    }
1333
1334    size_t frame_buffer_size = 2;
1335
1336    const uint8_t *unmasked_payload = NULL;
1337    if ([data isKindOfClass:[NSData class]]) {
1338        unmasked_payload = (uint8_t *)[data bytes];
1339    } else if ([data isKindOfClass:[NSString class]]) {
1340        unmasked_payload =  (const uint8_t *)[data UTF8String];
1341    } else {
1342        return;
1343    }
1344
1345    if (payloadLength < 126) {
1346        frame_buffer[1] |= payloadLength;
1347    } else if (payloadLength <= UINT16_MAX) {
1348        frame_buffer[1] |= 126;
1349        *((uint16_t *)(frame_buffer + frame_buffer_size)) = EndianU16_BtoN((uint16_t)payloadLength);
1350        frame_buffer_size += sizeof(uint16_t);
1351    } else {
1352        frame_buffer[1] |= 127;
1353        *((uint64_t *)(frame_buffer + frame_buffer_size)) = EndianU64_BtoN((uint64_t)payloadLength);
1354        frame_buffer_size += sizeof(uint64_t);
1355    }
1356
1357    if (!useMask) {
1358        for (size_t i = 0; i < payloadLength; i++) {
1359            frame_buffer[frame_buffer_size] = unmasked_payload[i];
1360            frame_buffer_size += 1;
1361        }
1362    } else {
1363        uint8_t *mask_key = frame_buffer + frame_buffer_size;
1364        SecRandomCopyBytes(kSecRandomDefault, sizeof(uint32_t), (uint8_t *)mask_key);
1365        frame_buffer_size += sizeof(uint32_t);
1366
1367        // TODO: could probably optimize this with SIMD
1368        for (size_t i = 0; i < payloadLength; i++) {
1369            frame_buffer[frame_buffer_size] = unmasked_payload[i] ^ mask_key[i % sizeof(uint32_t)];
1370            frame_buffer_size += 1;
1371        }
1372    }
1373
1374    assert(frame_buffer_size <= [frame length]);
1375    frame.length = frame_buffer_size;
1376
1377    [self _writeData:frame];
1378}
1379
1380- (void)stream:(NSStream *)aStream handleEvent:(NSStreamEvent)eventCode;
1381{
1382    if (_secure && !_pinnedCertFound && (eventCode == NSStreamEventHasBytesAvailable || eventCode == NSStreamEventHasSpaceAvailable)) {
1383
1384        NSArray *sslCerts = [_urlRequest SR_SSLPinnedCertificates];
1385        if (sslCerts) {
1386            SecTrustRef secTrust = (__bridge SecTrustRef)[aStream propertyForKey:(__bridge id)kCFStreamPropertySSLPeerTrust];
1387            if (secTrust) {
1388                NSInteger numCerts = SecTrustGetCertificateCount(secTrust);
1389                for (NSInteger i = 0; i < numCerts && !_pinnedCertFound; i++) {
1390                    SecCertificateRef cert = SecTrustGetCertificateAtIndex(secTrust, i);
1391                    NSData *certData = CFBridgingRelease(SecCertificateCopyData(cert));
1392
1393                    for (id ref in sslCerts) {
1394                        SecCertificateRef trustedCert = (__bridge SecCertificateRef)ref;
1395                        NSData *trustedCertData = CFBridgingRelease(SecCertificateCopyData(trustedCert));
1396
1397                        if ([trustedCertData isEqualToData:certData]) {
1398                            _pinnedCertFound = YES;
1399                            break;
1400                        }
1401                    }
1402                }
1403            }
1404
1405            if (!_pinnedCertFound) {
1406                dispatch_async(_workQueue, ^{
1407                    [self _failWithError:[NSError errorWithDomain:SRWebSocketErrorDomain code:23556 userInfo:[NSDictionary dictionaryWithObject:[NSString stringWithFormat:@"Invalid server cert"] forKey:NSLocalizedDescriptionKey]]];
1408                });
1409                return;
1410            }
1411        }
1412    }
1413
1414    dispatch_async(_workQueue, ^{
1415        switch (eventCode) {
1416            case NSStreamEventOpenCompleted: {
1417                SRFastLog(@"NSStreamEventOpenCompleted %@", aStream);
1418                if (self.readyState >= SR_CLOSING) {
1419                    return;
1420                }
1421                assert(_readBuffer);
1422
1423                if (self.readyState == SR_CONNECTING && aStream == _inputStream) {
1424                    [self didConnect];
1425                }
1426                [self _pumpWriting];
1427                [self _pumpScanner];
1428                break;
1429            }
1430
1431            case NSStreamEventErrorOccurred: {
1432                SRFastLog(@"NSStreamEventErrorOccurred %@ %@", aStream, [[aStream streamError] copy]);
1433                /// TODO specify error better!
1434                [self _failWithError:aStream.streamError];
1435                _readBufferOffset = 0;
1436                [_readBuffer setLength:0];
1437                break;
1438
1439            }
1440
1441            case NSStreamEventEndEncountered: {
1442                [self _pumpScanner];
1443                SRFastLog(@"NSStreamEventEndEncountered %@", aStream);
1444                if (aStream.streamError) {
1445                    [self _failWithError:aStream.streamError];
1446                } else {
1447                    if (self.readyState != SR_CLOSED) {
1448                        self.readyState = SR_CLOSED;
1449                        _selfRetain = nil;
1450                    }
1451
1452                    if (!_sentClose && !_failed) {
1453                        _sentClose = YES;
1454                        // If we get closed in this state it's probably not clean because we should be sending this when we send messages
1455                        [self _performDelegateBlock:^{
1456                            if ([self.delegate respondsToSelector:@selector(webSocket:didCloseWithCode:reason:wasClean:)]) {
1457                                [self.delegate webSocket:self didCloseWithCode:SRStatusCodeGoingAway reason:@"Stream end encountered" wasClean:NO];
1458                            }
1459                        }];
1460                    }
1461                }
1462
1463                break;
1464            }
1465
1466            case NSStreamEventHasBytesAvailable: {
1467                SRFastLog(@"NSStreamEventHasBytesAvailable %@", aStream);
1468                const int bufferSize = 2048;
1469                uint8_t buffer[bufferSize];
1470
1471                while (_inputStream.hasBytesAvailable) {
1472                    NSInteger bytes_read = [_inputStream read:buffer maxLength:bufferSize];
1473
1474                    if (bytes_read > 0) {
1475                        [_readBuffer appendBytes:buffer length:bytes_read];
1476                    } else if (bytes_read < 0) {
1477                        [self _failWithError:_inputStream.streamError];
1478                    }
1479
1480                    if (bytes_read != bufferSize) {
1481                        break;
1482                    }
1483                };
1484                [self _pumpScanner];
1485                break;
1486            }
1487
1488            case NSStreamEventHasSpaceAvailable: {
1489                SRFastLog(@"NSStreamEventHasSpaceAvailable %@", aStream);
1490                [self _pumpWriting];
1491                break;
1492            }
1493
1494            default:
1495                SRFastLog(@"(default)  %@", aStream);
1496                break;
1497        }
1498    });
1499}
1500
1501@end
1502
1503
1504@implementation SRIOConsumer
1505
1506@synthesize bytesNeeded = _bytesNeeded;
1507@synthesize consumer = _scanner;
1508@synthesize handler = _handler;
1509@synthesize readToCurrentFrame = _readToCurrentFrame;
1510@synthesize unmaskBytes = _unmaskBytes;
1511
1512- (void)setupWithScanner:(stream_scanner)scanner handler:(data_callback)handler bytesNeeded:(size_t)bytesNeeded readToCurrentFrame:(BOOL)readToCurrentFrame unmaskBytes:(BOOL)unmaskBytes;
1513{
1514    _scanner = [scanner copy];
1515    _handler = [handler copy];
1516    _bytesNeeded = bytesNeeded;
1517    _readToCurrentFrame = readToCurrentFrame;
1518    _unmaskBytes = unmaskBytes;
1519    assert(_scanner || _bytesNeeded);
1520}
1521
1522
1523@end
1524
1525
1526@implementation SRIOConsumerPool {
1527    NSUInteger _poolSize;
1528    NSMutableArray *_bufferedConsumers;
1529}
1530
1531- (id)initWithBufferCapacity:(NSUInteger)poolSize;
1532{
1533    self = [super init];
1534    if (self) {
1535        _poolSize = poolSize;
1536        _bufferedConsumers = [[NSMutableArray alloc] initWithCapacity:poolSize];
1537    }
1538    return self;
1539}
1540
1541- (id)init
1542{
1543    return [self initWithBufferCapacity:8];
1544}
1545
1546- (SRIOConsumer *)consumerWithScanner:(stream_scanner)scanner handler:(data_callback)handler bytesNeeded:(size_t)bytesNeeded readToCurrentFrame:(BOOL)readToCurrentFrame unmaskBytes:(BOOL)unmaskBytes;
1547{
1548    SRIOConsumer *consumer = nil;
1549    if (_bufferedConsumers.count) {
1550        consumer = [_bufferedConsumers lastObject];
1551        [_bufferedConsumers removeLastObject];
1552    } else {
1553        consumer = [[SRIOConsumer alloc] init];
1554    }
1555
1556    [consumer setupWithScanner:scanner handler:handler bytesNeeded:bytesNeeded readToCurrentFrame:readToCurrentFrame unmaskBytes:unmaskBytes];
1557
1558    return consumer;
1559}
1560
1561- (void)returnConsumer:(SRIOConsumer *)consumer;
1562{
1563    if (_bufferedConsumers.count < _poolSize) {
1564        [_bufferedConsumers addObject:consumer];
1565    }
1566}
1567
1568@end
1569
1570
1571@implementation  NSURLRequest (CertificateAdditions)
1572
1573- (NSArray *)SR_SSLPinnedCertificates;
1574{
1575    return [NSURLProtocol propertyForKey:@"SR_SSLPinnedCertificates" inRequest:self];
1576}
1577
1578@end
1579
1580@implementation  NSMutableURLRequest (CertificateAdditions)
1581
1582- (NSArray *)SR_SSLPinnedCertificates;
1583{
1584    return [NSURLProtocol propertyForKey:@"SR_SSLPinnedCertificates" inRequest:self];
1585}
1586
1587- (void)setSR_SSLPinnedCertificates:(NSArray *)SR_SSLPinnedCertificates;
1588{
1589    [NSURLProtocol setProperty:SR_SSLPinnedCertificates forKey:@"SR_SSLPinnedCertificates" inRequest:self];
1590}
1591
1592@end
1593
1594@implementation NSURL (SRWebSocket)
1595
1596- (NSString *)SR_origin;
1597{
1598    NSString *scheme = [self.scheme lowercaseString];
1599
1600    if ([scheme isEqualToString:@"wss"]) {
1601        scheme = @"https";
1602    } else if ([scheme isEqualToString:@"ws"]) {
1603        scheme = @"http";
1604    }
1605
1606    if (self.port) {
1607        return [NSString stringWithFormat:@"%@://%@:%@/", scheme, self.host, self.port];
1608    } else {
1609        return [NSString stringWithFormat:@"%@://%@/", scheme, self.host];
1610    }
1611}
1612
1613@end
1614
1615//#define SR_ENABLE_LOG
1616
1617static inline void SRFastLog(NSString *format, ...)  {
1618#ifdef SR_ENABLE_LOG
1619    __block va_list arg_list;
1620    va_start (arg_list, format);
1621
1622    NSString *formattedString = [[NSString alloc] initWithFormat:format arguments:arg_list];
1623
1624    va_end(arg_list);
1625
1626    NSLog(@"[SR] %@", formattedString);
1627#endif
1628}
1629
1630
1631#ifdef HAS_ICU
1632
1633static inline int32_t validate_dispatch_data_partial_string(NSData *data) {
1634    if ([data length] > INT32_MAX) {
1635        // INT32_MAX is the limit so long as this Framework is using 32 bit ints everywhere.
1636        return -1;
1637    }
1638
1639    int32_t size = (int32_t)[data length];
1640
1641    const void * contents = [data bytes];
1642    const uint8_t *str = (const uint8_t *)contents;
1643
1644    UChar32 codepoint = 1;
1645    int32_t offset = 0;
1646    int32_t lastOffset = 0;
1647    while(offset < size && codepoint > 0)  {
1648        lastOffset = offset;
1649        U8_NEXT(str, offset, size, codepoint);
1650    }
1651
1652    if (codepoint == -1) {
1653        // Check to see if the last byte is valid or whether it was just continuing
1654        if (!U8_IS_LEAD(str[lastOffset]) || U8_COUNT_TRAIL_BYTES(str[lastOffset]) + lastOffset < (int32_t)size) {
1655
1656            size = -1;
1657        } else {
1658            uint8_t leadByte = str[lastOffset];
1659            U8_MASK_LEAD_BYTE(leadByte, U8_COUNT_TRAIL_BYTES(leadByte));
1660
1661            for (int i = lastOffset + 1; i < offset; i++) {
1662                if (U8_IS_SINGLE(str[i]) || U8_IS_LEAD(str[i]) || !U8_IS_TRAIL(str[i])) {
1663                    size = -1;
1664                }
1665            }
1666
1667            if (size != -1) {
1668                size = lastOffset;
1669            }
1670        }
1671    }
1672
1673    if (size != -1 && ![[NSString alloc] initWithBytesNoCopy:(char *)[data bytes] length:size encoding:NSUTF8StringEncoding freeWhenDone:NO]) {
1674        size = -1;
1675    }
1676
1677    return size;
1678}
1679
1680#else
1681
1682// This is a hack, and probably not optimal
1683static inline int32_t validate_dispatch_data_partial_string(NSData *data) {
1684    static const int maxCodepointSize = 3;
1685
1686    for (int i = 0; i < maxCodepointSize; i++) {
1687        NSString *str = [[NSString alloc] initWithBytesNoCopy:(char *)data.bytes length:data.length - i encoding:NSUTF8StringEncoding freeWhenDone:NO];
1688        if (str) {
1689            return data.length - i;
1690        }
1691    }
1692
1693    return -1;
1694}
1695
1696#endif
1697
1698static _SRRunLoopThread *networkThread = nil;
1699static NSRunLoop *networkRunLoop = nil;
1700
1701@implementation NSRunLoop (SRWebSocket)
1702
1703+ (NSRunLoop *)SR_networkRunLoop {
1704    static dispatch_once_t onceToken;
1705    dispatch_once(&onceToken, ^{
1706        networkThread = [[_SRRunLoopThread alloc] init];
1707        networkThread.name = @"com.squareup.SocketRocket.NetworkThread";
1708        [networkThread start];
1709        networkRunLoop = networkThread.runLoop;
1710    });
1711
1712    return networkRunLoop;
1713}
1714
1715@end
1716
1717
1718@implementation _SRRunLoopThread {
1719    dispatch_group_t _waitGroup;
1720}
1721
1722@synthesize runLoop = _runLoop;
1723
1724- (void)dealloc
1725{
1726    sr_dispatch_release(_waitGroup);
1727}
1728
1729- (id)init
1730{
1731    self = [super init];
1732    if (self) {
1733        _waitGroup = dispatch_group_create();
1734        dispatch_group_enter(_waitGroup);
1735    }
1736    return self;
1737}
1738
1739- (void)main;
1740{
1741    @autoreleasepool {
1742        _runLoop = [NSRunLoop currentRunLoop];
1743        dispatch_group_leave(_waitGroup);
1744
1745        NSTimer *timer = [[NSTimer alloc] initWithFireDate:[NSDate distantFuture] interval:0.0 target:nil selector:nil userInfo:nil repeats:NO];
1746        [_runLoop addTimer:timer forMode:NSDefaultRunLoopMode];
1747
1748        while ([_runLoop runMode:NSDefaultRunLoopMode beforeDate:[NSDate distantFuture]]) {
1749
1750        }
1751        assert(NO);
1752    }
1753}
1754
1755- (NSRunLoop *)runLoop;
1756{
1757    dispatch_group_wait(_waitGroup, DISPATCH_TIME_FOREVER);
1758    return _runLoop;
1759}
1760
1761@end
1762