Allow server delegate to return custom HTTP error response

* Added -server:acceptWebSocketWithRequest:response: delegate method that allows the delegate to return an HTTP status code and headers. * PSWebSocketServer now sends an HTTP response on failure instead of just closing the socket. * PSWebSocket now captures an HTTP error response from the peer and stores it in custom userInfo properties of the returned NSError so the client can see the HTTP status and headers when the connection fails. * Replaced custom CRLF-scanning code with calls to memmem(). * Removed a no-op line that triggered a static analyzer warning. Fixes zwopple#28
couchbasedeps · Apr 12, 2015 · 3d4d16a · 3d4d16a
1 parent d42f035
commit 3d4d16a
Show file tree

Hide file tree

Showing 4 changed files with 60 additions and 41 deletions.
diff --git a/PocketSocket/PSWebSocketDriver.m b/PocketSocket/PSWebSocketDriver.m
@@ -182,7 +182,6 @@ - (void)sendCloseCode:(NSInteger)code reason:(NSString *)reason {
         NSAssert(success, @"Failed to write reason when sending close frame");
         NSAssert(remainingRange.length == 0, @"Failed to write reason when sending close frame");
 
-        success = YES;
         data.length = usedLength + sizeof(uint16_t);
     }
     [self writeMessageWithOpCode:PSWebSocketOpCodeClose data:data];
