temp

notes: Queue header has the same name as others, namely queue.hpp from boost this causes intellisense to fail to recognize symbols, bigger problems may arise. Some necessary changes to cmakes were made to allow compilation, aside from the queue header name, the cmake wasnt generating, some linking must be made public. Currently experimenting passing a lambda to the communicator with the strategy on how to get a message to send. This way the communciator can be kept decoupled from the queue implementation.
wazuh · Aug 12, 2024 · 0b062b7 · 0b062b7
1 parent 2e93ebd
commit 0b062b7
Show file tree

Hide file tree

Showing 12 changed files with 83 additions and 41 deletions.
diff --git a/src/agent/CMakeLists.txt b/src/agent/CMakeLists.txt
@@ -20,16 +20,16 @@ set(SOURCES
 add_subdirectory(agent_info)
 add_subdirectory(communicator)
 add_subdirectory(configuration_parser)
+add_subdirectory(queue)
 
 find_package(OpenSSL REQUIRED)
 find_package(Boost REQUIRED COMPONENTS asio beast)
 
 add_library(Agent ${SOURCES})
 target_include_directories(Agent PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include)
-target_link_libraries(Agent PUBLIC ConfigurationParser Communicator AgentInfo PRIVATE OpenSSL::SSL OpenSSL::Crypto Boost::asio Boost::beast)
+target_link_libraries(Agent PUBLIC ConfigurationParser Communicator AgentInfo queue PRIVATE OpenSSL::SSL OpenSSL::Crypto Boost::asio Boost::beast)
 
 if(BUILD_TESTS)
     enable_testing()
     add_subdirectory(tests)
-    add_subdirectory(queue)
 endif()
diff --git a/src/agent/communicator/include/communicator.hpp b/src/agent/communicator/include/communicator.hpp
@@ -20,9 +20,15 @@ namespace communicator
                      const std::function<std::string(std::string, std::string)> GetStringConfigValue);
 
         boost::asio::awaitable<void> WaitForTokenExpirationAndAuthenticate();
-        boost::asio::awaitable<void> GetCommandsFromManager(std::queue<std::string>& messageQueue);
-        boost::asio::awaitable<void> StatefulMessageProcessingTask(std::queue<std::string>& messageQueue);
-        boost::asio::awaitable<void> StatelessMessageProcessingTask(std::queue<std::string>& messageQueue);
+        boost::asio::awaitable<void>
+        GetCommandsFromManager(std::function<boost::asio::awaitable<std::string>()> getMessages,
+                               std::function<void(const std::string&)> onSuccess);
+        boost::asio::awaitable<void>
+        StatefulMessageProcessingTask(std::function<boost::asio::awaitable<std::string>()> getMessages,
+                                      std::function<void(const std::string&)> onSuccess);
+        boost::asio::awaitable<void>
+        StatelessMessageProcessingTask(std::function<boost::asio::awaitable<std::string>()> getMessages,
+                                       std::function<void(const std::string&)> onSuccess);
 
     private:
         long GetTokenRemainingSecs() const;

diff --git a/src/agent/communicator/include/http_client.hpp b/src/agent/communicator/include/http_client.hpp
@@ -46,11 +46,12 @@ namespace http_client
                           std::function<void()> onUnauthorized,
                           std::function<void(const std::string&)> onSuccess = {});
 
-    boost::asio::awaitable<void> Co_MessageProcessingTask(const std::string& token,
-                                                          HttpRequestParams params,
-                                                          std::function<std::string()> messageGetter,
-                                                          std::function<void()> onUnauthorized,
-                                                          std::function<void(const std::string&)> onSuccess = {});
+    boost::asio::awaitable<void>
+    Co_MessageProcessingTask(const std::string& token,
+                             HttpRequestParams params,
+                             std::function<boost::asio::awaitable<std::string>()> messageGetter,
+                             std::function<void()> onUnauthorized,
+                             std::function<void(const std::string&)> onSuccess = {});
 
     boost::beast::http::response<boost::beast::http::dynamic_body> PerformHttpRequest(const HttpRequestParams& params);
 

diff --git a/src/agent/communicator/src/communicator.cpp b/src/agent/communicator/src/communicator.cpp
@@ -67,15 +67,18 @@ namespace communicator
         return std::max(0L, static_cast<long>(m_tokenExpTimeInSeconds - now_seconds));
     }
 
