From ed75d14c868f2ac9c8413430cb6e1e5a76e42d05 Mon Sep 17 00:00:00 2001
From: Benjamin Sergeant <bsergean@gmail.com>
Date: Fri, 15 Nov 2019 17:07:31 -0800
Subject: [PATCH] proxy works but crash when the connection is refused

---
 ixsnake/ixsnake/IXSnakeConnectionState.h |  1 +
 ws/ws_proxy_server.cpp                   | 87 ++++++++++++++++++++++--
 2 files changed, 84 insertions(+), 4 deletions(-)

diff --git a/ixsnake/ixsnake/IXSnakeConnectionState.h b/ixsnake/ixsnake/IXSnakeConnectionState.h
index 713aad32..b4785608 100644
--- a/ixsnake/ixsnake/IXSnakeConnectionState.h
+++ b/ixsnake/ixsnake/IXSnakeConnectionState.h
@@ -20,6 +20,7 @@ namespace snake
         {
             return _nonce;
         }
+
         void setNonce(const std::string& nonce)
         {
             _nonce = nonce;
diff --git a/ws/ws_proxy_server.cpp b/ws/ws_proxy_server.cpp
index eaf073d7..f1cbaddb 100644
--- a/ws/ws_proxy_server.cpp
+++ b/ws/ws_proxy_server.cpp
@@ -10,6 +10,31 @@
 
 namespace ix
 {
+    class ProxyConnectionState : public ix::ConnectionState
+    {
+    public:
+        void setRemoteHost(const std::string& host)
+        {
+            _remoteHost = host;
+        }
+
+        void setRemotePort(int port)
+        {
+            _remotePort = port;
+        }
+
+        ix::WebSocket& webSocket()
+        {
+            return _serverWebSocket;
+        }
+
+    private:
+        std::string _remoteHost;
+        int _remotePort;
+
+        ix::WebSocket _serverWebSocket;
+    };
+
     int ws_proxy_server_main(int port,
                              const std::string& hostname,
                              const ix::SocketTLSOptions& tlsOptions,
@@ -21,15 +46,25 @@ namespace ix
         ix::WebSocketServer server(port, hostname);
         server.setTLSOptions(tlsOptions);
 
+        auto factory = []() -> std::shared_ptr<ix::ConnectionState> {
+            return std::make_shared<ProxyConnectionState>();
+        };
+        server.setConnectionStateFactory(factory);
+
         server.setOnConnectionCallback(
             [remoteHost, remotePort](std::shared_ptr<ix::WebSocket> webSocket,
                                      std::shared_ptr<ConnectionState> connectionState) {
-                webSocket->setOnMessageCallback(
-                    [webSocket, connectionState, remoteHost, remotePort](const WebSocketMessagePtr& msg) {
+                auto state = std::dynamic_pointer_cast<ProxyConnectionState>(connectionState);
+                state->setRemoteHost(remoteHost);
+                state->setRemotePort(remotePort);
+
+                // Server connection
+                state->webSocket().setOnMessageCallback(
+                    [webSocket, state](const WebSocketMessagePtr& msg) {
                         if (msg->type == ix::WebSocketMessageType::Open)
                         {
                             std::cerr << "New connection" << std::endl;
-                            std::cerr << "id: " << connectionState->getId() << std::endl;
+                            std::cerr << "id: " << state->getId() << std::endl;
                             std::cerr << "Uri: " << msg->openInfo.uri << std::endl;
                             std::cerr << "Headers:" << std::endl;
                             for (auto it : msg->openInfo.headers)
@@ -51,13 +86,57 @@ namespace ix
                             ss << "Wait time(ms): " << msg->errorInfo.wait_time << std::endl;
                             ss << "HTTP Status: " << msg->errorInfo.http_status << std::endl;
                             std::cerr << ss.str();
+                            webSocket->close(msg->closeInfo.code, msg->closeInfo.reason);
                         }
                         else if (msg->type == ix::WebSocketMessageType::Message)
                         {
-                            std::cerr << "Received " << msg->wireSize << " bytes" << std::endl;
+                            std::cerr << "Received " << msg->wireSize << " bytes from server" << std::endl;
                             webSocket->send(msg->str, msg->binary);
                         }
                     });
+
+                // Client connection
+                webSocket->setOnMessageCallback(
+                    [state, remoteHost](const WebSocketMessagePtr& msg) {
+                        if (msg->type == ix::WebSocketMessageType::Open)
+                        {
+                            std::cerr << "New connection" << std::endl;
+                            std::cerr << "id: " << state->getId() << std::endl;
+                            std::cerr << "Uri: " << msg->openInfo.uri << std::endl;
+                            std::cerr << "Headers:" << std::endl;
+                            for (auto it : msg->openInfo.headers)
+                            {
+                                std::cerr << it.first << ": " << it.second << std::endl;
+                            }
+
+                            // Connect to the 'real' server
+                            std::string url(remoteHost);
+                            url += msg->openInfo.uri;
+                            state->webSocket().setUrl(url);
+                            state->webSocket().start();
+                        }
+                        else if (msg->type == ix::WebSocketMessageType::Close)
+                        {
+                            std::cerr << "Closed connection"
+                                      << " code " << msg->closeInfo.code << " reason "
+                                      << msg->closeInfo.reason << std::endl;
+                            state->webSocket().close(msg->closeInfo.code, msg->closeInfo.reason);
+                        }
+                        else if (msg->type == ix::WebSocketMessageType::Error)
+                        {
+                            std::stringstream ss;
+                            ss << "Connection error: " << msg->errorInfo.reason << std::endl;
+                            ss << "#retries: " << msg->errorInfo.retries << std::endl;
+                            ss << "Wait time(ms): " << msg->errorInfo.wait_time << std::endl;
+                            ss << "HTTP Status: " << msg->errorInfo.http_status << std::endl;
+                            std::cerr << ss.str();
+                        }
+                        else if (msg->type == ix::WebSocketMessageType::Message)
+                        {
+                            std::cerr << "Received " << msg->wireSize << " bytes from client" << std::endl;
+                            state->webSocket().send(msg->str, msg->binary);
+                        }
+                    });
             });
 
         auto res = server.listen();