From f18f04c0ee237b668aef6550c2d20e24a2974e8f Mon Sep 17 00:00:00 2001
From: ozychhi <12700059+ozychhi@users.noreply.github.com>
Date: Mon, 26 Aug 2019 11:37:40 -0500
Subject: [PATCH] Add client handshake extra headers (#105)

Even though 6455 defines all the necessary headers needed for
client/server handshake, in practice most of the cases websocket servers
expect few more headers. Therefore adding this functionality.
---
 ixwebsocket/IXWebSocket.cpp          |  7 ++++++-
 ixwebsocket/IXWebSocket.h            |  5 +++++
 ixwebsocket/IXWebSocketHandshake.cpp |  4 ++++
 ixwebsocket/IXWebSocketHandshake.h   | 13 ++++++++-----
 ixwebsocket/IXWebSocketTransport.cpp | 10 ++++++----
 ixwebsocket/IXWebSocketTransport.h   |  7 +++++--
 6 files changed, 34 insertions(+), 12 deletions(-)

diff --git a/ixwebsocket/IXWebSocket.cpp b/ixwebsocket/IXWebSocket.cpp
index dab858af..610452ce 100644
--- a/ixwebsocket/IXWebSocket.cpp
+++ b/ixwebsocket/IXWebSocket.cpp
@@ -70,6 +70,11 @@ namespace ix
         std::lock_guard<std::mutex> lock(_configMutex);
         _url = url;
     }
+    void WebSocket::setExtraHeaders(const std::unordered_map<std::string, std::string>& headers)
+    {
+        std::lock_guard<std::mutex> lock(_configMutex);
+        _extraHeaders = headers;
+    }
 
     const std::string& WebSocket::getUrl() const
     {
@@ -176,7 +181,7 @@ namespace ix
                           _pingTimeoutSecs);
         }
 