-    boost::asio::awaitable<void> Communicator::GetCommandsFromManager(std::queue<std::string>& messageQueue)
+    boost::asio::awaitable<void>
+    Communicator::GetCommandsFromManager(std::function<boost::asio::awaitable<std::string>()> getMessages,
+                                         std::function<void(const std::string&)> onSuccess)
     {
         auto onAuthenticationFailed = [this]()
         {
             TryReAuthenticate();
         };
         const auto reqParams =
             http_client::HttpRequestParams(boost::beast::http::verb::get, m_managerIp, m_port, "/commands");
-        co_await http_client::Co_MessageProcessingTask(m_token, reqParams, {}, onAuthenticationFailed);
+        co_await http_client::Co_MessageProcessingTask(
+            m_token, reqParams, getMessages, onAuthenticationFailed, onSuccess);
     }
 
     boost::asio::awaitable<void> Communicator::WaitForTokenExpirationAndAuthenticate()
@@ -105,26 +108,32 @@ namespace communicator
         }
     }
 
-    boost::asio::awaitable<void> Communicator::StatefulMessageProcessingTask(std::queue<std::string>& messageQueue)
+    boost::asio::awaitable<void>
+    Communicator::StatefulMessageProcessingTask(std::function<boost::asio::awaitable<std::string>()> getMessages,
+                                                std::function<void(const std::string&)> onSuccess)
     {
         auto onAuthenticationFailed = [this]()
         {
             TryReAuthenticate();
         };
         const auto reqParams =
             http_client::HttpRequestParams(boost::beast::http::verb::post, m_managerIp, m_port, "/stateful");
-        co_await http_client::Co_MessageProcessingTask(m_token, reqParams, {}, onAuthenticationFailed);
+        co_await http_client::Co_MessageProcessingTask(
+            m_token, reqParams, getMessages, onAuthenticationFailed, onSuccess);
     }
 
-    boost::asio::awaitable<void> Communicator::StatelessMessageProcessingTask(std::queue<std::string>& messageQueue)
+    boost::asio::awaitable<void>
+    Communicator::StatelessMessageProcessingTask(std::function<boost::asio::awaitable<std::string>()> getMessages,
+                                                 std::function<void(const std::string&)> onSuccess)
     {
         auto onAuthenticationFailed = [this]()
         {
             TryReAuthenticate();
         };
         const auto reqParams =
             http_client::HttpRequestParams(boost::beast::http::verb::post, m_managerIp, m_port, "/stateless");
-        co_await http_client::Co_MessageProcessingTask(m_token, reqParams, {}, onAuthenticationFailed);
+        co_await http_client::Co_MessageProcessingTask(
+            m_token, reqParams, getMessages, onAuthenticationFailed, onSuccess);
     }
 
     void Communicator::TryReAuthenticate()

diff --git a/src/agent/communicator/src/http_client.cpp b/src/agent/communicator/src/http_client.cpp
@@ -96,11 +96,12 @@ namespace http_client
         std::cout << "Response body: " << boost::beast::buffers_to_string(res.body().data()) << std::endl;
     }
 
-    boost::asio::awaitable<void> Co_MessageProcessingTask(const std::string& token,
-                                                          HttpRequestParams reqParams,
-                                                          std::function<std::string()> messageGetter,
-                                                          std::function<void()> onUnauthorized,
-                                                          std::function<void(const std::string&)> onSuccess)
+    boost::asio::awaitable<void>
+    Co_MessageProcessingTask(const std::string& token,
+                             HttpRequestParams reqParams,
+                             std::function<boost::asio::awaitable<std::string>()> messageGetter,
+                             std::function<void()> onUnauthorized,
+                             std::function<void(const std::string&)> onSuccess)
     {
         using namespace std::chrono_literals;
 
@@ -129,7 +130,7 @@ namespace http_client
                 continue;
             }
 
-            reqParams.body = messageGetter ? messageGetter() : "";
+            reqParams.body = messageGetter ? co_await messageGetter() : "";
             reqParams.token = token;
             auto req = CreateHttpRequest(reqParams);
 

diff --git a/src/agent/include/agent.hpp b/src/agent/include/agent.hpp
@@ -3,10 +3,10 @@
 #include <agent_info.hpp>
 #include <communicator.hpp>
 #include <configuration_parser.hpp>
+#include <message_queue.hpp>
 #include <signal_handler.hpp>
 #include <task_manager.hpp>
 
-#include <queue>
 #include <string>
 
 class Agent
@@ -18,7 +18,7 @@ class Agent
     void Run();
 
 private:
-    std::queue<std::string> m_messageQueue;
+    MultiTypeQueue m_messageQueue;
 
     SignalHandler m_signalHandler;
     TaskManager m_taskManager;

diff --git a/src/agent/queue/CMakeLists.txt b/src/agent/queue/CMakeLists.txt
@@ -6,21 +6,15 @@ set(CMAKE_CXX_STANDARD_REQUIRED True)
 
 set(CMAKE_CXX_FLAGS "-Wall -Wextra -Wunused -pthread")
 
