feat: Adds timeout to HttpsSocket::Connect()

wazuh · Dec 12, 2024 · 59f5e23 · 59f5e23
1 parent 63d07a5
commit 59f5e23
Showing 1 changed file with 43 additions and 15 deletions.
diff --git a/src/agent/communicator/src/https_socket.hpp b/src/agent/communicator/src/https_socket.hpp
@@ -27,26 +27,54 @@ namespace http_client
         /// @brief Connects the socket to the given endpoints
         /// @param endpoints The endpoints to connect to
         /// @param ec The error code, if any occurred
-        void Connect([[maybe_unused]] boost::asio::io_context& io_context,
+        void Connect(boost::asio::io_context& io_context,
                      const boost::asio::ip::tcp::resolver::results_type& endpoints,
                      boost::system::error_code& ec,
-                     [[maybe_unused]] const std::chrono::seconds timeout = std::chrono::seconds(2)) override
+                     const std::chrono::seconds timeOut = std::chrono::seconds(TIMEOUT_DEFAULT)) override
         {
             try
             {
-                boost::asio::connect(m_ssl_socket.next_layer(), endpoints.begin(), endpoints.end(), ec);
-                if (ec)
-                {
-                    LogDebug("Connect failed: {}", ec.message());
-                    return;
-                }
-
-                m_ssl_socket.handshake(boost::asio::ssl::stream_base::client, ec);
-                if (ec)
-                {
-                    LogDebug("Handshake failed: {}", ec.message());
-                    return;
-                }
+                boost::asio::steady_timer timer(io_context);
+                bool connectionSuccess = false;
+                bool timer_expired = false;
+
+                // Start the asynchronous connect operation
+                boost::asio::async_connect(m_ssl_socket.lowest_layer(),
+                                           endpoints,
+                                           [&](const boost::system::error_code& errorCode, const auto&)
+                                           {
+                                               if (!timer_expired && !errorCode)
+                                               {
+                                                   timer.cancel(); // Cancel the timer if connect completes first
+                                                   connectionSuccess = true;
+                                                   ec = errorCode;
+
+                                                   m_ssl_socket.handshake(boost::asio::ssl::stream_base::client, ec);
+                                                   if (ec)
+                                                   {
+                                                       LogDebug("Handshake failed: {}", ec.message());
+                                                       return;
+                                                   }
+
+                                                   LogDebug("Connected successfully");
+                                               }
+                                           });
+
+                // Start the timer for the timeout
+                timer.expires_after(timeOut);
+                timer.async_wait(
+                    [&](const boost::system::error_code& errorCode)
+                    {
+                        if (!errorCode && !connectionSuccess)
+                        {
+                            timer_expired = true;
+                            // Close();
+                            ec = boost::asio::error::timed_out;
+                            LogDebug("Connect operation timed out");
+                        }
+                    });
+
+                io_context.run_for(timeOut); // Run for 2 seconds();
             }
             catch (const std::exception& e)
             {