From f17fc2c07998eecc4387137a4dd956e015bad37a Mon Sep 17 00:00:00 2001
From: Aleksey Kudryavtsev <aakudryavtsev@avito.ru>
Date: Fri, 6 Apr 2018 10:52:04 +0300
Subject: [PATCH] Handle close code if it was not handled before closing

---
 Source/WebSocket.swift | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/Source/WebSocket.swift b/Source/WebSocket.swift
index 6c22da2..037ee59 100644
--- a/Source/WebSocket.swift
+++ b/Source/WebSocket.swift
@@ -765,6 +765,7 @@ private class InnerWebSocket: Hashable {
                 privateReadyState = .closed
                 if rd != nil {
                     closeConn()
+                    handleCloseCodeFromFirstCloseFrameIfNeeded()
                     fire {
                         self.eclose()
                         self.event.close(Int(self.closeCode), self.closeReason, self.closeFinal)
@@ -901,6 +902,19 @@ private class InnerWebSocket: Hashable {
             }
         }
     }
+    
+    func handleCloseCodeFromFirstCloseFrameIfNeeded() {
+        guard closeCode == 0 else { return }
+        
+        for frame in frames {
+            if frame.code == .close {
+                closeCode = frame.statusCode
+                closeReason = frame.utf8.text
+                return
+            }
+        }
+    }
+    
     @inline(__always) func fire(_ block: ()->()){
         if let queue = eventQueue {
             queue.sync {