-include_directories(${CMAKE_SOURCE_DIR}/include)
-
 find_package(SQLiteCpp REQUIRED)
 find_package(nlohmann_json REQUIRED)
 find_package(fmt REQUIRED)
 find_package(Boost REQUIRED COMPONENTS asio beast)
 
-add_library(queue
-	src/sqlitestorage.cpp
-	src/queue.cpp
-)
-
-target_include_directories(queue PRIVATE include ${SQLiteCpp_INCLUDE_DIRS})
+add_library(queue src/sqlitestorage.cpp src/message_queue.cpp)
+target_include_directories(queue PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include PRIVATE ${SQLiteCpp_INCLUDE_DIRS})
 target_link_libraries(queue PRIVATE SQLiteCpp nlohmann_json::nlohmann_json fmt::fmt Boost::asio)
 
-if(BUILD_TESTS)
-    add_subdirectory(tests)
-endif()
+# if(BUILD_TESTS)
+# add_subdirectory(tests)
+# endif()
diff --git a/src/agent/queue/include/queue.hpp → src/agent/queue/include/message_queue.hpp b/src/agent/queue/include/queue.hpp → src/agent/queue/include/message_queue.hpp
diff --git a/src/agent/queue/src/main.cpp b/src/agent/queue/src/main.cpp
@@ -2,7 +2,7 @@
 #include <string>
 #include <thread>
 
-#include "queue.hpp"
+#include "message_queue.hpp"
 
 int main()
 {

diff --git a/src/agent/queue/src/queue.cpp → src/agent/queue/src/message_queue.cpp b/src/agent/queue/src/queue.cpp → src/agent/queue/src/message_queue.cpp
@@ -10,7 +10,7 @@
 #include <utility>
 #include <vector>
 
-#include "queue.hpp"
+#include "message_queue.hpp"
 
 int MultiTypeQueue::push(Message message, bool shouldWait)
 {

diff --git a/src/agent/queue/tests/queue_test.cpp b/src/agent/queue/tests/queue_test.cpp
@@ -12,7 +12,7 @@
 #include <boost/asio.hpp>
 #include <boost/asio/experimental/co_spawn.hpp>
 
-#include "queue.hpp"
+#include "message_queue.hpp"
 #include "queue_test.hpp"
 
 using json = nlohmann::json;

diff --git a/src/agent/src/agent.cpp b/src/agent/src/agent.cpp
@@ -1,7 +1,35 @@
 #include <agent.hpp>
 
+#include <shared.hpp>
+
+#include <string>
 #include <thread>
 
+namespace
+{
+    template<MessageType T>
+    auto getMessagesFromQueue(MultiTypeQueue& multiTypeQueue)
+    {
+        return [&multiTypeQueue]() -> boost::asio::awaitable<std::string>
+        {
+            std::cout << "Getting messages from queue\n";
+            const auto message = co_await multiTypeQueue.getNextAwaitable(T);
+            co_return message.data.dump();
+        };
+    }
+
+    template<MessageType T>
+    auto popMessagesFromQueue(MultiTypeQueue& multiTypeQueue)
+    {
+        return [&multiTypeQueue](const std::string& response) -> void
+        {
+            std::cout << "Response: " << response << '\n';
+            std::cout << "Popping messages from queue\n";
+            multiTypeQueue.pop(T);
+        };
+    }
+} // namespace
+
 Agent::Agent()
     : m_communicator(m_agentInfo.GetUUID(),
                      [this](std::string table, std::string key) -> std::string
@@ -18,9 +46,12 @@ Agent::~Agent()
 void Agent::Run()
 {
     m_taskManager.EnqueueTask(m_communicator.WaitForTokenExpirationAndAuthenticate());
-    m_taskManager.EnqueueTask(m_communicator.GetCommandsFromManager(m_messageQueue));
-    m_taskManager.EnqueueTask(m_communicator.StatefulMessageProcessingTask(m_messageQueue));
-    m_taskManager.EnqueueTask(m_communicator.StatelessMessageProcessingTask(m_messageQueue));
+    m_taskManager.EnqueueTask(m_communicator.GetCommandsFromManager(getMessagesFromQueue<COMMAND>(m_messageQueue),
+                                                                    popMessagesFromQueue<COMMAND>(m_messageQueue)));
+    m_taskManager.EnqueueTask(m_communicator.StatefulMessageProcessingTask(
+        getMessagesFromQueue<STATEFUL>(m_messageQueue), popMessagesFromQueue<STATEFUL>(m_messageQueue)));
+    m_taskManager.EnqueueTask(m_communicator.StatelessMessageProcessingTask(
+        getMessagesFromQueue<STATELESS>(m_messageQueue), popMessagesFromQueue<STATELESS>(m_messageQueue)));
 
     m_signalHandler.WaitForSignal();
 }