diff --git a/C2.md b/C2.md index f934cd0e77..1f171fcc4f 100644 --- a/C2.md +++ b/C2.md @@ -28,6 +28,7 @@ options defined are located in minifi.properties. - [Base Options](#base-options) - [Flow Id and URL](#flow-id-and-url) - [Agent Identifier Fallback](#agent-identifier-fallback) + - [C2 Connection Timeout](#c2-connection-timeout) - [Metrics](#metrics) - [UpdatePolicies](#updatepolicies) - [Triggers](#triggers) @@ -136,6 +137,15 @@ It is possible to set a persistent fallback agent id. This is needed so that the # in minifi.properties nifi.c2.agent.identifier.fallback=my_fallback_id + +#### C2 Connection Timeout + +There is an option to shut down the agent if the connection to the C2 server exceeds a specified timeout. This is controlled by the `nifi.c2.connection.timeout` property, which defines the maximum time allowed without a successful heartbeat response. If the connection is not established (or is lost) within this period, the agent will trigger a shutdown with exit code 90. If the agent is shut down with exit 90 neither the systemd service nor the Windows service will restart it. The feature is off by default, and can be enabled by setting a non-zero value. If the timeout is set to a value less than twice the heartbeat period, it will be automatically adjusted to twice the heartbeat period. + + # in minifi.properties + # timeout for C2 connection, setting to 0 disables the timeout + nifi.c2.connection.timeout=30 min + ### Metrics Command and Control metrics can be used to send metrics through the heartbeat or via the DESCRIBE diff --git a/behave_framework/src/minifi_behave/containers/container_linux.py b/behave_framework/src/minifi_behave/containers/container_linux.py index 034f53c69b..9df8567f68 100644 --- a/behave_framework/src/minifi_behave/containers/container_linux.py +++ b/behave_framework/src/minifi_behave/containers/container_linux.py @@ -330,6 +330,19 @@ def exited(self) -> bool: except Exception: return False + @property + def exit_code(self) -> int: + if not self.container: + return -1 + try: + self.container.reload() + return self.container.attrs["State"].get("ExitCode", -1) + except docker.errors.NotFound: + self.container = None + return -1 + except Exception: + return -1 + def get_number_of_files(self, directory_path: str) -> int: if not self.container: logger.warning("Container not running") diff --git a/behave_framework/src/minifi_behave/containers/container_protocol.py b/behave_framework/src/minifi_behave/containers/container_protocol.py index cc167c33b3..12e832c3b1 100644 --- a/behave_framework/src/minifi_behave/containers/container_protocol.py +++ b/behave_framework/src/minifi_behave/containers/container_protocol.py @@ -44,6 +44,9 @@ def get_logs(self) -> str: ... @property def exited(self) -> bool: ... + @property + def exit_code(self) -> int: ... + def get_number_of_files(self, directory_path: str) -> int: ... def verify_file_contents(self, directory_path: str, expected_contents: list[str]) -> bool: ... diff --git a/behave_framework/src/minifi_behave/containers/container_windows.py b/behave_framework/src/minifi_behave/containers/container_windows.py index 5117562b70..2e5f3d68c1 100644 --- a/behave_framework/src/minifi_behave/containers/container_windows.py +++ b/behave_framework/src/minifi_behave/containers/container_windows.py @@ -318,6 +318,19 @@ def exited(self) -> bool: except Exception: return False + @property + def exit_code(self) -> int: + if not self.container: + return -1 + try: + self.container.reload() + return self.container.attrs["State"].get("ExitCode", -1) + except docker.errors.NotFound: + self.container = None + return -1 + except Exception: + return -1 + def get_number_of_files(self, directory_path: str) -> int: if not self.container: return -1 diff --git a/behave_framework/src/minifi_behave/steps/checking_steps.py b/behave_framework/src/minifi_behave/steps/checking_steps.py index 83ff073b7d..a67f64c4d3 100644 --- a/behave_framework/src/minifi_behave/steps/checking_steps.py +++ b/behave_framework/src/minifi_behave/steps/checking_steps.py @@ -131,6 +131,13 @@ def verify_minifi_logs_do_not_contain_warnings(context: MinifiTestContext): ) +@then("the logs of the exited Minifi container contain the following message: '{message}'") +@then('the logs of the exited Minifi container contain the following message: "{message}"') +def verify_exited_minifi_container_logs_contain_message(context: MinifiTestContext, message: str): + assert context.containers[DEFAULT_MINIFI_CONTAINER_NAME].exited or log_due_to_failure(context) + assert message in context.containers[DEFAULT_MINIFI_CONTAINER_NAME].get_logs() or log_due_to_failure(context) + + @then("the logs of the '{container}' container contain the following message: '{message}' in less than {duration}") @then('the logs of the "{container}" container contain the following message: "{message}" in less than {duration}') def verify_container_logs_contain_message(context: MinifiTestContext, container: str, message: str, duration: str): diff --git a/bin/minifi.service.in b/bin/minifi.service.in index 2c7f84068f..0b8932cccf 100644 --- a/bin/minifi.service.in +++ b/bin/minifi.service.in @@ -12,6 +12,7 @@ RestartSec=3 KillSignal=SIGTERM TimeoutStopSec=20 RestartForceExitStatus=3 +RestartPreventExitStatus=90 [Install] WantedBy=multi-user.target diff --git a/conf/minifi.properties.in b/conf/minifi.properties.in index 722fc1ca71..8b55d19dbc 100644 --- a/conf/minifi.properties.in +++ b/conf/minifi.properties.in @@ -121,6 +121,8 @@ nifi.c2.full.heartbeat=false # nifi.c2.root.class.definitions.metrics.metrics.processorMetrics.name=ProcessorMetric # nifi.c2.root.class.definitions.metrics.metrics.processorMetrics.classes=GetFileMetrics +# nifi.c2.connection.timeout=30 min + ## Specify multiple heartbeat handlers. ## HeartbeatLogger logs the heartbeats on TRACE for debugging. #nifi.c2.agent.heartbeat.reporter.classes=HeartbeatLogger diff --git a/extensions/standard-processors/tests/features/minifi_c2_server.feature b/extensions/standard-processors/tests/features/minifi_c2_server.feature index c35ddeed95..9daba8f7df 100644 --- a/extensions/standard-processors/tests/features/minifi_c2_server.feature +++ b/extensions/standard-processors/tests/features/minifi_c2_server.feature @@ -55,3 +55,23 @@ Feature: MiNiFi can communicate with Apache NiFi MiNiFi C2 server Then the MiNiFi C2 server logs contain the following message: "acknowledged with a state of FULLY_APPLIED(DONE)" in less than 30 seconds And a single file with the content "test" is placed in the "/tmp/output" directory in less than 10 seconds And the Minifi logs do not contain the following message: "Failed to parse json response: The document is empty. at 0" after 0 seconds + + Scenario: MiNiFi shuts down when the connection to the MiNiFi C2 server is lost + Given a GenerateFlowFile processor + And C2 is enabled in MiNiFi + And MiNiFi configuration "nifi.c2.connection.timeout" is set to "10 sec" + And a MiNiFi C2 server is started + When the MiNiFi instance starts up + Then the MiNiFi C2 server logs contain the following message: "acknowledged with a state of FULLY_APPLIED(DONE)" in less than 30 seconds + And the Minifi logs do not contain the following message: "C2 connection timeout exceeded, shutting down..." after 15 seconds + When the MiNiFi C2 server is stopped + Then MiNiFi exits with exit code 90 in less than 60 seconds + And the logs of the exited Minifi container contain the following message: "C2 connection timeout exceeded, shutting down..." + + Scenario: MiNiFi shuts down when the connection to the MiNiFi C2 server is not available + Given a GenerateFlowFile processor + And C2 is enabled in MiNiFi + And MiNiFi configuration "nifi.c2.connection.timeout" is set to "10 sec" + When the MiNiFi instance starts up + Then MiNiFi exits with exit code 90 in less than 60 seconds + And the logs of the exited Minifi container contain the following message: "C2 connection timeout exceeded, shutting down..." diff --git a/extensions/standard-processors/tests/features/steps/steps.py b/extensions/standard-processors/tests/features/steps/steps.py index b4469fbe9a..df4b7940b5 100644 --- a/extensions/standard-processors/tests/features/steps/steps.py +++ b/extensions/standard-processors/tests/features/steps/steps.py @@ -16,12 +16,15 @@ # import humanfriendly -from behave import given, step, then +from behave import given, step, then, when from containers.diag_slave_container import DiagSlave from containers.syslog_container import SyslogContainer from containers.tcp_client_container import TcpClientContainer from minifi_behave.containers.minifi_protocol import set_up_ssl_properties -from minifi_behave.core.helpers import wait_for_condition +from minifi_behave.core.helpers import ( + log_due_to_failure, + wait_for_condition, +) from minifi_behave.core.minifi_test_context import MinifiTestContext from minifi_behave.steps import ( checking_steps, # noqa: F401 @@ -126,6 +129,24 @@ def start_minifi_c2_server(context: MinifiTestContext): assert context.containers["minifi-c2-server"].deploy(context) +@when("the MiNiFi C2 server is stopped") +def stop_minifi_c2_server(context: MinifiTestContext): + context.containers["minifi-c2-server"].stop() + + +@then("MiNiFi exits with exit code {exit_code:d} in less than {duration}") +def verify_minifi_exit_code(context: MinifiTestContext, exit_code: int, duration: str): + duration_seconds = humanfriendly.parse_timespan(duration) + minifi_container = context.get_or_create_default_minifi_container() + assert wait_for_condition( + condition=lambda: minifi_container.exited, + timeout_seconds=duration_seconds, + bail_condition=lambda: False, + context=context, + ) + assert minifi_container.exit_code == exit_code or log_due_to_failure(context) + + @then('the MiNiFi C2 server logs contain the following message: "{log_message}" in less than {duration}') def verify_c2_server_logs_contain_message(context: MinifiTestContext, log_message: str, duration: str): duration_seconds = humanfriendly.parse_timespan(duration) diff --git a/libminifi/include/FlowController.h b/libminifi/include/FlowController.h index 00ff3f756d..5177e29579 100644 --- a/libminifi/include/FlowController.h +++ b/libminifi/include/FlowController.h @@ -46,6 +46,7 @@ #include "minifi-cpp/core/state/nodes/MetricsBase.h" #include "core/state/UpdateController.h" #include "c2/C2Agent.h" +#include "RequestShutdown.h" #include "CronDrivenSchedulingAgent.h" #include "EventDrivenSchedulingAgent.h" #include "properties/Configure.h" @@ -69,7 +70,8 @@ class FlowController : public core::controller::ForwardingControllerServiceProvi FlowController(std::shared_ptr provenance_repo, std::shared_ptr flow_file_repo, std::shared_ptr configure, std::shared_ptr flow_configuration, std::shared_ptr content_repo, std::unique_ptr metrics_publisher_store = nullptr, - std::shared_ptr filesystem = std::make_shared(), std::function request_restart = []{}, + std::shared_ptr filesystem = std::make_shared(), + RequestShutdownCallback request_shutdown = [](ShutdownType, std::optional){}, utils::file::AssetManager* asset_manager = {}, core::BulletinStore* bulletin_store = {}); ~FlowController() override; diff --git a/libminifi/include/RequestShutdown.h b/libminifi/include/RequestShutdown.h new file mode 100644 index 0000000000..7451fb6a0a --- /dev/null +++ b/libminifi/include/RequestShutdown.h @@ -0,0 +1,32 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include +#include +#include + +namespace org::apache::nifi::minifi { + +enum class ShutdownType { + RESTART, + SHUTDOWN +}; + +using RequestShutdownCallback = std::function exit_code)>; + +} // namespace org::apache::nifi::minifi diff --git a/libminifi/include/c2/C2Agent.h b/libminifi/include/c2/C2Agent.h index 691d4be577..1267d68cd4 100644 --- a/libminifi/include/c2/C2Agent.h +++ b/libminifi/include/c2/C2Agent.h @@ -18,6 +18,7 @@ #pragma once +#include #include #include #include @@ -43,6 +44,7 @@ #include "utils/file/FileSystem.h" #include "C2Utils.h" #include "utils/file/AssetManager.h" +#include "RequestShutdown.h" namespace org::apache::nifi::minifi::c2 { @@ -62,7 +64,7 @@ class C2Agent : public state::UpdateController { C2Agent(std::shared_ptr configuration, std::weak_ptr node_reporter, std::shared_ptr filesystem, - std::function request_restart, + minifi::RequestShutdownCallback request_shutdown, utils::file::AssetManager* asset_manager); void initialize(core::controller::ControllerServiceProvider *controller, state::Pausable *pause_handler, state::StateMonitor* update_sink); @@ -235,12 +237,14 @@ class C2Agent : public state::UpdateController { const uint64_t C2RESPONSE_POLL_MS = 100; std::atomic restart_needed_ = false; - std::function request_restart_; + minifi::RequestShutdownCallback request_shutdown_; // time point the last time we performed a heartbeat. std::chrono::steady_clock::time_point last_run_; utils::file::AssetManager* asset_manager_; + std::optional c2_connection_timeout_; + std::chrono::steady_clock::time_point last_successful_heartbeat_{}; }; } // namespace org::apache::nifi::minifi::c2 diff --git a/libminifi/include/c2/C2ExitCodes.h b/libminifi/include/c2/C2ExitCodes.h new file mode 100644 index 0000000000..0a0c709176 --- /dev/null +++ b/libminifi/include/c2/C2ExitCodes.h @@ -0,0 +1,25 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include + +namespace org::apache::nifi::minifi::c2 { + +constexpr int32_t C2_CONNECTION_TIMEOUT_EXIT_CODE = 90; + +} // namespace org::apache::nifi::minifi::c2 diff --git a/libminifi/src/Configuration.cpp b/libminifi/src/Configuration.cpp index 73c9519c2e..f236ff814f 100644 --- a/libminifi/src/Configuration.cpp +++ b/libminifi/src/Configuration.cpp @@ -97,6 +97,8 @@ const std::unordered_map provenance_repo, std::shared_ptr flow_file_repo, std::shared_ptr configure, std::shared_ptr flow_configuration, std::shared_ptr content_repo, std::unique_ptr metrics_publisher_store, - std::shared_ptr filesystem, std::function request_restart, + std::shared_ptr filesystem, RequestShutdownCallback request_shutdown, utils::file::AssetManager* asset_manager, core::BulletinStore* bulletin_store) : core::controller::ForwardingControllerServiceProvider(core::className()), running_(false), @@ -82,7 +82,7 @@ FlowController::FlowController(std::shared_ptr provenance_repo if (auto publisher = metrics_publisher_store_->getMetricsPublisher(c2::C2_METRICS_PUBLISHER).lock()) { c2_metrics_publisher = std::dynamic_pointer_cast(publisher); } - c2_agent_ = std::make_unique(configuration_, c2_metrics_publisher, std::move(filesystem), std::move(request_restart), asset_manager); + c2_agent_ = std::make_unique(configuration_, c2_metrics_publisher, std::move(filesystem), std::move(request_shutdown), asset_manager); } if (c2::isControllerSocketEnabled(configuration_)) { diff --git a/libminifi/src/c2/C2Agent.cpp b/libminifi/src/c2/C2Agent.cpp index 67ec8b72a6..d4b44a12ba 100644 --- a/libminifi/src/c2/C2Agent.cpp +++ b/libminifi/src/c2/C2Agent.cpp @@ -43,6 +43,7 @@ #include "minifi-cpp/io/ArchiveStream.h" #include "io/StreamPipe.h" #include "utils/Id.h" +#include "c2/C2ExitCodes.h" #include "c2/C2Utils.h" #include "c2/protocols/RESTSender.h" #include "rapidjson/error/en.h" @@ -52,10 +53,29 @@ using namespace std::literals::chrono_literals; namespace org::apache::nifi::minifi::c2 { +namespace { +std::optional parseTimeoutPeriod(const std::shared_ptr& configuration, const std::string& config_param, + const std::shared_ptr& logger) { + return configuration->get(config_param) | utils::andThen([&](const auto& timeout_str) -> std::optional { + auto duration = parsing::parseDuration(timeout_str) | utils::toOptional(); + if (!duration) { + logger->log_error("Failed to parse C2 configuration parameter '{}': '{}'", config_param, timeout_str); + } + if (duration && *duration <= 0ms) { + if (*duration < 0ms) { + logger->log_warn("Negative value for C2 configuration parameter '{}': '{}', ignoring it", config_param, timeout_str); + } + return std::nullopt; + } + return duration; + }); +} +} // namespace + C2Agent::C2Agent(std::shared_ptr configuration, std::weak_ptr node_reporter, std::shared_ptr filesystem, - std::function request_restart, + minifi::RequestShutdownCallback request_shutdown, utils::file::AssetManager* asset_manager) : heart_beat_period_(3s), max_c2_responses(5), @@ -63,7 +83,7 @@ C2Agent::C2Agent(std::shared_ptr configuration, node_reporter_(std::move(node_reporter)), filesystem_(std::move(filesystem)), thread_pool_(2, "C2 threadpool"), - request_restart_(std::move(request_restart)), + request_shutdown_(std::move(request_shutdown)), last_run_(std::chrono::steady_clock::now()), asset_manager_(asset_manager) { if (!configuration_->getAgentClass()) { @@ -109,6 +129,7 @@ void C2Agent::start() { if (controller_running_) { return; } + last_successful_heartbeat_ = std::chrono::steady_clock::now(); task_ids_.clear(); for (const auto& function : functions_) { utils::Identifier uuid = utils::IdGenerator::getIdGenerator()->generate(); @@ -179,6 +200,17 @@ void C2Agent::configure(const std::shared_ptr &configure, bool reconf } logger_->log_debug("Using {} as the heartbeat period", heart_beat_period_); + c2_connection_timeout_ = parseTimeoutPeriod(configuration_, std::string{Configuration::nifi_c2_connection_timeout}, logger_); + if (c2_connection_timeout_ && *c2_connection_timeout_ < 2 * heart_beat_period_) { + logger_->log_warn( + "C2 connection timeout is less than twice the heartbeat period: C2 connection timeout={}, heartbeat period={}. Adjusting it to be twice the " + "heartbeat period", + *c2_connection_timeout_, + heart_beat_period_); + c2_connection_timeout_ = heart_beat_period_ * 2; + logger_->log_debug("Using {} as the C2 connection timeout", *c2_connection_timeout_); + } + std::string heartbeat_reporters; if (configure->get(Configuration::nifi_c2_agent_heartbeat_reporter_classes, "c2.agent.heartbeat.reporter.classes", heartbeat_reporters)) { std::vector reporters = utils::string::splitAndTrim(heartbeat_reporters, ","); @@ -234,6 +266,10 @@ void C2Agent::performHeartBeat() { } C2Payload response = protocol_->consumePayload(payload); + if (response.getStatus().getState() != state::UpdateState::READ_ERROR) { + last_successful_heartbeat_ = std::chrono::steady_clock::now(); + } + enqueue_c2_server_response(std::move(response)); std::lock_guard lock(heartbeat_mutex); @@ -944,7 +980,7 @@ utils::TaskRescheduleInfo C2Agent::produce() { if (restart_needed_ && requests.empty()) { configuration_->commitChanges(); - request_restart_(); + request_shutdown_(minifi::ShutdownType::RESTART, std::nullopt); return utils::TaskRescheduleInfo::Done(); } @@ -957,6 +993,12 @@ utils::TaskRescheduleInfo C2Agent::produce() { catch (...) { logger_->log_error("Unknown exception occurred while performing heartbeat, type: {}", getCurrentExceptionTypeName()); } + + if (c2_connection_timeout_ && controller_running_ && std::chrono::steady_clock::now() - last_successful_heartbeat_ > *c2_connection_timeout_) { + logger_->log_error("C2 connection timeout exceeded, shutting down..."); + request_shutdown_(minifi::ShutdownType::SHUTDOWN, C2_CONNECTION_TIMEOUT_EXIT_CODE); + return utils::TaskRescheduleInfo::Done(); + } } checkTriggers(); diff --git a/libminifi/test/integration/C2ConnectionTimeoutTest.cpp b/libminifi/test/integration/C2ConnectionTimeoutTest.cpp new file mode 100644 index 0000000000..2ad6447dc9 --- /dev/null +++ b/libminifi/test/integration/C2ConnectionTimeoutTest.cpp @@ -0,0 +1,93 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include + +#include "c2/C2Agent.h" +#include "c2/C2ExitCodes.h" +#include "integration/HTTPIntegrationBase.h" +#include "integration/ServerAwareHandler.h" +#include "properties/Configuration.h" +#include "unit/Catch.h" +#include "unit/TestBase.h" +#include "unit/TestUtils.h" + +using namespace std::literals::chrono_literals; + +namespace org::apache::nifi::minifi::test { + +class FailingHeartbeatHandler : public ServerAwareHandler { + public: + bool handlePost(CivetServer* /*server*/, struct mg_connection* conn) override { + ++heartbeat_count_; + if (fail_) { + mg_printf(conn, "HTTP/1.1 500 Internal Server Error\r\nContent-Length: 0\r\nConnection: close\r\n\r\n"); + } else { + static const std::string response = R"({"operation": "heartbeat", "requested_operations": []})"; + mg_printf(conn, "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: %lu\r\nConnection: close\r\n\r\n", response.length()); + mg_printf(conn, "%s", response.c_str()); + } + return true; + } + + std::atomic fail_{false}; + std::atomic heartbeat_count_{0}; +}; + +class VerifyC2ConnectionTimeout : public VerifyC2Base { + public: + VerifyC2ConnectionTimeout(const std::filesystem::path& test_file_path, FailingHeartbeatHandler& handler) + : VerifyC2Base(test_file_path), + handler_(handler) { + } + + void testSetup() override { + LogTestController::getInstance().setDebug(); + VerifyC2Base::testSetup(); + } + + void configureC2() override { + VerifyC2Base::configureC2(); + configuration->set(minifi::Configuration::nifi_c2_agent_heartbeat_period, "100 ms"); + configuration->set(minifi::Configuration::nifi_c2_connection_timeout, "1 sec"); + } + + void runAssertions() override { + // successful heartbeats keep the agent alive well past the connection timeout + std::this_thread::sleep_for(3s); + REQUIRE(handler_.heartbeat_count_ > 0); + REQUIRE(shutdown_requested_count_ == 0); + + handler_.fail_ = true; + REQUIRE(utils::verifyEventHappenedInPollTime(5s, [&] { return shutdown_requested_count_ > 0; })); + CHECK(shutdown_requested_count_ == 1); + CHECK(last_shutdown_exit_code_ == minifi::c2::C2_CONNECTION_TIMEOUT_EXIT_CODE); + CHECK(utils::verifyLogLinePresenceInPollTime(1s, "C2 connection timeout exceeded, shutting down...")); + } + + private: + FailingHeartbeatHandler& handler_; +}; + +TEST_CASE("C2 agent requests shutdown when the C2 server stops responding successfully", "[c2test]") { + FailingHeartbeatHandler handler; + VerifyC2ConnectionTimeout harness(std::filesystem::path(TEST_RESOURCES) / "TestEmpty.yml", handler); + harness.setUrl("http://localhost:0/heartbeat", &handler); + harness.run(); +} + +} // namespace org::apache::nifi::minifi::test diff --git a/libminifi/test/libtest/integration/IntegrationBase.cpp b/libminifi/test/libtest/integration/IntegrationBase.cpp index 8a767546a6..7d8d399873 100644 --- a/libminifi/test/libtest/integration/IntegrationBase.cpp +++ b/libminifi/test/libtest/integration/IntegrationBase.cpp @@ -129,7 +129,13 @@ void IntegrationBase::run() { std::shared_ptr pg(flow_config->getRoot()); queryRootProcessGroup(pg); - const auto request_restart = [&, this] { + const auto request_shutdown = [&, this] (ShutdownType shutdown_type, std::optional exit_code) { + if (shutdown_type == ShutdownType::SHUTDOWN) { + last_shutdown_exit_code_ = exit_code.value_or(0); + ++shutdown_requested_count_; + running = false; + return; + } ++restart_requested_count_; running = true; }; @@ -138,7 +144,7 @@ void IntegrationBase::run() { asset_manager_ = std::make_unique(*configuration); auto metrics_publisher_store = std::make_unique(configuration, repo_metric_sources, flow_config, asset_manager_.get(), bulletin_store_.get()); flowController_ = std::make_unique(test_repo, test_flow_repo, configuration, - std::move(flow_config), content_repo, std::move(metrics_publisher_store), filesystem, request_restart, asset_manager_.get()); + std::move(flow_config), content_repo, std::move(metrics_publisher_store), filesystem, request_shutdown, asset_manager_.get()); flowController_->load(); updateProperties(*flowController_); flowController_->start(); diff --git a/libminifi/test/libtest/integration/IntegrationBase.h b/libminifi/test/libtest/integration/IntegrationBase.h index 9f406d57ad..16c5286e9e 100644 --- a/libminifi/test/libtest/integration/IntegrationBase.h +++ b/libminifi/test/libtest/integration/IntegrationBase.h @@ -59,6 +59,8 @@ class IntegrationBase { key_dir{std::move(other.key_dir)}, state_dir{std::move(other.state_dir)}, restart_requested_count_{other.restart_requested_count_.load()}, + shutdown_requested_count_{other.shutdown_requested_count_.load()}, + last_shutdown_exit_code_{other.last_shutdown_exit_code_.load()}, home_path_{std::move(other.home_path_)}, flow_config_path_{std::move(other.flow_config_path_)} {} @@ -74,6 +76,8 @@ class IntegrationBase { key_dir = std::move(other.key_dir); state_dir = std::move(other.state_dir); restart_requested_count_ = other.restart_requested_count_.load(); + shutdown_requested_count_ = other.shutdown_requested_count_.load(); + last_shutdown_exit_code_ = other.last_shutdown_exit_code_.load(); home_path_ = std::move(other.home_path_); flow_config_path_ = std::move(other.flow_config_path_); return *this; @@ -141,6 +145,8 @@ class IntegrationBase { std::filesystem::path key_dir; std::filesystem::path state_dir; std::atomic restart_requested_count_{0}; + std::atomic shutdown_requested_count_{0}; + std::atomic last_shutdown_exit_code_{0}; std::optional home_path_; FlowConfigPath flow_config_path_; }; diff --git a/libminifi/test/unit/C2AgentConnectionTimeoutTests.cpp b/libminifi/test/unit/C2AgentConnectionTimeoutTests.cpp new file mode 100644 index 0000000000..9ef49654e2 --- /dev/null +++ b/libminifi/test/unit/C2AgentConnectionTimeoutTests.cpp @@ -0,0 +1,127 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "c2/C2Agent.h" +#include "c2/C2ExitCodes.h" +#include "properties/Configure.h" +#include "unit/Catch.h" +#include "unit/TestBase.h" +#include "unit/TestUtils.h" +#include "utils/file/FileSystem.h" + +using namespace std::literals::chrono_literals; + +namespace org::apache::nifi::minifi::test { + +class C2AgentConnectionTimeoutTestFixture { + public: + C2AgentConnectionTimeoutTestFixture() { + LogTestController::getInstance().reset(); + LogTestController::getInstance().setTrace(); + configuration_->set(Configuration::nifi_c2_agent_heartbeat_period, "100 ms"); + } + + C2AgentConnectionTimeoutTestFixture(C2AgentConnectionTimeoutTestFixture&&) = delete; + C2AgentConnectionTimeoutTestFixture(const C2AgentConnectionTimeoutTestFixture&) = delete; + C2AgentConnectionTimeoutTestFixture& operator=(C2AgentConnectionTimeoutTestFixture&&) = delete; + C2AgentConnectionTimeoutTestFixture& operator=(const C2AgentConnectionTimeoutTestFixture&) = delete; + + ~C2AgentConnectionTimeoutTestFixture() { + if (agent_) { + agent_->stop(); + } + LogTestController::getInstance().reset(); + } + + void startAgent() { + agent_ = std::make_unique(configuration_, + std::weak_ptr{}, + std::make_shared(), + [this](ShutdownType shutdown_type, std::optional exit_code) { + last_shutdown_type_ = shutdown_type; + last_exit_code_ = exit_code.value_or(0); + ++shutdown_request_count_; + }, + nullptr); + agent_->initialize(nullptr, nullptr, nullptr); + agent_->start(); + } + + protected: + std::shared_ptr configuration_ = std::make_shared(); + std::atomic shutdown_request_count_{0}; + std::atomic last_shutdown_type_{ShutdownType::RESTART}; + std::atomic last_exit_code_{0}; + std::unique_ptr agent_; +}; + +TEST_CASE_METHOD(C2AgentConnectionTimeoutTestFixture, "No shutdown is requested when the C2 connection timeout is not set", "[c2connectiontimeout]") { + startAgent(); + REQUIRE(LogTestController::getInstance().contains("Performing heartbeat", 1s)); + std::this_thread::sleep_for(1s); + REQUIRE(shutdown_request_count_ == 0); +} + +TEST_CASE_METHOD(C2AgentConnectionTimeoutTestFixture, "C2 connection timeout of 0 disables the C2 connection timeout", "[c2connectiontimeout]") { + configuration_->set(Configuration::nifi_c2_connection_timeout, "0 sec"); + startAgent(); + REQUIRE(LogTestController::getInstance().contains("Performing heartbeat", 1s)); + std::this_thread::sleep_for(1s); + REQUIRE(shutdown_request_count_ == 0); + REQUIRE_FALSE(LogTestController::getInstance().contains("Failed to parse C2 configuration parameter", 0ms)); +} + +TEST_CASE_METHOD(C2AgentConnectionTimeoutTestFixture, "Invalid C2 connection timeout disables the C2 connection timeout", "[c2connectiontimeout]") { + configuration_->set(Configuration::nifi_c2_connection_timeout, "abc"); + startAgent(); + REQUIRE(LogTestController::getInstance().contains("Failed to parse C2 configuration parameter 'nifi.c2.connection.timeout': 'abc'", 1s)); + std::this_thread::sleep_for(1s); + REQUIRE(shutdown_request_count_ == 0); +} + +TEST_CASE_METHOD(C2AgentConnectionTimeoutTestFixture, "C2 connection timeout shorter than twice the heartbeat period is increased", + "[c2connectiontimeout]") { + configuration_->set(Configuration::nifi_c2_connection_timeout, "150 ms"); + startAgent(); + REQUIRE(LogTestController::getInstance().contains("C2 connection timeout is less than twice the heartbeat period", 1s)); + REQUIRE(LogTestController::getInstance().contains("Using 200ms as the C2 connection timeout", 1s)); +} + +TEST_CASE_METHOD(C2AgentConnectionTimeoutTestFixture, "Failed heartbeats do not prevent the C2 connection timeout", "[c2connectiontimeout]") { + configuration_->set(Configuration::nifi_c2_connection_timeout, "300 ms"); + startAgent(); + REQUIRE(utils::verifyEventHappenedInPollTime(3s, [&] { return shutdown_request_count_ > 0; }, 10ms)); + REQUIRE(LogTestController::getInstance().contains("Performing heartbeat", 0ms)); + std::this_thread::sleep_for(500ms); + REQUIRE(shutdown_request_count_ == 1); + CHECK(last_shutdown_type_ == ShutdownType::SHUTDOWN); + CHECK(last_exit_code_ == c2::C2_CONNECTION_TIMEOUT_EXIT_CODE); +} + +TEST_CASE_METHOD(C2AgentConnectionTimeoutTestFixture, "Stopping the C2 agent stops the C2 connection timeout check", "[c2connectiontimeout]") { + configuration_->set(Configuration::nifi_c2_connection_timeout, "300 ms"); + startAgent(); + agent_->stop(); + std::this_thread::sleep_for(600ms); + REQUIRE(shutdown_request_count_ == 0); +} + +} // namespace org::apache::nifi::minifi::test diff --git a/minifi-api/include/minifi-cpp/properties/Configuration.h b/minifi-api/include/minifi-cpp/properties/Configuration.h index 1ec60afa51..9272c2152a 100644 --- a/minifi-api/include/minifi-cpp/properties/Configuration.h +++ b/minifi-api/include/minifi-cpp/properties/Configuration.h @@ -127,6 +127,7 @@ class Configuration : public virtual Properties { static constexpr const char *nifi_c2_rest_request_encoding = "nifi.c2.rest.request.encoding"; static constexpr const char *nifi_c2_flow_info_processor_bulletin_limit = "nifi.c2.flow.info.processor.bulletin.limit"; static constexpr const char *nifi_c2_asset_download_timeout = "nifi.c2.asset.download.timeout"; + static constexpr const char *nifi_c2_connection_timeout = "nifi.c2.connection.timeout"; // state management options static constexpr const char *nifi_state_storage_local = "nifi.state.storage.local"; diff --git a/minifi_main/MiNiFiMain.cpp b/minifi_main/MiNiFiMain.cpp index fcc2ba9d6a..1516a22f84 100644 --- a/minifi_main/MiNiFiMain.cpp +++ b/minifi_main/MiNiFiMain.cpp @@ -43,6 +43,7 @@ #include #include #include +#include #include #include "AgentDocs.h" @@ -51,6 +52,7 @@ #include "Fips.h" #include "FlowController.h" #include "MainHelper.h" +#include "RequestShutdown.h" #include "ResourceClaim.h" #include "agent/JsonSchema.h" #include "minifi-cpp/agent/agent_version.h" @@ -79,6 +81,45 @@ namespace utils = minifi::utils; static std::atomic_flag flow_controller_running; static std::atomic_flag process_running; +enum class LifecycleState : int32_t { + RUNNING, + RESTART_REQUESTED, + SHUTDOWN_REQUESTED +}; + +struct LifecycleRequest { + LifecycleState state = LifecycleState::RUNNING; + int32_t exit_code = 0; +}; + +// Ensure that LifecycleRequest can be safely used with atomic operations in signal handlers +// Extending LifecycleRequest in the future may break the lock-free guarantee +static_assert(std::atomic::is_always_lock_free, + "lifecycle_request is modified from signal handlers (see sigHandler), where only lock-free atomic operations are safe. " + "A non-lock-free std::atomic falls back to hidden mutex or a lock, which can deadlock if a signal interrupts a thread holding it."); +static std::atomic lifecycle_request; + +bool requestLifecycleChange(minifi::ShutdownType shutdown_type, int32_t exit_code) { + auto current_request = lifecycle_request.load(); + while (true) { + if (current_request.state == LifecycleState::SHUTDOWN_REQUESTED) { + return false; + } + if (shutdown_type == minifi::ShutdownType::RESTART && current_request.state == LifecycleState::RESTART_REQUESTED) { + return false; + } + const auto desired = shutdown_type == minifi::ShutdownType::SHUTDOWN + ? LifecycleRequest{.state = LifecycleState::SHUTDOWN_REQUESTED, .exit_code = exit_code} + : LifecycleRequest{.state = LifecycleState::RESTART_REQUESTED, .exit_code = 0}; + if (lifecycle_request.compare_exchange_weak(current_request, desired)) { + break; + } + } + flow_controller_running.clear(); + flow_controller_running.notify_all(); + return true; +} + /** * Removed the stop command from the signal handler so that we could trigger * unload after we exit the semaphore controlled critical section in main. @@ -93,8 +134,7 @@ static std::atomic_flag process_running; BOOL WINAPI consoleSignalHandler(DWORD signal) { if (!process_running.test()) { exit(0); return TRUE; } if (signal == CTRL_C_EVENT || signal == CTRL_BREAK_EVENT) { - flow_controller_running.clear(); - flow_controller_running.notify_all(); + requestLifecycleChange(minifi::ShutdownType::SHUTDOWN, 0); process_running.wait(true); return TRUE; } @@ -102,15 +142,13 @@ BOOL WINAPI consoleSignalHandler(DWORD signal) { } void SignalExitProcess() { - flow_controller_running.clear(); - flow_controller_running.notify_all(); + requestLifecycleChange(minifi::ShutdownType::SHUTDOWN, 0); } #endif void sigHandler(int signal) { if (signal == SIGINT || signal == SIGTERM) { - flow_controller_running.clear(); - flow_controller_running.notify_all(); + requestLifecycleChange(minifi::ShutdownType::SHUTDOWN, 0); } } @@ -175,6 +213,8 @@ void overridePropertiesFromCommandLine(const argparse::ArgumentParser& parser, c } int main(int argc, char **argv) { + flow_controller_running.test_and_set(); + argparse::ArgumentParser argument_parser("Apache MiNiFi C++", minifi::AgentBuild::VERSION); argument_parser.add_argument("-p", "--property") .append() @@ -289,19 +329,18 @@ int main(int argc, char **argv) { process_running.test_and_set(); - std::atomic restart_token{false}; - const auto request_restart = [&] { - if (!restart_token.exchange(true)) { - // only trigger if a restart is not already in progress (the flag was unset before the exchange) - flow_controller_running.clear(); - flow_controller_running.notify_all(); + const auto request_shutdown = [&] (minifi::ShutdownType shutdown_type, std::optional shutdown_exit_code) { + if (!requestLifecycleChange(shutdown_type, shutdown_exit_code.value_or(0))) { + return; + } + if (shutdown_type == minifi::ShutdownType::SHUTDOWN) { + logger->log_info("Initiating shutdown..."); + } else { logger->log_info("Initiating restart..."); } }; do { - flow_controller_running.test_and_set(); - std::string prov_repo_class = "provenancerepository"; std::string flow_repo_class = "flowfilerepository"; std::string nifi_configuration_class_name = "adaptiveconfiguration"; @@ -408,7 +447,7 @@ int main(int argc, char **argv) { auto metrics_publisher_store = std::make_unique(configure, repo_metric_sources, flow_configuration, asset_manager.get(), bulletin_store.get()); const auto controller = std::make_unique( prov_repo, flow_repo, configure, std::move(flow_configuration), content_repo, - std::move(metrics_publisher_store), filesystem, request_restart, asset_manager.get(), bulletin_store.get()); + std::move(metrics_publisher_store), filesystem, request_shutdown, asset_manager.get(), bulletin_store.get()); const bool disk_space_watchdog_enable = configure->get(minifi::Configure::minifi_disk_space_watchdog_enable) | utils::andThen(utils::string::toBool) @@ -495,15 +534,17 @@ int main(int argc, char **argv) { flow_repo = nullptr; prov_repo = nullptr; } while ([&] { - const auto restart_token_temp = restart_token.exchange(false); - if (restart_token_temp) { - logger->log_info("Restarting MiNiFi"); + flow_controller_running.test_and_set(); + auto expected = LifecycleRequest{.state = LifecycleState::RESTART_REQUESTED, .exit_code = 0}; + if (!lifecycle_request.compare_exchange_strong(expected, LifecycleRequest{})) { + return false; } - return restart_token_temp; + logger->log_info("Restarting MiNiFi"); + return true; }()); process_running.clear(); process_running.notify_all(); logger->log_info("MiNiFi exit"); - return 0; + return lifecycle_request.load().exit_code; } diff --git a/minifi_main/MiNiFiWindowsService.cpp b/minifi_main/MiNiFiWindowsService.cpp index dd5e446f3f..e3be008946 100644 --- a/minifi_main/MiNiFiWindowsService.cpp +++ b/minifi_main/MiNiFiWindowsService.cpp @@ -26,6 +26,7 @@ #include #include "MainHelper.h" +#include "c2/C2ExitCodes.h" namespace minifi = org::apache::nifi::minifi; @@ -176,6 +177,8 @@ void RunAsServiceIfNeeded() { return; } + DWORD serviceSpecificExitCode = 0; + do { Log()->log_debug("Start exe path {}", filePath); @@ -194,6 +197,11 @@ void RunAsServiceIfNeeded() { Log()->log_info("{} started", filePath); auto res = WaitForSingleObject(processInformation.hProcess, INFINITE); + DWORD exeExitCode = 0; + const bool exeExitCodeRetrieved = GetExitCodeProcess(processInformation.hProcess, &exeExitCode); + if (!exeExitCodeRetrieved) { + Log()->log_error("!GetExitCodeProcess lastError {:#x}", GetLastError()); + } CloseHandle(processInformation.hProcess); CloseHandle(processInformation.hThread); @@ -203,6 +211,12 @@ void RunAsServiceIfNeeded() { Log()->log_error("!WaitForSingleObject hProcess return {}", res); } + if (exeExitCodeRetrieved && exeExitCode == static_cast(minifi::c2::C2_CONNECTION_TIMEOUT_EXIT_CODE)) { + Log()->log_info("Exe exited with exit code {} due to C2 connection timeout, exe won't be restarted", exeExitCode); + serviceSpecificExitCode = exeExitCode; + break; + } + Log()->log_info("Sleep {} sec before restarting exe", WAIT_TIME_EXE_RESTART/1000); res = WaitForSingleObject(s_hEvent, WAIT_TIME_EXE_RESTART); @@ -220,7 +234,8 @@ void RunAsServiceIfNeeded() { s_serviceStatus.dwControlsAccepted = 0; s_serviceStatus.dwCurrentState = SERVICE_STOPPED; - s_serviceStatus.dwWin32ExitCode = 0; + s_serviceStatus.dwWin32ExitCode = serviceSpecificExitCode == 0 ? 0 : ERROR_SERVICE_SPECIFIC_ERROR; + s_serviceStatus.dwServiceSpecificExitCode = serviceSpecificExitCode; if (!SetServiceStatus(s_statusHandle, &s_serviceStatus)) { Log()->log_error("!SetServiceStatus SERVICE_STOPPED lastError {:#x}", GetLastError());