From 241f75ba235ba74e1ec8483930b3998e193248ab Mon Sep 17 00:00:00 2001
From: Andrew Lees <32634907+Andrew-Lees11@users.noreply.github.com>
Date: Tue, 17 Apr 2018 13:12:24 +0100
Subject: [PATCH] added user close message (#38)

---
 .../KituraWebSocket/WebSocketConnection.swift | 11 +++++++-
 Tests/KituraWebSocketTests/BasicTests.swift   | 25 ++++++++++++++++---
 .../ProtocolErrorTests.swift                  | 24 ++++++++++++++++++
 3 files changed, 56 insertions(+), 4 deletions(-)

diff --git a/Sources/KituraWebSocket/WebSocketConnection.swift b/Sources/KituraWebSocket/WebSocketConnection.swift
index 84099cc..b6a0071 100644
--- 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
index 930609c..cad31e1 100644
--- 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,7 +318,7 @@ 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)],
@@ -324,7 +326,7 @@ class BasicTests: KituraTest {
         }
     }
     
-    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
index 16f5886..d4e1805 100644
--- 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)