-        WebSocketInitResult status = _ws.connectToUrl(_url, timeoutSecs);
+        WebSocketInitResult status = _ws.connectToUrl(_url, _extraHeaders, timeoutSecs);
         if (!status.success)
         {
             return status;
diff --git a/ixwebsocket/IXWebSocket.h b/ixwebsocket/IXWebSocket.h
index 035fa05a..0d9acb77 100644
--- a/ixwebsocket/IXWebSocket.h
+++ b/ixwebsocket/IXWebSocket.h
@@ -21,6 +21,7 @@
 #include <mutex>
 #include <string>
 #include <thread>
+#include <unordered_map>
 
 namespace ix
 {
@@ -44,6 +45,8 @@ namespace ix
         ~WebSocket();
 
         void setUrl(const std::string& url);
+        void setExtraHeaders(const std::unordered_map<std::string, std::string>&
+                                 headers); // send extra headers in client handshake request
         void setPerMessageDeflateOptions(
             const WebSocketPerMessageDeflateOptions& perMessageDeflateOptions);
         void setHeartBeatPeriod(int heartBeatPeriodSecs);
@@ -111,6 +114,8 @@ namespace ix
         WebSocketTransport _ws;
 
         std::string _url;
+        std::unordered_map<std::string, std::string> _extraHeaders;
+
         WebSocketPerMessageDeflateOptions _perMessageDeflateOptions;
         mutable std::mutex _configMutex; // protect all config variables access
 
diff --git a/ixwebsocket/IXWebSocketHandshake.cpp b/ixwebsocket/IXWebSocketHandshake.cpp
index 70e004dc..e9e8c4a2 100644
--- a/ixwebsocket/IXWebSocketHandshake.cpp
+++ b/ixwebsocket/IXWebSocketHandshake.cpp
@@ -88,6 +88,7 @@ namespace ix
     }
 
     WebSocketInitResult WebSocketHandshake::clientHandshake(const std::string& url,
+                                                            const std::unordered_map<std::string, std::string>& extraHeaders,
                                                             const std::string& host,
                                                             const std::string& path,
                                                             int port,
@@ -127,6 +128,9 @@ namespace ix
         ss << "Sec-WebSocket-Version: 13\r\n";
         ss << "Sec-WebSocket-Key: " << secWebSocketKey << "\r\n";
 
+        for (auto& it : extraHeaders) {
+            ss << it.first << ":" << it.second << "\r\n";
+        }
         if (_enablePerMessageDeflate)
         {
             ss << _perMessageDeflateOptions.generateHeader();
diff --git a/ixwebsocket/IXWebSocketHandshake.h b/ixwebsocket/IXWebSocketHandshake.h
index 3d0b33fa..adc503df 100644
--- a/ixwebsocket/IXWebSocketHandshake.h
+++ b/ixwebsocket/IXWebSocketHandshake.h
@@ -16,6 +16,7 @@
 #include <memory>
 #include <string>
 #include <tuple>
+#include <unordered_map>
 
 namespace ix
 {
@@ -50,11 +51,13 @@ namespace ix
                            WebSocketPerMessageDeflateOptions& perMessageDeflateOptions,
                            std::atomic<bool>& enablePerMessageDeflate);
 
-        WebSocketInitResult clientHandshake(const std::string& url,
-                                            const std::string& host,
-                                            const std::string& path,
-                                            int port,
-                                            int timeoutSecs);
+        WebSocketInitResult clientHandshake(
+            const std::string& url,
+            const std::unordered_map<std::string, std::string>& extraHeaders,
+            const std::string& host,
+            const std::string& path,
+            int port,
+            int timeoutSecs);
 
         WebSocketInitResult serverHandshake(int fd, int timeoutSecs);
 
diff --git a/ixwebsocket/IXWebSocketTransport.cpp b/ixwebsocket/IXWebSocketTransport.cpp
index a799408e..ada516fd 100644
--- a/ixwebsocket/IXWebSocketTransport.cpp
+++ b/ixwebsocket/IXWebSocketTransport.cpp
@@ -127,8 +127,10 @@ namespace ix
     }
 
     // Client
-    WebSocketInitResult WebSocketTransport::connectToUrl(const std::string& url,
-                                                         int timeoutSecs)
+    WebSocketInitResult WebSocketTransport::connectToUrl(
+        const std::string& url,
+        const std::unordered_map<std::string, std::string>& headers,
+        int timeoutSecs)
     {
         std::lock_guard<std::mutex> lock(_socketMutex);
 
@@ -156,8 +158,8 @@ namespace ix
                                               _perMessageDeflateOptions,
                                               _enablePerMessageDeflate);
 
-        auto result = webSocketHandshake.clientHandshake(url, host, path, port,
-                                                         timeoutSecs);
+        auto result = webSocketHandshake.clientHandshake(url, headers, host, path,
+                                                         port, timeoutSecs);
         if (result.success)
         {
             setReadyState(ReadyState::OPEN);
diff --git a/ixwebsocket/IXWebSocketTransport.h b/ixwebsocket/IXWebSocketTransport.h
index 3da2fa5d..728d6c56 100644
--- a/ixwebsocket/IXWebSocketTransport.h
+++ b/ixwebsocket/IXWebSocketTransport.h
@@ -24,6 +24,7 @@
 #include <memory>
 #include <mutex>
 #include <string>
+#include <unordered_map>
 #include <vector>
 
 namespace ix
@@ -75,8 +76,10 @@ namespace ix
                        int pingIntervalSecs,
                        int pingTimeoutSecs);
 
-        WebSocketInitResult connectToUrl(const std::string& url, // Client
-                                         int timeoutSecs);
+        WebSocketInitResult connectToUrl( // Client
+            const std::string& url,
+            const std::unordered_map<std::string, std::string>& headers,
+            int timeoutSecs);
         WebSocketInitResult connectToSocket(int fd, // Server
                                             int timeoutSecs);