From 72bd1d62107ca3d194965312a9c1f879948def9e Mon Sep 17 00:00:00 2001 From: Luis Enrique Chico Capistrano Date: Tue, 28 Jan 2025 00:32:06 -0300 Subject: [PATCH] feat: Add test for self-restart functionality --- src/agent/restart_handler/CMakeLists.txt | 9 +- .../include/restart_handler.hpp | 3 + .../src/restart_handler_unix.cpp | 12 +- .../src/restart_handler_unix.hpp | 13 +- .../restart_handler/tests/CMakeLists.txt | 9 + .../tests/restart_handler_tests.cpp | 165 ++++++++++++++++++ 6 files changed, 197 insertions(+), 14 deletions(-) create mode 100644 src/agent/restart_handler/tests/CMakeLists.txt create mode 100644 src/agent/restart_handler/tests/restart_handler_tests.cpp diff --git a/src/agent/restart_handler/CMakeLists.txt b/src/agent/restart_handler/CMakeLists.txt index 70a411b568..78ec5f0c98 100644 --- a/src/agent/restart_handler/CMakeLists.txt +++ b/src/agent/restart_handler/CMakeLists.txt @@ -14,8 +14,15 @@ else() endif() add_library(RestartHandler ${SOURCES}) -target_include_directories(RestartHandler PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include) +target_include_directories(RestartHandler + PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src) target_link_libraries(RestartHandler PUBLIC Boost::asio CommandEntry PRIVATE Logger) include(../../cmake/ConfigureTarget.cmake) configure_target(RestartHandler) + +if(BUILD_TESTS) + enable_testing() + add_subdirectory(tests) +endif() diff --git a/src/agent/restart_handler/include/restart_handler.hpp b/src/agent/restart_handler/include/restart_handler.hpp index 64ac0af4e5..dcb6607d1f 100644 --- a/src/agent/restart_handler/include/restart_handler.hpp +++ b/src/agent/restart_handler/include/restart_handler.hpp @@ -15,6 +15,9 @@ namespace restart_handler /// @brief A list of command line arguments used to restart the agent process. static std::vector startupCmdLineArgs; + /// @brief Timeout for the restart operation until force a hard restart. + static const int timeoutInSecs = 30; + /// @brief RestartHandler class. explicit RestartHandler(); diff --git a/src/agent/restart_handler/src/restart_handler_unix.cpp b/src/agent/restart_handler/src/restart_handler_unix.cpp index a27410796e..df53c64e90 100644 --- a/src/agent/restart_handler/src/restart_handler_unix.cpp +++ b/src/agent/restart_handler/src/restart_handler_unix.cpp @@ -1,6 +1,5 @@ -#include "restart_handler_unix.hpp" -#include #include +#include #include #include @@ -31,13 +30,10 @@ namespace restart_handler } } - void StopAgent() + void StopAgent(const pid_t pid, const int timeoutInSecs) { - const int timeoutInSecs = 30; const time_t startTime = time(nullptr); - pid_t pid = getppid(); - // Shutdown Gracefully kill(pid, SIGTERM); @@ -45,7 +41,7 @@ namespace restart_handler { if (kill(pid, 0) != 0) { - LogInfo("Agent gracefully stopped."); + LogInfo("Agent stopped."); break; } @@ -70,7 +66,7 @@ namespace restart_handler else if (pid == 0) { // Child process - StopAgent(); + StopAgent(getppid(), RestartHandler::timeoutInSecs); LogInfo("Starting wazuh agent in a new process."); diff --git a/src/agent/restart_handler/src/restart_handler_unix.hpp b/src/agent/restart_handler/src/restart_handler_unix.hpp index 2d8c05f961..3ad189ec73 100644 --- a/src/agent/restart_handler/src/restart_handler_unix.hpp +++ b/src/agent/restart_handler/src/restart_handler_unix.hpp @@ -1,6 +1,5 @@ #pragma once #include -#include #include #include @@ -13,11 +12,15 @@ namespace restart_handler /// @return true if systemctl is available and running as a systemd service, otherwise returns false. bool UsingSystemctl(); - /// @brief Stops the agent by terminating the child process. + /// @brief Stops the agent by terminating its process. /// - /// This function sends a SIGTERM signal to the agent process to stop it. If the agent process - /// does not stop within a specified timeout period (30 seconds), it forces termination with a SIGKILL signal. - void StopAgent(); + /// This function sends a SIGTERM signal to the agent process to request its termination. If the agent + /// process does not stop within the specified timeout period (30 seconds), a SIGKILL signal is sent to force + /// termination. + /// + /// @param pid The process ID of the agent to be stopped. + /// @param timeout The maximum time (in seconds) to wait for the agent to stop before forcing termination. + void StopAgent(const pid_t pid, const int timeout); /// @brief Restarts the module by forking a new process. /// diff --git a/src/agent/restart_handler/tests/CMakeLists.txt b/src/agent/restart_handler/tests/CMakeLists.txt new file mode 100644 index 0000000000..c9b4a1eea1 --- /dev/null +++ b/src/agent/restart_handler/tests/CMakeLists.txt @@ -0,0 +1,9 @@ +find_package(GTest CONFIG REQUIRED) + +if(UNIX) + add_executable(restart_handler_test restart_handler_tests.cpp) + configure_target(restart_handler_test) + target_include_directories(restart_handler_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../src) + target_link_libraries(restart_handler_test PRIVATE RestartHandler GTest::gtest GTest::gmock GTest::gmock_main) + add_test(NAME RestartHandler COMMAND restart_handler_test) +endif() diff --git a/src/agent/restart_handler/tests/restart_handler_tests.cpp b/src/agent/restart_handler/tests/restart_handler_tests.cpp new file mode 100644 index 0000000000..d81cce4ec4 --- /dev/null +++ b/src/agent/restart_handler/tests/restart_handler_tests.cpp @@ -0,0 +1,165 @@ +#include +#include + +#include + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +using namespace testing; + +namespace +{ + // NOLINTBEGIN(cppcoreguidelines-avoid-reference-signal-handler-definitions) + void SigchldHandler(int status) + { + while (waitpid(-1, &status, WNOHANG) > 0) + { + // The child process has terminated, no action needed here + } + } + + // NOLINTEND(cppcoreguidelines-avoid-reference-signal-handler-definitions) + + void SigtermHandler([[maybe_unused]] int signum) + { + exit(0); + } +} // namespace + +TEST(TestUsingSystemctl, ReturnsTrueWhenSystemctlIsAvailable) +{ + setenv("INVOCATION_ID", "testing", 1); + if (0 == std::system("which systemctl > /dev/null 2>&1")) + { + EXPECT_TRUE(restart_handler::UsingSystemctl()); + } + else + { + EXPECT_FALSE(restart_handler::UsingSystemctl()); + } +} + +TEST(TestUsingSystemctl, ReturnsFalseWhenSystemctlIsNotAvailable) +{ + unsetenv("INVOCATION_ID"); + EXPECT_FALSE(restart_handler::UsingSystemctl()); +} + +TEST(StopAgentTest, GratefullyStop) +{ + testing::internal::CaptureStdout(); + signal(SIGCHLD, SigchldHandler); + + pid_t pid = fork(); + + if (pid < 0) + { + FAIL() << "Fork failed!"; + } + if (pid == 0) + { + // Child process (Simulate agent running) + const int timeout = 60; + std::this_thread::sleep_for(std::chrono::seconds(timeout)); + } + + // Parent process + std::this_thread::sleep_for(std::chrono::seconds(1)); + + const int timeout = 10; + restart_handler::StopAgent(pid, timeout); + + std::string stdout_logs = testing::internal::GetCapturedStdout(); + EXPECT_EQ(stdout_logs.find("Timeout reached! Forcing agent process termination."), std::string::npos); + EXPECT_NE(stdout_logs.find("Agent stopped"), std::string::npos); +} + +TEST(StopAgentTest, ForceStop) +{ + testing::internal::CaptureStdout(); + signal(SIGCHLD, SigchldHandler); + signal(SIGTERM, SIG_IGN); // Ignore SIGTERM + + pid_t pid = fork(); + + if (pid < 0) + { + FAIL() << "Fork failed!"; + } + if (pid == 0) + { + // Child process (Simulate agent running) + // NOLINTNEXTLINE(cppcoreguidelines-avoid-magic-numbers) + std::this_thread::sleep_for(std::chrono::seconds(60)); + } + + // Parent process + std::this_thread::sleep_for(std::chrono::seconds(1)); + + // NOLINTNEXTLINE(cppcoreguidelines-avoid-magic-numbers) + restart_handler::StopAgent(pid, 2); + + std::string stdout_logs = testing::internal::GetCapturedStdout(); + EXPECT_NE(stdout_logs.find("Timeout reached! Forcing agent process termination."), std::string::npos); + EXPECT_NE(stdout_logs.find("Agent stopped"), std::string::npos); +} + +TEST(RestartWithForkTest, ShouldRestartUsingFork) +{ + + testing::internal::CaptureStdout(); + pid_t pid = fork(); + + if (pid < 0) + { + FAIL() << "Fork failed!"; + } + if (pid == 0) + { + // Child process (Simulate agent running) + setsid(); + signal(SIGTERM, SigtermHandler); + + boost::asio::io_context io_context; + + boost::asio::co_spawn( + io_context, + []() -> boost::asio::awaitable + { + const char* args[] = {"/usr/bin/sleep", "3"}; + // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast) + restart_handler::RestartHandler::SetCommandLineArguments(2, const_cast(args)); + + const auto result = co_await restart_handler::RestartWithFork(); + }(), + boost::asio::detached); + + io_context.run(); + } + + // Parent process + int status {}; + waitpid(pid, &status, 0); + + std::string stdout_logs = testing::internal::GetCapturedStdout(); + EXPECT_EQ(stdout_logs.find("Agent stopped"), std::string::npos); + EXPECT_EQ(stdout_logs.find("Starting wazuh agent in a new process"), std::string::npos); +} + +int main(int argc, char** argv) +{ + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +}