MINIFICPP-2933 Stop MiNiFi C++ if C2 connection is lost for some time

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

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")

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: ...

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

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):

bin/minifi.service.in

@@ -12,6 +12,7 @@ RestartSec=3
 KillSignal=SIGTERM
 TimeoutStopSec=20
 RestartForceExitStatus=3
+RestartPreventExitStatus=90
 
 [Install]
 WantedBy=multi-user.target

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

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..."

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)

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<core::Repository> provenance_repo, std::shared_ptr<core::Repository> flow_file_repo,
                  std::shared_ptr<Configure> configure, std::shared_ptr<core::FlowConfiguration> flow_configuration,
                  std::shared_ptr<core::ContentRepository> content_repo, std::unique_ptr<state::MetricsPublisherStore> metrics_publisher_store = nullptr,
-                 std::shared_ptr<utils::file::FileSystem> filesystem = std::make_shared<utils::file::FileSystem>(), std::function<void()> request_restart = []{},
+                 std::shared_ptr<utils::file::FileSystem> filesystem = std::make_shared<utils::file::FileSystem>(),
+                 RequestShutdownCallback request_shutdown = [](ShutdownType, std::optional<int32_t>){},
                  utils::file::AssetManager* asset_manager = {}, core::BulletinStore* bulletin_store = {});
 
   ~FlowController() override;

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 <cstdint>
+#include <functional>
+#include <optional>
+
+namespace org::apache::nifi::minifi {
+
+enum class ShutdownType {
+  RESTART,
+  SHUTDOWN
+};
+
+using RequestShutdownCallback = std::function<void(ShutdownType shutdown_type, std::optional<int32_t> exit_code)>;
+
+}  // namespace org::apache::nifi::minifi

libminifi/include/c2/C2Agent.h

@@ -18,6 +18,7 @@
 
 #pragma once
 
+#include <chrono>
 #include <functional>
 #include <future>
 #include <map>
@@ -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<Configure> configuration,
           std::weak_ptr<state::response::NodeReporter> node_reporter,
           std::shared_ptr<utils::file::FileSystem> filesystem,
-          std::function<void()> 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<bool> restart_needed_ = false;
-  std::function<void()> 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<std::chrono::milliseconds> c2_connection_timeout_;
+  std::chrono::steady_clock::time_point last_successful_heartbeat_{};
 };
 
 }  // namespace org::apache::nifi::minifi::c2

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 <cstdint>
+
+namespace org::apache::nifi::minifi::c2 {
+
+constexpr int32_t C2_CONNECTION_TIMEOUT_EXIT_CODE = 90;
+
+}  // namespace org::apache::nifi::minifi::c2

libminifi/src/Configuration.cpp

@@ -97,6 +97,8 @@ const std::unordered_map<std::string_view, gsl::not_null<const core::PropertyVal
   {Configuration::nifi_c2_rest_request_encoding, gsl::make_not_null(&core::StandardPropertyValidators::ALWAYS_VALID_VALIDATOR)},
   {Configuration::nifi_c2_rest_heartbeat_minimize_updates, gsl::make_not_null(&core::StandardPropertyValidators::BOOLEAN_VALIDATOR)},
   {Configuration::nifi_c2_flow_info_processor_bulletin_limit, gsl::make_not_null(&core::StandardPropertyValidators::UNSIGNED_INTEGER_VALIDATOR)},
+  {Configuration::nifi_c2_asset_download_timeout, gsl::make_not_null(&core::StandardPropertyValidators::TIME_PERIOD_VALIDATOR)},
+  {Configuration::nifi_c2_connection_timeout, gsl::make_not_null(&core::StandardPropertyValidators::TIME_PERIOD_VALIDATOR)},
   {Configuration::nifi_state_storage_local, gsl::make_not_null(&core::StandardPropertyValidators::ALWAYS_VALID_VALIDATOR)},
   {Configuration::nifi_state_storage_local_old, gsl::make_not_null(&core::StandardPropertyValidators::ALWAYS_VALID_VALIDATOR)},
   {Configuration::nifi_state_storage_local_class_name, gsl::make_not_null(&core::StandardPropertyValidators::ALWAYS_VALID_VALIDATOR)},

libminifi/src/FlowController.cpp

@@ -49,7 +49,7 @@ namespace org::apache::nifi::minifi {
 FlowController::FlowController(std::shared_ptr<core::Repository> provenance_repo, std::shared_ptr<core::Repository> flow_file_repo,
                                std::shared_ptr<Configure> configure, std::shared_ptr<core::FlowConfiguration> flow_configuration,
                                std::shared_ptr<core::ContentRepository> content_repo, std::unique_ptr<state::MetricsPublisherStore> metrics_publisher_store,
-                               std::shared_ptr<utils::file::FileSystem> filesystem, std::function<void()> request_restart,
+                               std::shared_ptr<utils::file::FileSystem> filesystem, RequestShutdownCallback request_shutdown,
                                utils::file::AssetManager* asset_manager, core::BulletinStore* bulletin_store)
     : core::controller::ForwardingControllerServiceProvider(core::className<FlowController>()),
       running_(false),
@@ -82,7 +82,7 @@ FlowController::FlowController(std::shared_ptr<core::Repository> provenance_repo
     if (auto publisher = metrics_publisher_store_->getMetricsPublisher(c2::C2_METRICS_PUBLISHER).lock()) {
       c2_metrics_publisher = std::dynamic_pointer_cast<c2::C2MetricsPublisher>(publisher);
     }
-    c2_agent_ = std::make_unique<c2::C2Agent>(configuration_, c2_metrics_publisher, std::move(filesystem), std::move(request_restart), asset_manager);
+    c2_agent_ = std::make_unique<c2::C2Agent>(configuration_, c2_metrics_publisher, std::move(filesystem), std::move(request_shutdown), asset_manager);
   }
 
   if (c2::isControllerSocketEnabled(configuration_)) {