@@ -80,11 +80,14 @@ namespace aws { namespace iot { namespace securedtunneling {
8080 wss{ nullptr },
8181 wss_resolver{ io_ctx },
8282 wss_response{ },
83+ num_active_connections{ 0 },
8384 stream_id{ -1 },
8485 service_id{ " " },
8586 serviceId_to_streamId_map{},
8687 serviceId_to_tcp_server_map{},
8788 serviceId_to_tcp_client_map{},
89+ serviceId_to_control_message_handler_map{},
90+ serviceId_to_data_message_handler_map{},
8891 bind_address_actual{ },
8992 is_web_socket_reading{ false },
9093 is_service_ids_received{ false },
@@ -105,6 +108,8 @@ namespace aws { namespace iot { namespace securedtunneling {
105108 // debuggability.
106109 boost::beast::websocket::response_type wss_response;
107110
111+ std::atomic_uint16_t num_active_connections;
112+
108113 // represents the current stream ID to expect data from
109114 // care should be taken how(if) this is updated directly
110115 // To be deleted
@@ -113,6 +118,8 @@ namespace aws { namespace iot { namespace securedtunneling {
113118 std::unordered_map<std::string, std::int32_t > serviceId_to_streamId_map;
114119 std::unordered_map<std::string, tcp_server::pointer> serviceId_to_tcp_server_map;
115120 std::unordered_map<std::string, tcp_client::pointer> serviceId_to_tcp_client_map;
121+ std::unordered_map<std::string, std::function<bool (message const &)>> serviceId_to_control_message_handler_map;
122+ std::unordered_map<std::string, std::function<bool (message const &)>> serviceId_to_data_message_handler_map;
116123 std::string bind_address_actual;
117124 // flag set to true while web socket data is being drained
118125 // necessary for better TCP socket recovery rather than destroying
@@ -152,15 +159,15 @@ namespace aws { namespace iot { namespace securedtunneling {
152159
153160 int run_proxy ();
154161 private:
162+ void update_message_handlers (tcp_adapter_context &tac, std::function<bool (message const &)> handler);
155163 void setup_tcp_socket (tcp_adapter_context &tac, std::string const & service_id);
156164 void setup_tcp_sockets (tcp_adapter_context &tac);
157165 // setup async io flow to connect tcp socket to the adapter config's data host/port
158- void async_setup_dest_tcp_socket (tcp_adapter_context &tac, std::string const & service_id);
159- void async_setup_dest_tcp_socket_retry (tcp_adapter_context &tac, std::shared_ptr<basic_retry_config> retry_config, std::string const & service_id);
166+ void async_setup_dest_tcp_socket (tcp_adapter_context &tac, std::string const & service_id, uint32_t const & connection_id, bool is_first_connection );
167+ void async_setup_dest_tcp_socket_retry (tcp_adapter_context &tac, std::shared_ptr<basic_retry_config> retry_config, std::string const & service_id, uint32_t const & connection_id, bool is_first_connection );
160168 void async_setup_source_tcp_sockets (tcp_adapter_context &tac);
161169 void async_setup_source_tcp_socket_retry (tcp_adapter_context &tac, std::shared_ptr<basic_retry_config> retry_config, std::string service_id);
162- void initialize_tcp_clients (tcp_adapter_context &tac);
163- void initialize_tcp_servers (tcp_adapter_context &tac);
170+ void do_accept_tcp_connection (tcp_adapter_context &tac, std::shared_ptr<basic_retry_config> retry_config, std::string service_id, std::uint16_t local_port, bool is_first_connection);
164171 void setup_web_socket (tcp_adapter_context &tac);
165172 // setup async web socket, and as soon as connection is up, setup async ping schedule
166173 void async_setup_web_socket (tcp_adapter_context &tac);
@@ -169,10 +176,13 @@ namespace aws { namespace iot { namespace securedtunneling {
169176 // then the reset is intentionally reset via web socket, and retries
170177 // occur definitely (regardless of retry configuration)
171178 void tcp_socket_reset_all (tcp_adapter_context &tac, std::function<void ()> post_reset_operation);
172- void tcp_socket_reset (tcp_adapter_context &tac, std::string service_id, std::function<void ()> post_reset_operation);
173- tcp_connection::pointer get_tcp_connection (tcp_adapter_context &tac, std::string service_id);
179+ void tcp_socket_reset_init (tcp_adapter_context &tac, std::string service_id, std::function<void ()> post_reset_operation);
180+ void tcp_socket_reset (tcp_adapter_context &tac, std::string service_id, uint32_t connection_id, std::function<void ()> post_reset_operation);
181+ void tcp_socket_close (tcp_adapter_context &tac, std::string service_id, uint32_t connection_id);
182+ tcp_connection::pointer get_tcp_connection (tcp_adapter_context &tac, std::string service_id, uint32_t connection_id);
174183
175- void tcp_socket_error (tcp_adapter_context &tac, boost::system::error_code const &_ec, std::string const & service_id);
184+ void delete_tcp_socket (tcp_adapter_context &tac, std::string const & service_id, uint32_t const & connection_id);
185+ void tcp_socket_error (tcp_adapter_context &tac, boost::system::error_code const &_ec, std::string const & service_id, uint32_t const & connection_id);
176186
177187 // sets up a web socket read loop that will read, and ignore most messages until a stream start
178188 // is read and then do something with it (likely, connect to configured endpoint)
@@ -197,22 +207,21 @@ namespace aws { namespace iot { namespace securedtunneling {
197207 // invokes after_setup_web_socket_read_until_stream_start() after stream start is encountered
198208 bool async_wait_for_stream_start (tcp_adapter_context &tac, message const &message);
199209 bool async_wait_for_service_ids (tcp_adapter_context &tac);
200- void async_tcp_socket_read_loop (tcp_adapter_context &tac, std::string const & service_id);
210+ void async_tcp_socket_read_loop (tcp_adapter_context &tac, std::string const & service_id, uint32_t const & connection_id );
201211
202212 // below loop does continuous writes to TCP socket from the TCP adapter
203213 // context's tcp_write_buffer. After consuming chunks out of the buffer
204- // the behavior will be to check
205- void async_tcp_write_buffer_drain (tcp_adapter_context &tac, std::string service_id);
214+ // the behavior will be to check
215+ void async_tcp_write_buffer_drain (tcp_adapter_context &tac, std::string service_id, uint32_t connection_id );
206216
207- void async_setup_bidirectional_data_transfers (tcp_adapter_context &tac, std::string const & service_id);
208- void async_setup_web_socket_write_buffer_drain (tcp_adapter_context &tac, std::string const & service_id);
217+ void async_setup_bidirectional_data_transfers (tcp_adapter_context &tac, std::string const & service_id, uint32_t const & connection_id );
218+ void async_setup_web_socket_write_buffer_drain (tcp_adapter_context &tac, std::string const & service_id, uint32_t const & connection_id );
209219
210220 // returns a boolean that indicates if another web socket data read message can be put
211221 // onto the tcp write buffer. We have no way of knowing what the next message is and if
212222 // it will be too big to process, thus we don't do the read applying back pressure on
213223 // the socket. Implicitly, this means that an async_read is not happening on the web socket
214224 bool tcp_has_enough_write_buffer_space (tcp_connection::pointer connection);
215- bool tcp_has_enough_write_buffer_space (tcp_adapter_context const &tac);
216225
217226 // returns a boolean that indicates if another tcp socket read's data can be put on the
218227 // web socket write buffer. It's a bit different from tcp write buffer space requirements
@@ -226,8 +235,11 @@ namespace aws { namespace iot { namespace securedtunneling {
226235 bool is_valid_stream_id (tcp_adapter_context const & tac, message const &message);
227236
228237 void async_send_message (tcp_adapter_context &tac, message const &message);
229- void async_send_stream_start (tcp_adapter_context &tac, std::string const & service_id);
230- void async_send_stream_reset (tcp_adapter_context &tac, std::string const & service_id);
238+ void async_send_message (tcp_adapter_context &tac, message const &message, std::string const & service_id, uint32_t const & connection_id);
239+ void async_send_stream_start (tcp_adapter_context &tac, std::string const & service_id, uint32_t const & connection_id);
240+ void async_send_stream_reset (tcp_adapter_context &tac, std::string const & service_id, uint32_t const & connection_id);
241+ void async_send_connection_start (tcp_adapter_context &tac, std::string const & service_id, uint32_t const & connection_id);
242+ void async_send_connection_reset (tcp_adapter_context &tac, std::string const & service_id, uint32_t const & connection_id);
231243
232244 // handler for successfully sent ping will delay the next one
233245 void async_ping_handler_loop (tcp_adapter_context &tac,
@@ -239,16 +251,14 @@ namespace aws { namespace iot { namespace securedtunneling {
239251 void clear_ws_buffers (tcp_adapter_context &tac);
240252 void clear_tcp_connection_buffers (tcp_connection::pointer connection);
241253
242- void tcp_socket_ensure_closed (boost::asio::ip::tcp::socket & tcp_socket);
243-
244254 // closes the websocket connection
245255 // 1 - shutdown the receive side of TCP
246256 // 2 - drain the web socket write buffer
247257 // 3 - send a web socket close frame
248258 // 4 - perform teardown procedure on websocket
249259 void web_socket_close_and_stop (tcp_adapter_context &tac);
250260
251- void async_resolve_destination_for_connect (tcp_adapter_context &tac, std::shared_ptr<basic_retry_config> retry_config, std::string const & service_id, boost::system::error_code const &ec, tcp::resolver::results_type results);
261+ void async_resolve_destination_for_connect (tcp_adapter_context &tac, std::shared_ptr<basic_retry_config> retry_config, std::string const & service_id, uint32_t const & connection_id, boost::system::error_code const &ec, tcp::resolver::results_type results);
252262
253263 bool process_incoming_websocket_buffer (tcp_adapter_context &tac, boost::beast::multi_buffer &message_buffer);
254264
@@ -264,7 +274,7 @@ namespace aws { namespace iot { namespace securedtunneling {
264274
265275 bool fall_back_to_v1_message_format (std::unordered_map<std::string, std::string> const & serviceId_to_endpoint_map);
266276
267- void async_send_message_to_web_socket (tcp_adapter_context &tac, std::shared_ptr<boost::beast::flat_buffer> const & ss, std::string const & service_id);
277+ void async_send_message_to_web_socket (tcp_adapter_context &tac, std::shared_ptr<boost::beast::flat_buffer> const & ss, std::string const & service_id, uint32_t const & connection_id );
268278
269279 void async_setup_destination_tcp_sockets (tcp_adapter_context &tac);
270280
0 commit comments