1/*
2 * IXWebSocketProxyServer.cpp
3 * Author: Benjamin Sergeant
4 * Copyright (c) 2018 Machine Zone, Inc. All rights reserved.
5 */
6
7#include "IXWebSocketProxyServer.h"
8
9#include "IXWebSocketServer.h"
10#include <sstream>
11
12namespace ix
13{
14 class ProxyConnectionState : public ix::ConnectionState
15 {
16 public:
17 ProxyConnectionState()
18 : _connected(false)
19 {
20 }
21
22 ix::WebSocket& webSocket()
23 {
24 return _serverWebSocket;
25 }
26
27 bool isConnected()
28 {
29 return _connected;
30 }
31
32 void setConnected()
33 {
34 _connected = true;
35 }
36
37 private:
38 ix::WebSocket _serverWebSocket;
39 bool _connected;
40 };
41
42 int websocket_proxy_server_main(int port,
43 const std::string& hostname,
44 const ix::SocketTLSOptions& tlsOptions,
45 const std::string& remoteUrl,
46 const RemoteUrlsMapping& remoteUrlsMapping,
47 bool /*verbose*/)
48 {
49 ix::WebSocketServer server(port, hostname);
50 server.setTLSOptions(tlsOptions);
51
52 auto factory = []() -> std::shared_ptr<ix::ConnectionState> {
53 return std::make_shared<ProxyConnectionState>();
54 };
55 server.setConnectionStateFactory(factory);
56
57 server.setOnConnectionCallback(
58 [remoteUrl, remoteUrlsMapping](std::weak_ptr<ix::WebSocket> webSocket,
59 std::shared_ptr<ConnectionState> connectionState) {
60 auto state = std::dynamic_pointer_cast<ProxyConnectionState>(connectionState);
61 auto remoteIp = connectionState->getRemoteIp();
62
63 // Server connection
64 state->webSocket().setOnMessageCallback(
65 [webSocket, state, remoteIp](const WebSocketMessagePtr& msg) {
66 if (msg->type == ix::WebSocketMessageType::Close)
67 {
68 state->setTerminated();
69 }
70 else if (msg->type == ix::WebSocketMessageType::Message)
71 {
72 auto ws = webSocket.lock();
73 if (ws)
74 {
75 ws->send(msg->str, msg->binary);
76 }
77 }
78 });
79
80 // Client connection
81 auto ws = webSocket.lock();
82 if (ws)
83 {
84 ws->setOnMessageCallback([state, remoteUrl, remoteUrlsMapping](
85 const WebSocketMessagePtr& msg) {
86 if (msg->type == ix::WebSocketMessageType::Open)
87 {
88 // Connect to the 'real' server
89 std::string url(remoteUrl);
90
91 // maybe we want a different url based on the mapping
92 std::string host = msg->openInfo.headers["Host"];
93 auto it = remoteUrlsMapping.find(host);
94 if (it != remoteUrlsMapping.end())
95 {
96 url = it->second;
97 }
98
99 // append the uri to form the full url
100 // (say ws://localhost:1234/foo/?bar=baz)
101 url += msg->openInfo.uri;
102
103 state->webSocket().setUrl(url);
104 state->webSocket().disableAutomaticReconnection();
105 state->webSocket().start();
106
107 // we should sleep here for a bit until we've established the
108 // connection with the remote server
109 while (state->webSocket().getReadyState() != ReadyState::Open)
110 {
111 std::this_thread::sleep_for(std::chrono::milliseconds(10));
112 }
113 }
114 else if (msg->type == ix::WebSocketMessageType::Close)
115 {
116 state->webSocket().close(msg->closeInfo.code, msg->closeInfo.reason);
117 }
118 else if (msg->type == ix::WebSocketMessageType::Message)
119 {
120 state->webSocket().send(msg->str, msg->binary);
121 }
122 });
123 }
124 });
125
126 auto res = server.listen();
127 if (!res.first)
128 {
129 return 1;
130 }
131
132 server.start();
133 server.wait();
134
135 return 0;
136 }
137} // namespace ix
138