@@ -401,29 +400,16 @@ - (NSInteger)readBytes:(void *)bytes maxLength:(NSUInteger)maxLength error:(NSEr
             NSAssert(maxLength > 0, @"Must have 1 or more bytes");
             NSAssert(_state == PSWebSocketDriverStateHandshakeResponse, @"Invalid state for reading handshake response");
 
-            uint8_t boundary[] = {'\r', '\n','\r', '\n'};
-            NSUInteger preBoundaryLength = 0;
-            NSUInteger matched = 0;
-            for(NSUInteger i = 0; i < maxLength; ++i) {
-                const uint8_t byte = ((const uint8_t *)bytes)[i];
-                const uint8_t boundaryByte = boundary[matched];
-                if(byte == boundaryByte) {
-                    if(++matched == sizeof(boundary)) {
-                        preBoundaryLength = i + 1;
-                        break;
-                    }
-                } else {
-                    matched = 0;
-                }
-            }
-            if(preBoundaryLength == 0) {
+            void* boundary = memmem(bytes, maxLength, "\r\n\r\n", 4);
+            if (boundary == NULL) {
                 // do not allow too much data for headers
                 if(maxLength >= 16384) {
                     PSWebSocketSetOutError(outError, PSWebSocketErrorCodeHandshakeFailed, @"HTTP headers did not finish after reading 16384 bytes");
                     return -1;
                 }
                 return 0;
             }
+            NSUInteger preBoundaryLength = boundary + 4 - bytes;
 
             // create handshake
             CFHTTPMessageRef msg = CFHTTPMessageCreateEmpty(NULL, NO);
@@ -444,7 +430,21 @@ - (NSInteger)readBytes:(void *)bytes maxLength:(NSUInteger)maxLength error:(NSEr
 
             // validate status
             if(statusCode != 101) {
-                PSWebSocketSetOutError(outError, PSWebSocketErrorCodeHandshakeFailed, @"Handshake failed");
+                if(outError) {
+                    NSString* message = CFBridgingRelease(CFHTTPMessageCopyResponseStatusLine(msg));
+                    if (!message)
+                        message = [NSHTTPURLResponse localizedStringForStatusCode:statusCode];
+                    else if ([message hasPrefix:@"HTTP/1.1 "])
+                        message = [message substringFromIndex:9];
+                    NSString* desc = [@"Handshake failed: " stringByAppendingString:message];
+                    NSDictionary* userInfo = @{NSLocalizedDescriptionKey: desc,
+                                               NSLocalizedFailureReasonErrorKey: message,
+                                               PSHTTPStatusErrorKey: @(statusCode),
+                                               PSHTTPHeadersErrorKey: headers};
+                    *outError = [NSError errorWithDomain:PSWebSocketErrorDomain
+                                                    code:PSWebSocketErrorCodeHandshakeFailed
+                                                userInfo:userInfo];
+                }
                 return - 1;
             }
 

diff --git a/PocketSocket/PSWebSocketServer.h b/PocketSocket/PSWebSocketServer.h
@@ -25,13 +25,16 @@
 - (void)server:(PSWebSocketServer *)server didFailWithError:(NSError *)error;
 - (void)serverDidStop:(PSWebSocketServer *)server;
 
-- (BOOL)server:(PSWebSocketServer *)server acceptWebSocketWithRequest:(NSURLRequest *)request;
 - (void)server:(PSWebSocketServer *)server webSocketDidOpen:(PSWebSocket *)webSocket;
 - (void)server:(PSWebSocketServer *)server webSocket:(PSWebSocket *)webSocket didReceiveMessage:(id)message;
 - (void)server:(PSWebSocketServer *)server webSocket:(PSWebSocket *)webSocket didFailWithError:(NSError *)error;
 - (void)server:(PSWebSocketServer *)server webSocket:(PSWebSocket *)webSocket didCloseWithCode:(NSInteger)code reason:(NSString *)reason wasClean:(BOOL)wasClean;
 
 @optional
+// Delegate may implement either one of these; variant with response is preferred:
+- (BOOL)server:(PSWebSocketServer *)server acceptWebSocketWithRequest:(NSURLRequest *)request;
+- (BOOL)server:(PSWebSocketServer *)server acceptWebSocketWithRequest:(NSURLRequest *)request response:(NSHTTPURLResponse **)response;
+
 - (void)server:(PSWebSocketServer *)server webSocketIsHungry:(PSWebSocket *)webSocket;
 @end
 

diff --git a/PocketSocket/PSWebSocketServer.m b/PocketSocket/PSWebSocketServer.m
@@ -200,7 +200,7 @@ - (void)disconnectGracefully:(BOOL)silent {
     }
 
     for(PSWebSocketServerConnection *connection in _connections.allObjects) {
-        [self disconnectConnectionGracefully:connection statusCode:500 description:@"Service Going Away"];
+        [self disconnectConnectionGracefully:connection statusCode:500 description:@"Service Going Away" headers: nil];
     }
     for(PSWebSocket *webSocket in _webSockets.allObjects) {
         [webSocket close];
@@ -346,12 +346,22 @@ - (void)detatchConnection:(PSWebSocketServerConnection *)connection {
     connection.inputStream.delegate = nil;
     connection.outputStream.delegate = nil;
 }
-- (void)disconnectConnectionGracefully:(PSWebSocketServerConnection *)connection statusCode:(NSInteger)statusCode description:(NSString *)description {
+- (void)disconnectConnectionGracefully:(PSWebSocketServerConnection *)connection
+                            statusCode:(NSInteger)statusCode
+                           description:(NSString *)description
+                               headers:(NSDictionary*)headers
+{
     if(connection.readyState >= PSWebSocketServerConnectionReadyStateClosing) {
         return;
     }
     connection.readyState = PSWebSocketServerConnectionReadyStateClosing;
+    if (!description)
+        description = [NSHTTPURLResponse localizedStringForStatusCode:statusCode];
     CFHTTPMessageRef msg = CFHTTPMessageCreateResponse(kCFAllocatorDefault, statusCode, (__bridge CFStringRef)description, kCFHTTPVersion1_1);
+    for (NSString* name in headers) {
+        CFHTTPMessageSetHeaderFieldValue(msg, (__bridge CFStringRef)name,
+                                         (__bridge CFStringRef)headers[name]);
+    }
     NSData *data = CFBridgingRelease(CFHTTPMessageCopySerializedMessage(msg));
     CFRelease(msg);
     [connection.outputBuffer appendData:data];
@@ -397,27 +407,17 @@ - (void)pumpInput {
         }
 
         if(connection.inputBuffer.bytesAvailable > 4) {
-            uint8_t boundary[] = {'\r', '\n','\r', '\n'};
-            NSUInteger boundaryOffset = 0;
-            NSUInteger matched = 0;
-            for(NSUInteger i = 0; i < connection.inputBuffer.bytesAvailable; ++i) {
-                const uint8_t byte = ((const uint8_t *)connection.inputBuffer.bytes)[i];
-                const uint8_t boundaryByte = boundary[matched];
-                if(byte == boundaryByte) {
-                    if(++matched == sizeof(boundary)) {
-                        boundaryOffset = i + 1;
-                        break;
-                    }
-                } else {
-                    matched = 0;
-                }
-            }
-            if(boundaryOffset == 0) {
+            void* boundary = memmem(connection.inputBuffer.bytes,
+                                    connection.inputBuffer.bytesAvailable,
+                                    "\r\n\r\n", 4);
+            if (boundary == NULL) {
+                // Haven't reached end of HTTP headers yet
                 if(connection.inputBuffer.bytesAvailable >= 16384) {
                     [self disconnectConnection:connection];
                 }
                 continue;
             }
+            NSUInteger boundaryOffset = boundary + 4 - connection.inputBuffer.bytes;
 
             CFHTTPMessageRef msg = CFHTTPMessageCreateEmpty(kCFAllocatorDefault, YES);
             CFHTTPMessageAppendBytes(msg, connection.inputBuffer.bytes, connection.inputBuffer.bytesAvailable);
@@ -444,18 +444,30 @@ - (void)pumpInput {
             }];
 
             if(![PSWebSocket isWebSocketRequest:request]) {
-                [self disconnectConnection:connection];
+                [self disconnectConnectionGracefully:connection
+                                          statusCode:501 description:@"WebSockets only, please"
+                                             headers:nil];
                 CFRelease(msg);
                 continue;
             }
 
             if(_delegate) {
-                __block BOOL accept = NO;
+                __block BOOL accept;
+                __block NSHTTPURLResponse* response = nil;
                 [self executeDelegateAndWait:^{
-                    accept = [_delegate server:self acceptWebSocketWithRequest:request];
+                    if ([_delegate respondsToSelector: @selector(server:acceptWebSocketWithRequest:response:)]) {
+                        accept = [_delegate server:self acceptWebSocketWithRequest:request response:&response];
+                    } else if ([_delegate respondsToSelector: @selector(server:acceptWebSocketWithRequest:)]) {
+                        accept = [_delegate server:self acceptWebSocketWithRequest:request];
+                    } else {
+                        accept = YES;
+                    }
                 }];
                 if(!accept) {
-                    [self disconnectConnection:connection];
+                    [self disconnectConnectionGracefully:connection
+                                              statusCode:(response.statusCode ?: 403)
+                                             description:nil
+                                                 headers:response.allHeaderFields];
                     CFRelease(msg);
                     continue;
                 }

diff --git a/PocketSocket/PSWebSocketTypes.h b/PocketSocket/PSWebSocketTypes.h
@@ -41,3 +41,7 @@ typedef NS_ENUM(NSInteger, PSWebSocketStatusCode) {
 
 #define PSWebSocketGUID @"258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
 #define PSWebSocketErrorDomain @"PSWebSocketErrorDomain"
+
+// NSError userInfo keys, used with PSWebSocketErrorCodeHandshakeFailed:
+#define PSHTTPStatusErrorKey @"HTTPStatus"
+#define PSHTTPHeadersErrorKey @"HTTPHeaders"