PageRenderTime 47ms CodeModel.GetById 14ms RepoModel.GetById 1ms app.codeStats 0ms

/Pods/SocketRocket/SocketRocket/SRWebSocket.m

https://gitlab.com/trungminhnt/sampleShinobi
Objective C | 1511 lines | 1097 code | 334 blank | 80 comment | 257 complexity | 17060def50942d84a56e6d34449cb9ba MD5 | raw file
  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. #import "SRWebSocket.h"
  17. #if TARGET_OS_IPHONE
  18. #define HAS_ICU
  19. #endif
  20. #ifdef HAS_ICU
  21. #import <unicode/utf8.h>
  22. #endif
  23. #if TARGET_OS_IPHONE
  24. #import <Endian.h>
  25. #else
  26. #import <CoreServices/CoreServices.h>
  27. #endif
  28. #import <CommonCrypto/CommonDigest.h>
  29. #import <Security/SecRandom.h>
  30. #if OS_OBJECT_USE_OBJC_RETAIN_RELEASE
  31. #define sr_dispatch_retain(x)
  32. #define sr_dispatch_release(x)
  33. #define maybe_bridge(x) ((__bridge void *) x)
  34. #else
  35. #define sr_dispatch_retain(x) dispatch_retain(x)
  36. #define sr_dispatch_release(x) dispatch_release(x)
  37. #define maybe_bridge(x) (x)
  38. #endif
  39. #if !__has_feature(objc_arc)
  40. #error SocketRocket must be compiled with ARC enabled
  41. #endif
  42. typedef enum {
  43. SROpCodeTextFrame = 0x1,
  44. SROpCodeBinaryFrame = 0x2,
  45. // 3-7 reserved.
  46. SROpCodeConnectionClose = 0x8,
  47. SROpCodePing = 0x9,
  48. SROpCodePong = 0xA,
  49. // B-F reserved.
  50. } SROpCode;
  51. typedef struct {
  52. BOOL fin;
  53. // BOOL rsv1;
  54. // BOOL rsv2;
  55. // BOOL rsv3;
  56. uint8_t opcode;
  57. BOOL masked;
  58. uint64_t payload_length;
  59. } frame_header;
  60. static NSString *const SRWebSocketAppendToSecKeyString = @"258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
  61. static inline int32_t validate_dispatch_data_partial_string(NSData *data);
  62. static inline void SRFastLog(NSString *format, ...);
  63. @interface NSData (SRWebSocket)
  64. - (NSString *)stringBySHA1ThenBase64Encoding;
  65. @end
  66. @interface NSString (SRWebSocket)
  67. - (NSString *)stringBySHA1ThenBase64Encoding;
  68. @end
  69. @interface NSURL (SRWebSocket)
  70. // The origin isn't really applicable for a native application.
  71. // So instead, just map ws -> http and wss -> https.
  72. - (NSString *)SR_origin;
  73. @end
  74. @interface _SRRunLoopThread : NSThread
  75. @property (nonatomic, readonly) NSRunLoop *runLoop;
  76. @end
  77. static NSString *newSHA1String(const char *bytes, size_t length) {
  78. uint8_t md[CC_SHA1_DIGEST_LENGTH];
  79. assert(length >= 0);
  80. assert(length <= UINT32_MAX);
  81. CC_SHA1(bytes, (CC_LONG)length, md);
  82. NSData *data = [NSData dataWithBytes:md length:CC_SHA1_DIGEST_LENGTH];
  83. if ([data respondsToSelector:@selector(base64EncodedStringWithOptions:)]) {
  84. return [data base64EncodedStringWithOptions:0];
  85. }
  86. #pragma clang diagnostic push
  87. #pragma clang diagnostic ignored "-Wdeprecated-declarations"
  88. return [data base64Encoding];
  89. #pragma clang diagnostic pop
  90. }
  91. @implementation NSData (SRWebSocket)
  92. - (NSString *)stringBySHA1ThenBase64Encoding;
  93. {
  94. return newSHA1String(self.bytes, self.length);
  95. }
  96. @end
  97. @implementation NSString (SRWebSocket)
  98. - (NSString *)stringBySHA1ThenBase64Encoding;
  99. {
  100. return newSHA1String(self.UTF8String, self.length);
  101. }
  102. @end
  103. NSString *const SRWebSocketErrorDomain = @"SRWebSocketErrorDomain";
  104. NSString *const SRHTTPResponseErrorKey = @"HTTPResponseStatusCode";
  105. // Returns number of bytes consumed. Returning 0 means you didn't match.
  106. // Sends bytes to callback handler;
  107. typedef size_t (^stream_scanner)(NSData *collected_data);
  108. typedef void (^data_callback)(SRWebSocket *webSocket, NSData *data);
  109. @interface SRIOConsumer : NSObject {
  110. stream_scanner _scanner;
  111. data_callback _handler;
  112. size_t _bytesNeeded;
  113. BOOL _readToCurrentFrame;
  114. BOOL _unmaskBytes;
  115. }
  116. @property (nonatomic, copy, readonly) stream_scanner consumer;
  117. @property (nonatomic, copy, readonly) data_callback handler;
  118. @property (nonatomic, assign) size_t bytesNeeded;
  119. @property (nonatomic, assign, readonly) BOOL readToCurrentFrame;
  120. @property (nonatomic, assign, readonly) BOOL unmaskBytes;
  121. @end
  122. // This class is not thread-safe, and is expected to always be run on the same queue.
  123. @interface SRIOConsumerPool : NSObject
  124. - (id)initWithBufferCapacity:(NSUInteger)poolSize;
  125. - (SRIOConsumer *)consumerWithScanner:(stream_scanner)scanner handler:(data_callback)handler bytesNeeded:(size_t)bytesNeeded readToCurrentFrame:(BOOL)readToCurrentFrame unmaskBytes:(BOOL)unmaskBytes;
  126. - (void)returnConsumer:(SRIOConsumer *)consumer;
  127. @end
  128. @interface SRWebSocket () <NSStreamDelegate>
  129. @property (nonatomic) SRReadyState readyState;
  130. @property (nonatomic) NSOperationQueue *delegateOperationQueue;
  131. @property (nonatomic) dispatch_queue_t delegateDispatchQueue;
  132. @end
  133. @implementation SRWebSocket {
  134. NSInteger _webSocketVersion;
  135. NSOperationQueue *_delegateOperationQueue;
  136. dispatch_queue_t _delegateDispatchQueue;
  137. dispatch_queue_t _workQueue;
  138. NSMutableArray *_consumers;
  139. NSInputStream *_inputStream;
  140. NSOutputStream *_outputStream;
  141. NSMutableData *_readBuffer;
  142. NSUInteger _readBufferOffset;
  143. NSMutableData *_outputBuffer;
  144. NSUInteger _outputBufferOffset;
  145. uint8_t _currentFrameOpcode;
  146. size_t _currentFrameCount;
  147. size_t _readOpCount;
  148. uint32_t _currentStringScanPosition;
  149. NSMutableData *_currentFrameData;
  150. NSString *_closeReason;
  151. NSString *_secKey;
  152. BOOL _pinnedCertFound;
  153. uint8_t _currentReadMaskKey[4];
  154. size_t _currentReadMaskOffset;
  155. BOOL _consumerStopped;
  156. BOOL _closeWhenFinishedWriting;
  157. BOOL _failed;
  158. BOOL _secure;
  159. NSURLRequest *_urlRequest;
  160. CFHTTPMessageRef _receivedHTTPHeaders;
  161. BOOL _sentClose;
  162. BOOL _didFail;
  163. int _closeCode;
  164. BOOL _isPumping;
  165. NSMutableSet *_scheduledRunloops;
  166. // We use this to retain ourselves.
  167. __strong SRWebSocket *_selfRetain;
  168. NSArray *_requestedProtocols;
  169. SRIOConsumerPool *_consumerPool;
  170. }
  171. @synthesize delegate = _delegate;
  172. @synthesize url = _url;
  173. @synthesize readyState = _readyState;
  174. @synthesize protocol = _protocol;
  175. static __strong NSData *CRLFCRLF;
  176. + (void)initialize;
  177. {
  178. CRLFCRLF = [[NSData alloc] initWithBytes:"\r\n\r\n" length:4];
  179. }
  180. - (id)initWithURLRequest:(NSURLRequest *)request protocols:(NSArray *)protocols;
  181. {
  182. self = [super init];
  183. if (self) {
  184. assert(request.URL);
  185. _url = request.URL;
  186. _urlRequest = request;
  187. _requestedProtocols = [protocols copy];
  188. [self _SR_commonInit];
  189. }
  190. return self;
  191. }
  192. - (id)initWithURLRequest:(NSURLRequest *)request;
  193. {
  194. return [self initWithURLRequest:request protocols:nil];
  195. }
  196. - (id)initWithURL:(NSURL *)url;
  197. {
  198. return [self initWithURL:url protocols:nil];
  199. }
  200. - (id)initWithURL:(NSURL *)url protocols:(NSArray *)protocols;
  201. {
  202. NSMutableURLRequest *request = [[NSMutableURLRequest alloc] initWithURL:url];
  203. return [self initWithURLRequest:request protocols:protocols];
  204. }
  205. - (void)_SR_commonInit;
  206. {
  207. NSString *scheme = _url.scheme.lowercaseString;
  208. assert([scheme isEqualToString:@"ws"] || [scheme isEqualToString:@"http"] || [scheme isEqualToString:@"wss"] || [scheme isEqualToString:@"https"]);
  209. if ([scheme isEqualToString:@"wss"] || [scheme isEqualToString:@"https"]) {
  210. _secure = YES;
  211. }
  212. _readyState = SR_CONNECTING;
  213. _consumerStopped = YES;
  214. _webSocketVersion = 13;
  215. _workQueue = dispatch_queue_create(NULL, DISPATCH_QUEUE_SERIAL);
  216. // Going to set a specific on the queue so we can validate we're on the work queue
  217. dispatch_queue_set_specific(_workQueue, (__bridge void *)self, maybe_bridge(_workQueue), NULL);
  218. _delegateDispatchQueue = dispatch_get_main_queue();
  219. sr_dispatch_retain(_delegateDispatchQueue);
  220. _readBuffer = [[NSMutableData alloc] init];
  221. _outputBuffer = [[NSMutableData alloc] init];
  222. _currentFrameData = [[NSMutableData alloc] init];
  223. _consumers = [[NSMutableArray alloc] init];
  224. _consumerPool = [[SRIOConsumerPool alloc] init];
  225. _scheduledRunloops = [[NSMutableSet alloc] init];
  226. [self _initializeStreams];
  227. // default handlers
  228. }
  229. - (void)assertOnWorkQueue;
  230. {
  231. assert(dispatch_get_specific((__bridge void *)self) == maybe_bridge(_workQueue));
  232. }
  233. - (void)dealloc
  234. {
  235. _inputStream.delegate = nil;
  236. _outputStream.delegate = nil;
  237. [_inputStream close];
  238. [_outputStream close];
  239. sr_dispatch_release(_workQueue);
  240. _workQueue = NULL;
  241. if (_receivedHTTPHeaders) {
  242. CFRelease(_receivedHTTPHeaders);
  243. _receivedHTTPHeaders = NULL;
  244. }
  245. if (_delegateDispatchQueue) {
  246. sr_dispatch_release(_delegateDispatchQueue);
  247. _delegateDispatchQueue = NULL;
  248. }
  249. }
  250. #ifndef NDEBUG
  251. - (void)setReadyState:(SRReadyState)aReadyState;
  252. {
  253. [self willChangeValueForKey:@"readyState"];
  254. assert(aReadyState > _readyState);
  255. _readyState = aReadyState;
  256. [self didChangeValueForKey:@"readyState"];
  257. }
  258. #endif
  259. - (void)open;
  260. {
  261. assert(_url);
  262. NSAssert(_readyState == SR_CONNECTING, @"Cannot call -(void)open on SRWebSocket more than once");
  263. _selfRetain = self;
  264. [self openConnection];
  265. }
  266. // Calls block on delegate queue
  267. - (void)_performDelegateBlock:(dispatch_block_t)block;
  268. {
  269. if (_delegateOperationQueue) {
  270. [_delegateOperationQueue addOperationWithBlock:block];
  271. } else {
  272. assert(_delegateDispatchQueue);
  273. dispatch_async(_delegateDispatchQueue, block);
  274. }
  275. }
  276. - (void)setDelegateDispatchQueue:(dispatch_queue_t)queue;
  277. {
  278. if (queue) {
  279. sr_dispatch_retain(queue);
  280. }
  281. if (_delegateDispatchQueue) {
  282. sr_dispatch_release(_delegateDispatchQueue);
  283. }
  284. _delegateDispatchQueue = queue;
  285. }
  286. - (BOOL)_checkHandshake:(CFHTTPMessageRef)httpMessage;
  287. {
  288. NSString *acceptHeader = CFBridgingRelease(CFHTTPMessageCopyHeaderFieldValue(httpMessage, CFSTR("Sec-WebSocket-Accept")));
  289. if (acceptHeader == nil) {
  290. return NO;
  291. }
  292. NSString *concattedString = [_secKey stringByAppendingString:SRWebSocketAppendToSecKeyString];
  293. NSString *expectedAccept = [concattedString stringBySHA1ThenBase64Encoding];
  294. return [acceptHeader isEqualToString:expectedAccept];
  295. }
  296. - (void)_HTTPHeadersDidFinish;
  297. {
  298. NSInteger responseCode = CFHTTPMessageGetResponseStatusCode(_receivedHTTPHeaders);
  299. if (responseCode >= 400) {
  300. SRFastLog(@"Request failed with response code %d", responseCode);
  301. [self _failWithError:[NSError errorWithDomain:SRWebSocketErrorDomain code:2132 userInfo:@{NSLocalizedDescriptionKey:[NSString stringWithFormat:@"received bad response code from server %ld", (long)responseCode], SRHTTPResponseErrorKey:@(responseCode)}]];
  302. return;
  303. }
  304. if(![self _checkHandshake:_receivedHTTPHeaders]) {
  305. [self _failWithError:[NSError errorWithDomain:SRWebSocketErrorDomain code:2133 userInfo:[NSDictionary dictionaryWithObject:[NSString stringWithFormat:@"Invalid Sec-WebSocket-Accept response"] forKey:NSLocalizedDescriptionKey]]];
  306. return;
  307. }
  308. NSString *negotiatedProtocol = CFBridgingRelease(CFHTTPMessageCopyHeaderFieldValue(_receivedHTTPHeaders, CFSTR("Sec-WebSocket-Protocol")));
  309. if (negotiatedProtocol) {
  310. // Make sure we requested the protocol
  311. if ([_requestedProtocols indexOfObject:negotiatedProtocol] == NSNotFound) {
  312. [self _failWithError:[NSError errorWithDomain:SRWebSocketErrorDomain code:2133 userInfo:[NSDictionary dictionaryWithObject:[NSString stringWithFormat:@"Server specified Sec-WebSocket-Protocol that wasn't requested"] forKey:NSLocalizedDescriptionKey]]];
  313. return;
  314. }
  315. _protocol = negotiatedProtocol;
  316. }
  317. self.readyState = SR_OPEN;
  318. if (!_didFail) {
  319. [self _readFrameNew];
  320. }
  321. [self _performDelegateBlock:^{
  322. if ([self.delegate respondsToSelector:@selector(webSocketDidOpen:)]) {
  323. [self.delegate webSocketDidOpen:self];
  324. };
  325. }];
  326. }
  327. - (void)_readHTTPHeader;
  328. {
  329. if (_receivedHTTPHeaders == NULL) {
  330. _receivedHTTPHeaders = CFHTTPMessageCreateEmpty(NULL, NO);
  331. }
  332. [self _readUntilHeaderCompleteWithCallback:^(SRWebSocket *self, NSData *data) {
  333. CFHTTPMessageAppendBytes(_receivedHTTPHeaders, (const UInt8 *)data.bytes, data.length);
  334. if (CFHTTPMessageIsHeaderComplete(_receivedHTTPHeaders)) {
  335. SRFastLog(@"Finished reading headers %@", CFBridgingRelease(CFHTTPMessageCopyAllHeaderFields(_receivedHTTPHeaders)));
  336. [self _HTTPHeadersDidFinish];
  337. } else {
  338. [self _readHTTPHeader];
  339. }
  340. }];
  341. }
  342. - (void)didConnect
  343. {
  344. SRFastLog(@"Connected");
  345. CFHTTPMessageRef request = CFHTTPMessageCreateRequest(NULL, CFSTR("GET"), (__bridge CFURLRef)_url, kCFHTTPVersion1_1);
  346. // Set host first so it defaults
  347. CFHTTPMessageSetHeaderFieldValue(request, CFSTR("Host"), (__bridge CFStringRef)(_url.port ? [NSString stringWithFormat:@"%@:%@", _url.host, _url.port] : _url.host));
  348. NSMutableData *keyBytes = [[NSMutableData alloc] initWithLength:16];
  349. SecRandomCopyBytes(kSecRandomDefault, keyBytes.length, keyBytes.mutableBytes);
  350. if ([keyBytes respondsToSelector:@selector(base64EncodedStringWithOptions:)]) {
  351. _secKey = [keyBytes base64EncodedStringWithOptions:0];
  352. } else {
  353. #pragma clang diagnostic push
  354. #pragma clang diagnostic ignored "-Wdeprecated-declarations"
  355. _secKey = [keyBytes base64Encoding];
  356. #pragma clang diagnostic pop
  357. }
  358. assert([_secKey length] == 24);
  359. CFHTTPMessageSetHeaderFieldValue(request, CFSTR("Upgrade"), CFSTR("websocket"));
  360. CFHTTPMessageSetHeaderFieldValue(request, CFSTR("Connection"), CFSTR("Upgrade"));
  361. CFHTTPMessageSetHeaderFieldValue(request, CFSTR("Sec-WebSocket-Key"), (__bridge CFStringRef)_secKey);
  362. CFHTTPMessageSetHeaderFieldValue(request, CFSTR("Sec-WebSocket-Version"), (__bridge CFStringRef)[NSString stringWithFormat:@"%ld", (long)_webSocketVersion]);
  363. CFHTTPMessageSetHeaderFieldValue(request, CFSTR("Origin"), (__bridge CFStringRef)_url.SR_origin);
  364. if (_requestedProtocols) {
  365. CFHTTPMessageSetHeaderFieldValue(request, CFSTR("Sec-WebSocket-Protocol"), (__bridge CFStringRef)[_requestedProtocols componentsJoinedByString:@", "]);
  366. }
  367. [_urlRequest.allHTTPHeaderFields enumerateKeysAndObjectsUsingBlock:^(id key, id obj, BOOL *stop) {
  368. CFHTTPMessageSetHeaderFieldValue(request, (__bridge CFStringRef)key, (__bridge CFStringRef)obj);
  369. }];
  370. NSData *message = CFBridgingRelease(CFHTTPMessageCopySerializedMessage(request));
  371. CFRelease(request);
  372. [self _writeData:message];
  373. [self _readHTTPHeader];
  374. }
  375. - (void)_initializeStreams;
  376. {
  377. assert(_url.port.unsignedIntValue <= UINT32_MAX);
  378. uint32_t port = _url.port.unsignedIntValue;
  379. if (port == 0) {
  380. if (!_secure) {
  381. port = 80;
  382. } else {
  383. port = 443;
  384. }
  385. }
  386. NSString *host = _url.host;
  387. CFReadStreamRef readStream = NULL;
  388. CFWriteStreamRef writeStream = NULL;
  389. CFStreamCreatePairWithSocketToHost(NULL, (__bridge CFStringRef)host, port, &readStream, &writeStream);
  390. _outputStream = CFBridgingRelease(writeStream);
  391. _inputStream = CFBridgingRelease(readStream);
  392. if (_secure) {
  393. NSMutableDictionary *SSLOptions = [[NSMutableDictionary alloc] init];
  394. [_outputStream setProperty:(__bridge id)kCFStreamSocketSecurityLevelNegotiatedSSL forKey:(__bridge id)kCFStreamPropertySocketSecurityLevel];
  395. // If we're using pinned certs, don't validate the certificate chain
  396. if ([_urlRequest SR_SSLPinnedCertificates].count) {
  397. [SSLOptions setValue:[NSNumber numberWithBool:NO] forKey:(__bridge id)kCFStreamSSLValidatesCertificateChain];
  398. }
  399. #if DEBUG
  400. [SSLOptions setValue:[NSNumber numberWithBool:NO] forKey:(__bridge id)kCFStreamSSLValidatesCertificateChain];
  401. NSLog(@"SocketRocket: In debug mode. Allowing connection to any root cert");
  402. #endif
  403. [_outputStream setProperty:SSLOptions
  404. forKey:(__bridge id)kCFStreamPropertySSLSettings];
  405. }
  406. _inputStream.delegate = self;
  407. _outputStream.delegate = self;
  408. }
  409. - (void)openConnection;
  410. {
  411. if (!_scheduledRunloops.count) {
  412. [self scheduleInRunLoop:[NSRunLoop SR_networkRunLoop] forMode:NSDefaultRunLoopMode];
  413. }
  414. [_outputStream open];
  415. [_inputStream open];
  416. }
  417. - (void)scheduleInRunLoop:(NSRunLoop *)aRunLoop forMode:(NSString *)mode;
  418. {
  419. [_outputStream scheduleInRunLoop:aRunLoop forMode:mode];
  420. [_inputStream scheduleInRunLoop:aRunLoop forMode:mode];
  421. [_scheduledRunloops addObject:@[aRunLoop, mode]];
  422. }
  423. - (void)unscheduleFromRunLoop:(NSRunLoop *)aRunLoop forMode:(NSString *)mode;
  424. {
  425. [_outputStream removeFromRunLoop:aRunLoop forMode:mode];
  426. [_inputStream removeFromRunLoop:aRunLoop forMode:mode];
  427. [_scheduledRunloops removeObject:@[aRunLoop, mode]];
  428. }
  429. - (void)close;
  430. {
  431. [self closeWithCode:SRStatusCodeNormal reason:nil];
  432. }
  433. - (void)closeWithCode:(NSInteger)code reason:(NSString *)reason;
  434. {
  435. assert(code);
  436. dispatch_async(_workQueue, ^{
  437. if (self.readyState == SR_CLOSING || self.readyState == SR_CLOSED) {
  438. return;
  439. }
  440. BOOL wasConnecting = self.readyState == SR_CONNECTING;
  441. self.readyState = SR_CLOSING;
  442. SRFastLog(@"Closing with code %d reason %@", code, reason);
  443. if (wasConnecting) {
  444. [self closeConnection];
  445. return;
  446. }
  447. size_t maxMsgSize = [reason maximumLengthOfBytesUsingEncoding:NSUTF8StringEncoding];
  448. NSMutableData *mutablePayload = [[NSMutableData alloc] initWithLength:sizeof(uint16_t) + maxMsgSize];
  449. NSData *payload = mutablePayload;
  450. ((uint16_t *)mutablePayload.mutableBytes)[0] = EndianU16_BtoN(code);
  451. if (reason) {
  452. NSRange remainingRange = {0};
  453. NSUInteger usedLength = 0;
  454. 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];
  455. #pragma unused (success)
  456. assert(success);
  457. assert(remainingRange.length == 0);
  458. if (usedLength != maxMsgSize) {
  459. payload = [payload subdataWithRange:NSMakeRange(0, usedLength + sizeof(uint16_t))];
  460. }
  461. }
  462. [self _sendFrameWithOpcode:SROpCodeConnectionClose data:payload];
  463. });
  464. }
  465. - (void)_closeWithProtocolError:(NSString *)message;
  466. {
  467. // Need to shunt this on the _callbackQueue first to see if they received any messages
  468. [self _performDelegateBlock:^{
  469. [self closeWithCode:SRStatusCodeProtocolError reason:message];
  470. dispatch_async(_workQueue, ^{
  471. [self closeConnection];
  472. });
  473. }];
  474. }
  475. - (void)_failWithError:(NSError *)error;
  476. {
  477. dispatch_async(_workQueue, ^{
  478. if (self.readyState != SR_CLOSED) {
  479. _failed = YES;
  480. [self _performDelegateBlock:^{
  481. if ([self.delegate respondsToSelector:@selector(webSocket:didFailWithError:)]) {
  482. [self.delegate webSocket:self didFailWithError:error];
  483. }
  484. }];
  485. self.readyState = SR_CLOSED;
  486. _selfRetain = nil;
  487. SRFastLog(@"Failing with error %@", error.localizedDescription);
  488. [self closeConnection];
  489. }
  490. });
  491. }
  492. - (void)_writeData:(NSData *)data;
  493. {
  494. [self assertOnWorkQueue];
  495. if (_closeWhenFinishedWriting) {
  496. return;
  497. }
  498. [_outputBuffer appendData:data];
  499. [self _pumpWriting];
  500. }
  501. - (void)send:(id)data;
  502. {
  503. NSAssert(self.readyState != SR_CONNECTING, @"Invalid State: Cannot call send: until connection is open");
  504. // TODO: maybe not copy this for performance
  505. data = [data copy];
  506. dispatch_async(_workQueue, ^{
  507. if ([data isKindOfClass:[NSString class]]) {
  508. [self _sendFrameWithOpcode:SROpCodeTextFrame data:[(NSString *)data dataUsingEncoding:NSUTF8StringEncoding]];
  509. } else if ([data isKindOfClass:[NSData class]]) {
  510. [self _sendFrameWithOpcode:SROpCodeBinaryFrame data:data];
  511. } else if (data == nil) {
  512. [self _sendFrameWithOpcode:SROpCodeTextFrame data:data];
  513. } else {
  514. assert(NO);
  515. }
  516. });
  517. }
  518. - (void)sendPing:(NSData *)data;
  519. {
  520. NSAssert(self.readyState == SR_OPEN, @"Invalid State: Cannot call send: until connection is open");
  521. // TODO: maybe not copy this for performance
  522. data = [data copy] ?: [NSData data]; // It's okay for a ping to be empty
  523. dispatch_async(_workQueue, ^{
  524. [self _sendFrameWithOpcode:SROpCodePing data:data];
  525. });
  526. }
  527. - (void)handlePing:(NSData *)pingData;
  528. {
  529. // Need to pingpong this off _callbackQueue first to make sure messages happen in order
  530. [self _performDelegateBlock:^{
  531. dispatch_async(_workQueue, ^{
  532. [self _sendFrameWithOpcode:SROpCodePong data:pingData];
  533. });
  534. }];
  535. }
  536. - (void)handlePong:(NSData *)pongData;
  537. {
  538. SRFastLog(@"Received pong");
  539. [self _performDelegateBlock:^{
  540. if ([self.delegate respondsToSelector:@selector(webSocket:didReceivePong:)]) {
  541. [self.delegate webSocket:self didReceivePong:pongData];
  542. }
  543. }];
  544. }
  545. - (void)_handleMessage:(id)message
  546. {
  547. SRFastLog(@"Received message");
  548. [self _performDelegateBlock:^{
  549. [self.delegate webSocket:self didReceiveMessage:message];
  550. }];
  551. }
  552. static inline BOOL closeCodeIsValid(int closeCode) {
  553. if (closeCode < 1000) {
  554. return NO;
  555. }
  556. if (closeCode >= 1000 && closeCode <= 1011) {
  557. if (closeCode == 1004 ||
  558. closeCode == 1005 ||
  559. closeCode == 1006) {
  560. return NO;
  561. }
  562. return YES;
  563. }
  564. if (closeCode >= 3000 && closeCode <= 3999) {
  565. return YES;
  566. }
  567. if (closeCode >= 4000 && closeCode <= 4999) {
  568. return YES;
  569. }
  570. return NO;
  571. }
  572. // Note from RFC:
  573. //
  574. // If there is a body, the first two
  575. // bytes of the body MUST be a 2-byte unsigned integer (in network byte
  576. // order) representing a status code with value /code/ defined in
  577. // Section 7.4. Following the 2-byte integer the body MAY contain UTF-8
  578. // encoded data with value /reason/, the interpretation of which is not
  579. // defined by this specification.
  580. - (void)handleCloseWithData:(NSData *)data;
  581. {
  582. size_t dataSize = data.length;
  583. __block uint16_t closeCode = 0;
  584. SRFastLog(@"Received close frame");
  585. if (dataSize == 1) {
  586. // TODO handle error
  587. [self _closeWithProtocolError:@"Payload for close must be larger than 2 bytes"];
  588. return;
  589. } else if (dataSize >= 2) {
  590. [data getBytes:&closeCode length:sizeof(closeCode)];
  591. _closeCode = EndianU16_BtoN(closeCode);
  592. if (!closeCodeIsValid(_closeCode)) {
  593. [self _closeWithProtocolError:[NSString stringWithFormat:@"Cannot have close code of %d", _closeCode]];
  594. return;
  595. }
  596. if (dataSize > 2) {
  597. _closeReason = [[NSString alloc] initWithData:[data subdataWithRange:NSMakeRange(2, dataSize - 2)] encoding:NSUTF8StringEncoding];
  598. if (!_closeReason) {
  599. [self _closeWithProtocolError:@"Close reason MUST be valid UTF-8"];
  600. return;
  601. }
  602. }
  603. } else {
  604. _closeCode = SRStatusNoStatusReceived;
  605. }
  606. [self assertOnWorkQueue];
  607. if (self.readyState == SR_OPEN) {
  608. [self closeWithCode:1000 reason:nil];
  609. }
  610. dispatch_async(_workQueue, ^{
  611. [self closeConnection];
  612. });
  613. }
  614. - (void)closeConnection;
  615. {
  616. [self assertOnWorkQueue];
  617. SRFastLog(@"Trying to disconnect");
  618. _closeWhenFinishedWriting = YES;
  619. [self _pumpWriting];
  620. }
  621. - (void)_handleFrameWithData:(NSData *)frameData opCode:(NSInteger)opcode;
  622. {
  623. // Check that the current data is valid UTF8
  624. BOOL isControlFrame = (opcode == SROpCodePing || opcode == SROpCodePong || opcode == SROpCodeConnectionClose);
  625. if (!isControlFrame) {
  626. [self _readFrameNew];
  627. } else {
  628. dispatch_async(_workQueue, ^{
  629. [self _readFrameContinue];
  630. });
  631. }
  632. switch (opcode) {
  633. case SROpCodeTextFrame: {
  634. NSString *str = [[NSString alloc] initWithData:frameData encoding:NSUTF8StringEncoding];
  635. if (str == nil && frameData) {
  636. [self closeWithCode:SRStatusCodeInvalidUTF8 reason:@"Text frames must be valid UTF-8"];
  637. dispatch_async(_workQueue, ^{
  638. [self closeConnection];
  639. });
  640. return;
  641. }
  642. [self _handleMessage:str];
  643. break;
  644. }
  645. case SROpCodeBinaryFrame:
  646. [self _handleMessage:[frameData copy]];
  647. break;
  648. case SROpCodeConnectionClose:
  649. [self handleCloseWithData:frameData];
  650. break;
  651. case SROpCodePing:
  652. [self handlePing:frameData];
  653. break;
  654. case SROpCodePong:
  655. [self handlePong:frameData];
  656. break;
  657. default:
  658. [self _closeWithProtocolError:[NSString stringWithFormat:@"Unknown opcode %ld", (long)opcode]];
  659. // TODO: Handle invalid opcode
  660. break;
  661. }
  662. }
  663. - (void)_handleFrameHeader:(frame_header)frame_header curData:(NSData *)curData;
  664. {
  665. assert(frame_header.opcode != 0);
  666. if (self.readyState != SR_OPEN) {
  667. return;
  668. }
  669. BOOL isControlFrame = (frame_header.opcode == SROpCodePing || frame_header.opcode == SROpCodePong || frame_header.opcode == SROpCodeConnectionClose);
  670. if (isControlFrame && !frame_header.fin) {
  671. [self _closeWithProtocolError:@"Fragmented control frames not allowed"];
  672. return;
  673. }
  674. if (isControlFrame && frame_header.payload_length >= 126) {
  675. [self _closeWithProtocolError:@"Control frames cannot have payloads larger than 126 bytes"];
  676. return;
  677. }
  678. if (!isControlFrame) {
  679. _currentFrameOpcode = frame_header.opcode;
  680. _currentFrameCount += 1;
  681. }
  682. if (frame_header.payload_length == 0) {
  683. if (isControlFrame) {
  684. [self _handleFrameWithData:curData opCode:frame_header.opcode];
  685. } else {
  686. if (frame_header.fin) {
  687. [self _handleFrameWithData:_currentFrameData opCode:frame_header.opcode];
  688. } else {
  689. // TODO add assert that opcode is not a control;
  690. [self _readFrameContinue];
  691. }
  692. }
  693. } else {
  694. assert(frame_header.payload_length <= SIZE_T_MAX);
  695. [self _addConsumerWithDataLength:(size_t)frame_header.payload_length callback:^(SRWebSocket *self, NSData *newData) {
  696. if (isControlFrame) {
  697. [self _handleFrameWithData:newData opCode:frame_header.opcode];
  698. } else {
  699. if (frame_header.fin) {
  700. [self _handleFrameWithData:self->_currentFrameData opCode:frame_header.opcode];
  701. } else {
  702. // TODO add assert that opcode is not a control;
  703. [self _readFrameContinue];
  704. }
  705. }
  706. } readToCurrentFrame:!isControlFrame unmaskBytes:frame_header.masked];
  707. }
  708. }
  709. /* From RFC:
  710. 0 1 2 3
  711. 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
  712. +-+-+-+-+-------+-+-------------+-------------------------------+
  713. |F|R|R|R| opcode|M| Payload len | Extended payload length |
  714. |I|S|S|S| (4) |A| (7) | (16/64) |
  715. |N|V|V|V| |S| | (if payload len==126/127) |
  716. | |1|2|3| |K| | |
  717. +-+-+-+-+-------+-+-------------+ - - - - - - - - - - - - - - - +
  718. | Extended payload length continued, if payload len == 127 |
  719. + - - - - - - - - - - - - - - - +-------------------------------+
  720. | |Masking-key, if MASK set to 1 |
  721. +-------------------------------+-------------------------------+
  722. | Masking-key (continued) | Payload Data |
  723. +-------------------------------- - - - - - - - - - - - - - - - +
  724. : Payload Data continued ... :
  725. + - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - +
  726. | Payload Data continued ... |
  727. +---------------------------------------------------------------+
  728. */
  729. static const uint8_t SRFinMask = 0x80;
  730. static const uint8_t SROpCodeMask = 0x0F;
  731. static const uint8_t SRRsvMask = 0x70;
  732. static const uint8_t SRMaskMask = 0x80;
  733. static const uint8_t SRPayloadLenMask = 0x7F;
  734. - (void)_readFrameContinue;
  735. {
  736. assert((_currentFrameCount == 0 && _currentFrameOpcode == 0) || (_currentFrameCount > 0 && _currentFrameOpcode > 0));
  737. [self _addConsumerWithDataLength:2 callback:^(SRWebSocket *self, NSData *data) {
  738. __block frame_header header = {0};
  739. const uint8_t *headerBuffer = data.bytes;
  740. assert(data.length >= 2);
  741. if (headerBuffer[0] & SRRsvMask) {
  742. [self _closeWithProtocolError:@"Server used RSV bits"];
  743. return;
  744. }
  745. uint8_t receivedOpcode = (SROpCodeMask & headerBuffer[0]);
  746. BOOL isControlFrame = (receivedOpcode == SROpCodePing || receivedOpcode == SROpCodePong || receivedOpcode == SROpCodeConnectionClose);
  747. if (!isControlFrame && receivedOpcode != 0 && self->_currentFrameCount > 0) {
  748. [self _closeWithProtocolError:@"all data frames after the initial data frame must have opcode 0"];
  749. return;
  750. }
  751. if (receivedOpcode == 0 && self->_currentFrameCount == 0) {
  752. [self _closeWithProtocolError:@"cannot continue a message"];
  753. return;
  754. }
  755. header.opcode = receivedOpcode == 0 ? self->_currentFrameOpcode : receivedOpcode;
  756. header.fin = !!(SRFinMask & headerBuffer[0]);
  757. header.masked = !!(SRMaskMask & headerBuffer[1]);
  758. header.payload_length = SRPayloadLenMask & headerBuffer[1];
  759. headerBuffer = NULL;
  760. if (header.masked) {
  761. [self _closeWithProtocolError:@"Client must receive unmasked data"];
  762. }
  763. size_t extra_bytes_needed = header.masked ? sizeof(_currentReadMaskKey) : 0;
  764. if (header.payload_length == 126) {
  765. extra_bytes_needed += sizeof(uint16_t);
  766. } else if (header.payload_length == 127) {
  767. extra_bytes_needed += sizeof(uint64_t);
  768. }
  769. if (extra_bytes_needed == 0) {
  770. [self _handleFrameHeader:header curData:self->_currentFrameData];
  771. } else {
  772. [self _addConsumerWithDataLength:extra_bytes_needed callback:^(SRWebSocket *self, NSData *data) {
  773. size_t mapped_size = data.length;
  774. #pragma unused (mapped_size)
  775. const void *mapped_buffer = data.bytes;
  776. size_t offset = 0;
  777. if (header.payload_length == 126) {
  778. assert(mapped_size >= sizeof(uint16_t));
  779. uint16_t newLen = EndianU16_BtoN(*(uint16_t *)(mapped_buffer));
  780. header.payload_length = newLen;
  781. offset += sizeof(uint16_t);
  782. } else if (header.payload_length == 127) {
  783. assert(mapped_size >= sizeof(uint64_t));
  784. header.payload_length = EndianU64_BtoN(*(uint64_t *)(mapped_buffer));
  785. offset += sizeof(uint64_t);
  786. } else {
  787. assert(header.payload_length < 126 && header.payload_length >= 0);
  788. }
  789. if (header.masked) {
  790. assert(mapped_size >= sizeof(_currentReadMaskOffset) + offset);
  791. memcpy(self->_currentReadMaskKey, ((uint8_t *)mapped_buffer) + offset, sizeof(self->_currentReadMaskKey));
  792. }
  793. [self _handleFrameHeader:header curData:self->_currentFrameData];
  794. } readToCurrentFrame:NO unmaskBytes:NO];
  795. }
  796. } readToCurrentFrame:NO unmaskBytes:NO];
  797. }
  798. - (void)_readFrameNew;
  799. {
  800. dispatch_async(_workQueue, ^{
  801. [_currentFrameData setLength:0];
  802. _currentFrameOpcode = 0;
  803. _currentFrameCount = 0;
  804. _readOpCount = 0;
  805. _currentStringScanPosition = 0;
  806. [self _readFrameContinue];
  807. });
  808. }
  809. - (void)_pumpWriting;
  810. {
  811. [self assertOnWorkQueue];
  812. NSUInteger dataLength = _outputBuffer.length;
  813. if (dataLength - _outputBufferOffset > 0 && _outputStream.hasSpaceAvailable) {
  814. NSInteger bytesWritten = [_outputStream write:_outputBuffer.bytes + _outputBufferOffset maxLength:dataLength - _outputBufferOffset];
  815. if (bytesWritten == -1) {
  816. [self _failWithError:[NSError errorWithDomain:SRWebSocketErrorDomain code:2145 userInfo:[NSDictionary dictionaryWithObject:@"Error writing to stream" forKey:NSLocalizedDescriptionKey]]];
  817. return;
  818. }
  819. _outputBufferOffset += bytesWritten;
  820. if (_outputBufferOffset > 4096 && _outputBufferOffset > (_outputBuffer.length >> 1)) {
  821. _outputBuffer = [[NSMutableData alloc] initWithBytes:(char *)_outputBuffer.bytes + _outputBufferOffset length:_outputBuffer.length - _outputBufferOffset];
  822. _outputBufferOffset = 0;
  823. }
  824. }
  825. if (_closeWhenFinishedWriting &&
  826. _outputBuffer.length - _outputBufferOffset == 0 &&
  827. (_inputStream.streamStatus != NSStreamStatusNotOpen &&
  828. _inputStream.streamStatus != NSStreamStatusClosed) &&
  829. !_sentClose) {
  830. _sentClose = YES;
  831. [_outputStream close];
  832. [_inputStream close];
  833. for (NSArray *runLoop in [_scheduledRunloops copy]) {
  834. [self unscheduleFromRunLoop:[runLoop objectAtIndex:0] forMode:[runLoop objectAtIndex:1]];
  835. }
  836. if (!_failed) {
  837. [self _performDelegateBlock:^{
  838. if ([self.delegate respondsToSelector:@selector(webSocket:didCloseWithCode:reason:wasClean:)]) {
  839. [self.delegate webSocket:self didCloseWithCode:_closeCode reason:_closeReason wasClean:YES];
  840. }
  841. }];
  842. }
  843. _selfRetain = nil;
  844. }
  845. }
  846. - (void)_addConsumerWithScanner:(stream_scanner)consumer callback:(data_callback)callback;
  847. {
  848. [self assertOnWorkQueue];
  849. [self _addConsumerWithScanner:consumer callback:callback dataLength:0];
  850. }
  851. - (void)_addConsumerWithDataLength:(size_t)dataLength callback:(data_callback)callback readToCurrentFrame:(BOOL)readToCurrentFrame unmaskBytes:(BOOL)unmaskBytes;
  852. {
  853. [self assertOnWorkQueue];
  854. assert(dataLength);
  855. [_consumers addObject:[_consumerPool consumerWithScanner:nil handler:callback bytesNeeded:dataLength readToCurrentFrame:readToCurrentFrame unmaskBytes:unmaskBytes]];
  856. [self _pumpScanner];
  857. }
  858. - (void)_addConsumerWithScanner:(stream_scanner)consumer callback:(data_callback)callback dataLength:(size_t)dataLength;
  859. {
  860. [self assertOnWorkQueue];
  861. [_consumers addObject:[_consumerPool consumerWithScanner:consumer handler:callback bytesNeeded:dataLength readToCurrentFrame:NO unmaskBytes:NO]];
  862. [self _pumpScanner];
  863. }
  864. static const char CRLFCRLFBytes[] = {'\r', '\n', '\r', '\n'};
  865. - (void)_readUntilHeaderCompleteWithCallback:(data_callback)dataHandler;
  866. {
  867. [self _readUntilBytes:CRLFCRLFBytes length:sizeof(CRLFCRLFBytes) callback:dataHandler];
  868. }
  869. - (void)_readUntilBytes:(const void *)bytes length:(size_t)length callback:(data_callback)dataHandler;
  870. {
  871. // TODO optimize so this can continue from where we last searched
  872. stream_scanner consumer = ^size_t(NSData *data) {
  873. __block size_t found_size = 0;
  874. __block size_t match_count = 0;
  875. size_t size = data.length;
  876. const unsigned char *buffer = data.bytes;
  877. for (size_t i = 0; i < size; i++ ) {
  878. if (((const unsigned char *)buffer)[i] == ((const unsigned char *)bytes)[match_count]) {
  879. match_count += 1;
  880. if (match_count == length) {
  881. found_size = i + 1;
  882. break;
  883. }
  884. } else {
  885. match_count = 0;
  886. }
  887. }
  888. return found_size;
  889. };
  890. [self _addConsumerWithScanner:consumer callback:dataHandler];
  891. }
  892. // Returns true if did work
  893. - (BOOL)_innerPumpScanner {
  894. BOOL didWork = NO;
  895. if (self.readyState >= SR_CLOSING) {
  896. return didWork;
  897. }
  898. if (!_consumers.count) {
  899. return didWork;
  900. }
  901. size_t curSize = _readBuffer.length - _readBufferOffset;
  902. if (!curSize) {
  903. return didWork;
  904. }
  905. SRIOConsumer *consumer = [_consumers objectAtIndex:0];
  906. size_t bytesNeeded = consumer.bytesNeeded;
  907. size_t foundSize = 0;
  908. if (consumer.consumer) {
  909. NSData *tempView = [NSData dataWithBytesNoCopy:(char *)_readBuffer.bytes + _readBufferOffset length:_readBuffer.length - _readBufferOffset freeWhenDone:NO];
  910. foundSize = consumer.consumer(tempView);
  911. } else {
  912. assert(consumer.bytesNeeded);
  913. if (curSize >= bytesNeeded) {
  914. foundSize = bytesNeeded;
  915. } else if (consumer.readToCurrentFrame) {
  916. foundSize = curSize;
  917. }
  918. }
  919. NSData *slice = nil;
  920. if (consumer.readToCurrentFrame || foundSize) {
  921. NSRange sliceRange = NSMakeRange(_readBufferOffset, foundSize);
  922. slice = [_readBuffer subdataWithRange:sliceRange];
  923. _readBufferOffset += foundSize;
  924. if (_readBufferOffset > 4096 && _readBufferOffset > (_readBuffer.length >> 1)) {
  925. _readBuffer = [[NSMutableData alloc] initWithBytes:(char *)_readBuffer.bytes + _readBufferOffset length:_readBuffer.length - _readBufferOffset]; _readBufferOffset = 0;
  926. }
  927. if (consumer.unmaskBytes) {
  928. NSMutableData *mutableSlice = [slice mutableCopy];
  929. NSUInteger len = mutableSlice.length;
  930. uint8_t *bytes = mutableSlice.mutableBytes;
  931. for (NSUInteger i = 0; i < len; i++) {
  932. bytes[i] = bytes[i] ^ _currentReadMaskKey[_currentReadMaskOffset % sizeof(_currentReadMaskKey)];
  933. _currentReadMaskOffset += 1;
  934. }
  935. slice = mutableSlice;
  936. }
  937. if (consumer.readToCurrentFrame) {
  938. [_currentFrameData appendData:slice];
  939. _readOpCount += 1;
  940. if (_currentFrameOpcode == SROpCodeTextFrame) {
  941. // Validate UTF8 stuff.
  942. size_t currentDataSize = _currentFrameData.length;
  943. if (_currentFrameOpcode == SROpCodeTextFrame && currentDataSize > 0) {
  944. // TODO: Optimize the crap out of this. Don't really have to copy all the data each time
  945. size_t scanSize = currentDataSize - _currentStringScanPosition;
  946. NSData *scan_data = [_currentFrameData subdataWithRange:NSMakeRange(_currentStringScanPosition, scanSize)];
  947. int32_t valid_utf8_size = validate_dispatch_data_partial_string(scan_data);
  948. if (valid_utf8_size == -1) {
  949. [self closeWithCode:SRStatusCodeInvalidUTF8 reason:@"Text frames must be valid UTF-8"];
  950. dispatch_async(_workQueue, ^{
  951. [self closeConnection];
  952. });
  953. return didWork;
  954. } else {
  955. _currentStringScanPosition += valid_utf8_size;
  956. }
  957. }
  958. }
  959. consumer.bytesNeeded -= foundSize;
  960. if (consumer.bytesNeeded == 0) {
  961. [_consumers removeObjectAtIndex:0];
  962. consumer.handler(self, nil);
  963. [_consumerPool returnConsumer:consumer];
  964. didWork = YES;
  965. }
  966. } else if (foundSize) {
  967. [_consumers removeObjectAtIndex:0];
  968. consumer.handler(self, slice);
  969. [_consumerPool returnConsumer:consumer];
  970. didWork = YES;
  971. }
  972. }
  973. return didWork;
  974. }
  975. -(void)_pumpScanner;
  976. {
  977. [self assertOnWorkQueue];
  978. if (!_isPumping) {
  979. _isPumping = YES;
  980. } else {
  981. return;
  982. }
  983. while ([self _innerPumpScanner]) {
  984. }
  985. _isPumping = NO;
  986. }
  987. //#define NOMASK
  988. static const size_t SRFrameHeaderOverhead = 32;
  989. - (void)_sendFrameWithOpcode:(SROpCode)opcode data:(id)data;
  990. {
  991. [self assertOnWorkQueue];
  992. if (nil == data) {
  993. return;
  994. }
  995. NSAssert([data isKindOfClass:[NSData class]] || [data isKindOfClass:[NSString class]], @"NSString or NSData");
  996. size_t payloadLength = [data isKindOfClass:[NSString class]] ? [(NSString *)data lengthOfBytesUsingEncoding:NSUTF8StringEncoding] : [data length];
  997. NSMutableData *frame = [[NSMutableData alloc] initWithLength:payloadLength + SRFrameHeaderOverhead];
  998. if (!frame) {
  999. [self closeWithCode:SRStatusCodeMessageTooBig reason:@"Message too big"];
  1000. return;
  1001. }
  1002. uint8_t *frame_buffer = (uint8_t *)[frame mutableBytes];
  1003. // set fin
  1004. frame_buffer[0] = SRFinMask | opcode;
  1005. BOOL useMask = YES;
  1006. #ifdef NOMASK
  1007. useMask = NO;
  1008. #endif
  1009. if (useMask) {
  1010. // set the mask and header
  1011. frame_buffer[1] |= SRMaskMask;
  1012. }
  1013. size_t frame_buffer_size = 2;
  1014. const uint8_t *unmasked_payload = NULL;
  1015. if ([data isKindOfClass:[NSData class]]) {
  1016. unmasked_payload = (uint8_t *)[data bytes];
  1017. } else if ([data isKindOfClass:[NSString class]]) {
  1018. unmasked_payload = (const uint8_t *)[data UTF8String];
  1019. } else {
  1020. return;
  1021. }
  1022. if (payloadLength < 126) {
  1023. frame_buffer[1] |= payloadLength;
  1024. } else if (payloadLength <= UINT16_MAX) {
  1025. frame_buffer[1] |= 126;
  1026. *((uint16_t *)(frame_buffer + frame_buffer_size)) = EndianU16_BtoN((uint16_t)payloadLength);
  1027. frame_buffer_size += sizeof(uint16_t);
  1028. } else {
  1029. frame_buffer[1] |= 127;
  1030. *((uint64_t *)(frame_buffer + frame_buffer_size)) = EndianU64_BtoN((uint64_t)payloadLength);
  1031. frame_buffer_size += sizeof(uint64_t);
  1032. }
  1033. if (!useMask) {
  1034. for (size_t i = 0; i < payloadLength; i++) {
  1035. frame_buffer[frame_buffer_size] = unmasked_payload[i];
  1036. frame_buffer_size += 1;
  1037. }
  1038. } else {
  1039. uint8_t *mask_key = frame_buffer + frame_buffer_size;
  1040. SecRandomCopyBytes(kSecRandomDefault, sizeof(uint32_t), (uint8_t *)mask_key);
  1041. frame_buffer_size += sizeof(uint32_t);
  1042. // TODO: could probably optimize this with SIMD
  1043. for (size_t i = 0; i < payloadLength; i++) {
  1044. frame_buffer[frame_buffer_size] = unmasked_payload[i] ^ mask_key[i % sizeof(uint32_t)];
  1045. frame_buffer_size += 1;
  1046. }
  1047. }
  1048. assert(frame_buffer_size <= [frame length]);
  1049. frame.length = frame_buffer_size;
  1050. [self _writeData:frame];
  1051. }
  1052. - (void)stream:(NSStream *)aStream handleEvent:(NSStreamEvent)eventCode;
  1053. {
  1054. if (_secure && !_pinnedCertFound && (eventCode == NSStreamEventHasBytesAvailable || eventCode == NSStreamEventHasSpaceAvailable)) {
  1055. NSArray *sslCerts = [_urlRequest SR_SSLPinnedCertificates];
  1056. if (sslCerts) {
  1057. SecTrustRef secTrust = (__bridge SecTrustRef)[aStream propertyForKey:(__bridge id)kCFStreamPropertySSLPeerTrust];
  1058. if (secTrust) {
  1059. NSInteger numCerts = SecTrustGetCertificateCount(secTrust);
  1060. for (NSInteger i = 0; i < numCerts && !_pinnedCertFound; i++) {
  1061. SecCertificateRef cert = SecTrustGetCertificateAtIndex(secTrust, i);
  1062. NSData *certData = CFBridgingRelease(SecCertificateCopyData(cert));
  1063. for (id ref in sslCerts) {
  1064. SecCertificateRef trustedCert = (__bridge SecCertificateRef)ref;
  1065. NSData *trustedCertData = CFBridgingRelease(SecCertificateCopyData(trustedCert));
  1066. if ([trustedCertData isEqualToData:certData]) {
  1067. _pinnedCertFound = YES;
  1068. break;
  1069. }
  1070. }
  1071. }
  1072. }
  1073. if (!_pinnedCertFound) {
  1074. dispatch_async(_workQueue, ^{
  1075. [self _failWithError:[NSError errorWithDomain:SRWebSocketErrorDomain code:23556 userInfo:[NSDictionary dictionaryWithObject:[NSString stringWithFormat:@"Invalid server cert"] forKey:NSLocalizedDescriptionKey]]];
  1076. });
  1077. return;
  1078. }
  1079. }
  1080. }
  1081. dispatch_async(_workQueue, ^{
  1082. switch (eventCode) {
  1083. case NSStreamEventOpenCompleted: {
  1084. SRFastLog(@"NSStreamEventOpenCompleted %@", aStream);
  1085. if (self.readyState >= SR_CLOSING) {
  1086. return;
  1087. }
  1088. assert(_readBuffer);
  1089. if (self.readyState == SR_CONNECTING && aStream == _inputStream) {
  1090. [self didConnect];
  1091. }
  1092. [self _pumpWriting];
  1093. [self _pumpScanner];
  1094. break;
  1095. }
  1096. case NSStreamEventErrorOccurred: {
  1097. SRFastLog(@"NSStreamEventErrorOccurred %@ %@", aStream, [[aStream streamError] copy]);
  1098. /// TODO specify error better!
  1099. [self _failWithError:aStream.streamError];
  1100. _readBufferOffset = 0;
  1101. [_readBuffer setLength:0];
  1102. break;
  1103. }
  1104. case NSStreamEventEndEncountered: {
  1105. [self _pumpScanner];
  1106. SRFastLog(@"NSStreamEventEndEncountered %@", aStream);
  1107. if (aStream.streamError) {
  1108. [self _failWithError:aStream.streamError];
  1109. } else {
  1110. if (self.readyState != SR_CLOSED) {
  1111. self.readyState = SR_CLOSED;
  1112. _selfRetain = nil;
  1113. }
  1114. if (!_sentClose && !_failed) {
  1115. _sentClose = YES;
  1116. // If we get closed in this state it's probably not clean because we should be sending this when we send messages
  1117. [self _performDelegateBlock:^{
  1118. if ([self.delegate respondsToSelector:@selector(webSocket:didCloseWithCode:reason:wasClean:)]) {
  1119. [self.delegate webSocket:self didCloseWithCode:SRStatusCodeGoingAway reason:@"Stream end encountered" wasClean:NO];
  1120. }
  1121. }];
  1122. }
  1123. }
  1124. break;
  1125. }
  1126. case NSStreamEventHasBytesAvailable: {
  1127. SRFastLog(@"NSStreamEventHasBytesAvailable %@", aStream);
  1128. const int bufferSize = 2048;
  1129. uint8_t buffer[bufferSize];
  1130. while (_inputStream.hasBytesAvailable) {
  1131. NSInteger bytes_read = [_inputStream read:buffer maxLength:bufferSize];
  1132. if (bytes_read > 0) {
  1133. [_readBuffer appendBytes:buffer length:bytes_read];
  1134. } else if (bytes_read < 0) {
  1135. [self _failWithError:_inputStream.streamError];
  1136. }
  1137. if (bytes_read != bufferSize) {
  1138. break;
  1139. }
  1140. };
  1141. [self _pumpScanner];
  1142. break;
  1143. }
  1144. case NSStreamEventHasSpaceAvailable: {
  1145. SRFastLog(@"NSStreamEventHasSpaceAvailable %@", aStream);
  1146. [self _pumpWriting];
  1147. break;
  1148. }
  1149. default:
  1150. SRFastLog(@"(default) %@", aStream);
  1151. break;
  1152. }
  1153. });
  1154. }
  1155. @end
  1156. @implementation SRIOConsumer
  1157. @synthesize bytesNeeded = _bytesNeeded;
  1158. @synthesize consumer = _scanner;
  1159. @synthesize handler = _handler;
  1160. @synthesize readToCurrentFrame = _readToCurrentFrame;
  1161. @synthesize unmaskBytes = _unmaskBytes;
  1162. - (void)setupWithScanner:(stream_scanner)scanner handler:(data_callback)handler bytesNeeded:(size_t)bytesNeeded readToCurrentFrame:(BOOL)readToCurrentFrame unmaskBytes:(BOOL)unmaskBytes;
  1163. {
  1164. _scanner = [scanner copy];
  1165. _handler = [handler copy];
  1166. _bytesNeeded = bytesNeeded;
  1167. _readToCurrentFrame = readToCurrentFrame;
  1168. _unmaskBytes = unmaskBytes;
  1169. assert(_scanner || _bytesNeeded);
  1170. }
  1171. @end
  1172. @implementation SRIOConsumerPool {
  1173. NSUInteger _poolSize;
  1174. NSMutableArray *_bufferedConsumers;
  1175. }
  1176. - (id)in