added user close message (#38)

Kitura · Apr 17, 2018 · 241f75b · 241f75b
1 parent 592bd2a
commit 241f75b
Show file tree

Hide file tree

Showing 3 changed files with 56 additions and 4 deletions.
diff --git a/Sources/KituraWebSocket/WebSocketConnection.swift b/Sources/KituraWebSocket/WebSocketConnection.swift
@@ -221,6 +221,7 @@ public class WebSocketConnection {
         case .close:
             if active {
                 let reasonCode: WebSocketCloseReasonCode
+                var description: String?
                 if frame.payload.length >= 2 && frame.payload.length < 126 {
                     let networkOrderedReasonCode = UnsafeRawPointer(frame.payload.bytes).assumingMemoryBound(to: UInt16.self)[0]
                     let hostOrderedReasonCode: UInt16
@@ -230,13 +231,21 @@ public class WebSocketConnection {
                         hostOrderedReasonCode = UInt16(CFSwapInt16BigToHost(networkOrderedReasonCode))
                     #endif
                     reasonCode = WebSocketCloseReasonCode.from(code: hostOrderedReasonCode)
+
+                    var closeMessage = Data(referencing: frame.payload)
+                    _ = closeMessage.removeFirst(2)
+                    description = String(data: closeMessage, encoding: .utf8)
+                    if description == nil {
+                        closeConnection(reason: .invalidDataContents, description: "Failed to convert received close message to UTF-8 String", hard: true)
+                        return
+                    }
                 } else if frame.payload.length == 0 {
                     reasonCode = .normal
                 } else {
                     connectionClosed(reason: .protocolError, description: "Close frames, which contain a payload, must be between 2 and 125 octets inclusive")
                     return
                 }
-                connectionClosed(reason: reasonCode)
+                connectionClosed(reason: reasonCode, description: description)
             }
             break
 

diff --git a/Tests/KituraWebSocketTests/BasicTests.swift b/Tests/KituraWebSocketTests/BasicTests.swift
@@ -39,7 +39,9 @@ class BasicTests: KituraTest {
             ("testTextLongMessage", testTextLongMessage),
             ("testTextMediumMessage", testTextMediumMessage),
             ("testTextShortMessage", testTextShortMessage),
-            ("testUserDefinedCloseMessage", testUserDefinedCloseMessage)
+//            ("testNullCharacter", testNullCharacter),
+            ("testUserDefinedCloseCode", testUserDefinedCloseCode),
+            ("testUserCloseMessage", testUserCloseMessage)
         ]
     }
 
@@ -316,15 +318,15 @@ class BasicTests: KituraTest {
 
         performServerTest() { expectation in
 
-            let textPayload = self.payload(text: "\u{0}")
+            let textPayload = self.payload(text: "\u{00}")
 
             self.performTest(framesToSend: [(true, self.opcodeText, textPayload)],
                              expectedFrames: [(true, self.opcodeText, textPayload)],
                              expectation: expectation)
         }
     }
 
-    func testUserDefinedCloseMessage() {
+    func testUserDefinedCloseCode() {
         register(closeReason: .userDefined(65535))
 
         performServerTest() { expectation in
@@ -337,4 +339,21 @@ class BasicTests: KituraTest {
                              expectation: expectation)
         }
     }
+
+    func testUserCloseMessage() {
+        register(closeReason: .normal)
+
+        performServerTest() { expectation in
+            let testString = "Testing, 1,2,3"
+            let dataPayload = testString.data(using: String.Encoding.utf8)!
+            let payload = NSMutableData()
+            let closeReasonCode = self.payload(closeReasonCode: .normal)
+            payload.append(closeReasonCode.bytes, length: closeReasonCode.length)
+            payload.append(dataPayload)
+
+            self.performTest(framesToSend: [(true, self.opcodeClose, payload)],
+                             expectedFrames: [(true, self.opcodeClose, payload)],
+                             expectation: expectation)
+        }
+    }
 }
diff --git a/Tests/KituraWebSocketTests/ProtocolErrorTests.swift b/Tests/KituraWebSocketTests/ProtocolErrorTests.swift
@@ -34,6 +34,7 @@ class ProtocolErrorTests: KituraTest {
             ("testJustContinuationFrame", testJustContinuationFrame),
             ("testJustFinalContinuationFrame", testJustFinalContinuationFrame),
             ("testInvalidUTF", testInvalidUTF),
+            ("testInvalidUTFCloseMessage", testInvalidUTFCloseMessage),
             ("testTextAndBinaryFrames", testTextAndBinaryFrames),
             ("testUnmaskedFrame", testUnmaskedFrame)
         ]
@@ -240,6 +241,29 @@ class ProtocolErrorTests: KituraTest {
         }
     }
 
+    func testInvalidUTFCloseMessage() {
+        register(closeReason: .noReasonCodeSent)
+
+        performServerTest() { expectation in
+            let testString = "Testing, 1,2,3"
+            let dataPayload = testString.data(using: String.Encoding.utf16)!
+            let payload = NSMutableData()
+            let closeReasonCode = self.payload(closeReasonCode: .normal)
+            payload.append(closeReasonCode.bytes, length: closeReasonCode.length)
+            payload.append(dataPayload)
+
+            let expectedPayload = NSMutableData()
+            var part = self.payload(closeReasonCode: .invalidDataContents)
+            expectedPayload.append(part.bytes, length: part.length)
+            part = self.payload(text: "Failed to convert received close message to UTF-8 String")
+            expectedPayload.append(part.bytes, length: part.length)
+
+            self.performTest(framesToSend: [(true, self.opcodeClose, payload)],
+                             expectedFrames: [(true, self.opcodeClose, expectedPayload)],
+                             expectation: expectation)
+        }
+    }
+
     func testTextAndBinaryFrames() {
         register(closeReason: .protocolError)