From d821e865ed4fb252c0c3b2d6b8133f09369f8388 Mon Sep 17 00:00:00 2001
From: John Jones <jmjatlanta@gmail.com>
Date: Fri, 24 May 2019 08:53:45 -0500
Subject: [PATCH] Add configurable header

---
 include/fc/network/http/websocket.hpp |  4 ++--
 src/network/http/websocket.cpp        | 21 ++++++++++++++-------
 tests/network/http/websocket_test.cpp |  2 +-
 tests/ws_test_client.cpp              |  6 +++---
 4 files changed, 20 insertions(+), 13 deletions(-)

diff --git a/include/fc/network/http/websocket.hpp b/include/fc/network/http/websocket.hpp
index db67b29..5dd2a8d 100644
--- a/include/fc/network/http/websocket.hpp
+++ b/include/fc/network/http/websocket.hpp
@@ -30,7 +30,7 @@ namespace fc { namespace http {
          boost::any& get_session_data() { return _session_data; }
 
          virtual std::string get_request_header(const std::string& key) = 0;
-         virtual std::string get_remote_hostname() = 0;
+         virtual std::string get_remote_hostname(const std::string& forward_header_key) = 0;
 
          fc::signal<void()> closed;
       private:
@@ -45,7 +45,7 @@ namespace fc { namespace http {
    class websocket_server
    {
       public:
-         websocket_server();
+         websocket_server( std::string forward_header_key = std::string() );
          ~websocket_server();
 
          void on_connection( const on_connection_handler& handler);
diff --git a/src/network/http/websocket.cpp b/src/network/http/websocket.cpp
index a41627e..f315b7d 100644
--- a/src/network/http/websocket.cpp
+++ b/src/network/http/websocket.cpp
@@ -165,11 +165,16 @@ namespace fc { namespace http {
               return _ws_connection->get_request_header(key);
             }
 
-            virtual std::string get_remote_hostname()
+            virtual std::string get_remote_hostname(const std::string& forward_header_key)
             {
                // TODO: check headers, revert to the raw connection details
-               // Notes: T is a pointer to a websocketpp::connection (see connection.hpp)
-               return _ws_connection->get_uri()->get_host();
+               if (!forward_header_key.empty())
+               {
+                  std::string header_value = _ws_connection->get_request_header(forward_header_key);
+                  if (!header_value.empty())
+                     return header_value;
+               }
+               return _ws_connection->get_remote_endpoint();
             }
 
             T _ws_connection;
@@ -180,8 +185,8 @@ namespace fc { namespace http {
       class websocket_server_impl
       {
          public:
-            websocket_server_impl()
-            :_server_thread( fc::thread::current() )
+            websocket_server_impl(const std::string& forward_header_key = std::string() )
+               :_server_thread( fc::thread::current() )
             {
 
                _server.clear_access_channels( websocketpp::log::alevel::all );
@@ -199,7 +204,8 @@ namespace fc { namespace http {
                        assert( current_con != _connections.end() );
                        auto payload = msg->get_payload();
                        std::shared_ptr<websocket_connection> con = current_con->second;
-                       wdump( ("server") ( con->get_remote_hostname() ) (msg->get_payload()));
+                       wlog( "Websocket Server Remote: ${host} Payload: ${body}", 
+                           ("host", con->get_remote_hostname(forward_header_key)) ("body", msg->get_payload()));
                        ++_pending_messages;
                        auto f = fc::async([this,con,payload](){ if( _pending_messages ) --_pending_messages; con->on_message( payload ); });
                        if( _pending_messages > 100 ) 
@@ -584,7 +590,8 @@ namespace fc { namespace http {
 
    } // namespace detail
 
-   websocket_server::websocket_server():my( new detail::websocket_server_impl() ) {}
+   websocket_server::websocket_server(std::string forward_header_key)
+         :my( new detail::websocket_server_impl(forward_header_key) ) {}
    websocket_server::~websocket_server(){}
 
    void websocket_server::on_connection( const on_connection_handler& handler )
diff --git a/tests/network/http/websocket_test.cpp b/tests/network/http/websocket_test.cpp
index 6200b08..cfa78d0 100644
--- a/tests/network/http/websocket_test.cpp
+++ b/tests/network/http/websocket_test.cpp
@@ -19,7 +19,7 @@ BOOST_AUTO_TEST_CASE(websocket_test)
     fc::http::websocket_connection_ptr s_conn, c_conn;
     int port;
     {
-        fc::http::websocket_server server;
+        fc::http::websocket_server server("MyProxyHeaderKey");
         server.on_connection([&]( const fc::http::websocket_connection_ptr& c ){
                 s_conn = c;
                 c->on_message_handler([&](const std::string& s){
diff --git a/tests/ws_test_client.cpp b/tests/ws_test_client.cpp
index 55006d5..e841a34 100644
--- a/tests/ws_test_client.cpp
+++ b/tests/ws_test_client.cpp
@@ -17,9 +17,9 @@ int main(int argc, char** argv)
       
       fc::http::websocket_client client;
       fc::http::websocket_connection_ptr s_conn, c_conn;
-      int port = std::stoi(argv[1]);
-      wlog( "Connecting to server at port ${port}", ("port", argv[1]) );
-      c_conn = client.connect( "ws://127.0.0.1:" + fc::to_string(port) );
+      std::string url = argv[1];
+      wlog( "Connecting to server at url ${url}", ("url", url) );
+      c_conn = client.connect( "ws://" + url );
 
       std::string echo;
       c_conn->on_message_handler([&](const std::string& s){