libminifi/src/c2/ControllerSocketProtocol.cpp - nifi-minifi-cpp - Git at Google

 /**
  * 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 "c2/ControllerSocketProtocol.h"

 #include <fstream>
 #include <sstream>
 #include <string>
 #include <utility>
 #include <vector>

 #include "asio/detached.hpp"
 #include "asio/ssl/stream.hpp"
 #include "c2/C2Payload.h"
 #include "c2/C2Utils.h"
 #include "controllers/SSLContextService.h"
 #include "io/AsioStream.h"
 #include "properties/Configuration.h"
 #include "utils/ConfigurationUtils.h"
 #include "utils/StringUtils.h"
 #include "minifi-cpp/utils/gsl.h"
 #include "utils/net/AsioSocketUtils.h"

 namespace org::apache::nifi::minifi::c2 {

 ControllerSocketProtocol::SocketRestartCommandProcessor::SocketRestartCommandProcessor(state::StateMonitor& update_sink, const std::shared_ptr<core::logging::Logger>& logger) :
     update_sink_(update_sink),
     logger_(logger) {
   command_queue_.start();
   command_processor_thread_ = std::thread([this] {
     while (running_) {
       CommandData command_data;
       if (command_queue_.dequeueWait(command_data)) {
         if (command_data.command == Command::FLOW_UPDATE) {
           auto result = update_sink_.applyUpdate("ControllerSocketProtocol", command_data.data, true);
           if (!result) {
             logger_->log_error("Failed to apply flow update: {}", result.error());
           }
         } else if (command_data.command == Command::START) {
           update_sink_.executeOnComponent(command_data.data, [](state::StateController& component) {
             component.start();
           });
         }
       }
       is_socket_restarting_ = false;
     }
   });
 }

 ControllerSocketProtocol::SocketRestartCommandProcessor::~SocketRestartCommandProcessor() {
   running_ = false;
   command_queue_.stop();
   if (command_processor_thread_.joinable()) {
     command_processor_thread_.join();
   }
 }

 ControllerSocketProtocol::ControllerSocketProtocol(state::StateMonitor& update_sink, std::shared_ptr<Configure> configuration,
   const std::shared_ptr<ControllerSocketReporter>& controller_socket_reporter)
     : update_sink_(update_sink),
       controller_socket_reporter_(controller_socket_reporter),
       configuration_(std::move(configuration)),
       socket_restart_processor_(update_sink_, logger_) {
   gsl_Expects(configuration_);
 }

 ControllerSocketProtocol::~ControllerSocketProtocol() {
   stopListener();
 }

 void ControllerSocketProtocol::stopListener() {
   if (acceptor_) {
     asio::post(io_context_, [this] {
       acceptor_->close();
     });
   }
   if (server_thread_.joinable()) {
     server_thread_.join();
   }
   io_context_.restart();
 }

 asio::awaitable<void> ControllerSocketProtocol::startAccept() {
   while (true) {
     auto [accept_error, socket] = co_await acceptor_->async_accept(utils::net::use_nothrow_awaitable);
     if (accept_error) {
       if (accept_error == asio::error::operation_aborted || accept_error == asio::error::bad_descriptor) {
         logger_->log_debug("Controller socket accept aborted");
         co_return;
       }
       logger_->log_error("Controller socket accept failed with the following message: '{}'", accept_error.message());
       continue;
     }
     auto stream = std::make_unique<io::AsioStream<asio::ip::tcp::socket>>(std::move(socket));
     co_spawn(io_context_, handleCommand(std::move(stream)), asio::detached);
   }
 }

 asio::awaitable<void> ControllerSocketProtocol::handshakeAndHandleCommand(asio::ip::tcp::socket&& socket, std::shared_ptr<minifi::controllers::SSLContextServiceInterface> ssl_context_service) {
   asio::ssl::context ssl_context = utils::net::getSslContext(*ssl_context_service, asio::ssl::context::tls_server);
   ssl_context.set_options(utils::net::MINIFI_SSL_OPTIONS);
   asio::ssl::stream<asio::ip::tcp::socket> ssl_socket(std::move(socket), ssl_context);

   auto [handshake_error] = co_await ssl_socket.async_handshake(utils::net::HandshakeType::server, utils::net::use_nothrow_awaitable);
   if (handshake_error) {
     logger_->log_error("Controller socket handshake failed with the following message: '{}'", handshake_error.message());
     co_return;
   }

   auto stream = std::make_unique<io::AsioStream<asio::ssl::stream<asio::ip::tcp::socket>>>(std::move(ssl_socket));
   co_return co_await handleCommand(std::move(stream));
 }

 asio::awaitable<void> ControllerSocketProtocol::startAcceptSsl(std::shared_ptr<minifi::controllers::SSLContextServiceInterface> ssl_context_service) {
   while (true) {  // NOLINT(clang-analyzer-core.NullDereference) suppressing asio library linter warning
     auto [accept_error, socket] = co_await acceptor_->async_accept(utils::net::use_nothrow_awaitable);
     if (accept_error) {
       if (accept_error == asio::error::operation_aborted || accept_error == asio::error::bad_descriptor) {
         logger_->log_debug("Controller socket accept aborted");
         co_return;
       }
       logger_->log_error("Controller socket accept failed with the following message: '{}'", accept_error.message());
       continue;
     }

     co_spawn(io_context_, handshakeAndHandleCommand(std::move(socket), ssl_context_service), asio::detached);
   }
 }

 void ControllerSocketProtocol::initialize() {
   std::unique_lock<std::mutex> lock(initialization_mutex_);
   std::shared_ptr<minifi::controllers::SSLContextServiceInterface> secure_context;
   std::string secure_str;
   if (configuration_->get(Configure::nifi_remote_input_secure, secure_str) && org::apache::nifi::minifi::utils::string::toBool(secure_str).value_or(false)) {
     secure_context = std::make_shared<minifi::controllers::SSLContextService>("ControllerSocketProtocolSSL", configuration_);
     secure_context->onEnable();
   }

   std::string limit_str;
   const bool any_interface = configuration_->get(Configuration::controller_socket_local_any_interface, limit_str) && utils::string::toBool(limit_str).value_or(false);

   // if host name isn't defined we will use localhost
   std::string host = "localhost";
   configuration_->get(Configuration::controller_socket_host, host);

   std::string port;
   stopListener();
   if (configuration_->get(Configuration::controller_socket_port, port)) {
     // if we have a localhost hostname and we did not manually specify any.interface we will bind only to the loopback adapter
     if ((host == "localhost" || host == "127.0.0.1" || host == "::") && !any_interface) {
       acceptor_ = std::make_unique<asio::ip::tcp::acceptor>(io_context_, asio::ip::tcp::endpoint(asio::ip::address_v4::loopback(), std::stoi(port)));
     } else {
       acceptor_ = std::make_unique<asio::ip::tcp::acceptor>(io_context_, asio::ip::tcp::endpoint(asio::ip::tcp::v4(), std::stoi(port)));
     }

     if (secure_context) {
       co_spawn(io_context_, startAcceptSsl(std::move(secure_context)), asio::detached);
     } else {
       co_spawn(io_context_, startAccept(), asio::detached);
     }
     server_thread_ = std::thread([this] {
       io_context_.run();
     });
   }
 }

 void ControllerSocketProtocol::setRoot(core::ProcessGroup* root) {
   if (auto controller_socket_reporter = controller_socket_reporter_.lock()) {
     controller_socket_reporter->setRoot(root);
   }
 }

 void ControllerSocketProtocol::handleStart(io::BaseStream &stream) {
   std::string component_str;
   const auto size = stream.read(component_str);
   if (!io::isError(size)) {
     if (component_str == "FlowController") {
       // Starting flow controller resets socket
       socket_restart_processor_.enqueue({SocketRestartCommandProcessor::Command::START, component_str});
     } else {
       update_sink_.executeOnComponent(component_str, [](state::StateController& component) {
         component.start();
       });
     }
   } else {
     logger_->log_error("Connection broke");
   }
 }

 void ControllerSocketProtocol::handleStop(io::BaseStream &stream) {
   std::string component_str;
   const auto size = stream.read(component_str);
   if (!io::isError(size)) {
     update_sink_.executeOnComponent(component_str, [](state::StateController& component) {
       component.stop();
     });
   } else {
     logger_->log_error("Connection broke");
   }
 }

 void ControllerSocketProtocol::handleClear(io::BaseStream &stream) {
   std::string connection;
   const auto size = stream.read(connection);
   if (!io::isError(size)) {
     update_sink_.clearConnection(connection);
   }
 }

 void ControllerSocketProtocol::handleUpdate(io::BaseStream &stream) {
   std::string what;
   {
     const auto size = stream.read(what);
     if (io::isError(size)) {
       logger_->log_error("Connection broke");
       return;
     }
   }
   if (what == "flow") {
     std::string ff_loc;
     {
       const auto size = stream.read(ff_loc);
       if (io::isError(size)) {
         logger_->log_error("Connection broke");
         return;
       }
     }
     std::ifstream tf(ff_loc);
     std::string flow_configuration((std::istreambuf_iterator<char>(tf)),
         std::istreambuf_iterator<char>());
     socket_restart_processor_.enqueue({SocketRestartCommandProcessor::Command::FLOW_UPDATE, flow_configuration});
   }
 }

 void ControllerSocketProtocol::writeQueueSizesResponse(io::BaseStream &stream) {
   std::string connection;
   const auto size_ = stream.read(connection);
   if (io::isError(size_)) {
     logger_->log_error("Connection broke");
     return;
   }
   std::unordered_map<std::string, ControllerSocketReporter::QueueSize> sizes;
   if (auto controller_socket_reporter = controller_socket_reporter_.lock()) {
     sizes = controller_socket_reporter->getQueueSizes();
   }
   std::stringstream response;
   if (sizes.contains(connection)) {
     response << sizes[connection].queue_size << " / " << sizes[connection].queue_size_max;
   } else {
     response << "not found";
   }
   io::BufferStream resp;
   auto op = static_cast<uint8_t>(Operation::describe);
   resp.write(&op, 1);
   resp.write(response.str());
   stream.write(resp.getBuffer());
 }

 void ControllerSocketProtocol::writeComponentsResponse(io::BaseStream &stream) {
   std::vector<std::pair<std::string, bool>> components;
   update_sink_.executeOnAllComponents([&components](state::StateController& component) {
     components.emplace_back(component.getComponentName(), component.isRunning());
   });
   io::BufferStream resp;
   auto op = static_cast<uint8_t>(Operation::describe);
   resp.write(&op, 1);
   resp.write(gsl::narrow<uint16_t>(components.size()));
   for (const auto& [name, is_running] : components) {
     resp.write(name);
     resp.write(is_running ? "true" : "false");
   }

   stream.write(resp.getBuffer());
 }

 void ControllerSocketProtocol::writeConnectionsResponse(io::BaseStream &stream) {
   io::BufferStream resp;
   auto op = static_cast<uint8_t>(Operation::describe);
   resp.write(&op, 1);
   std::unordered_set<std::string> connections;
   if (auto controller_socket_reporter = controller_socket_reporter_.lock()) {
     connections = controller_socket_reporter->getConnections();
   }

   const auto size = gsl::narrow<uint16_t>(connections.size());
   resp.write(size);
   for (const auto &connection : connections) {
     resp.write(connection, false);
   }
   stream.write(resp.getBuffer());
 }

 void ControllerSocketProtocol::writeGetFullResponse(io::BaseStream &stream) {
   io::BufferStream resp;
   auto op = static_cast<uint8_t>(Operation::describe);
   resp.write(&op, 1);
   std::unordered_set<std::string> full_connections;
   if (auto controller_socket_reporter = controller_socket_reporter_.lock()) {
     full_connections = controller_socket_reporter->getFullConnections();
   }

   const auto size = gsl::narrow<uint16_t>(full_connections.size());
   resp.write(size);
   for (const auto &connection : full_connections) {
     resp.write(connection, false);
   }
   stream.write(resp.getBuffer());
 }

 void ControllerSocketProtocol::writeManifestResponse(io::BaseStream &stream) {
   io::BufferStream resp;
   auto op = static_cast<uint8_t>(Operation::describe);
   resp.write(&op, 1);
   std::string manifest;
   if (auto controller_socket_reporter = controller_socket_reporter_.lock()) {
     manifest = controller_socket_reporter->getAgentManifest();
   }
   resp.write(manifest, true);
   stream.write(resp.getBuffer());
 }

 std::string ControllerSocketProtocol::getJstack() {
   if (!update_sink_.isRunning()) {
     return {};
   }
   std::stringstream result;
   const auto traces = update_sink_.getTraces();
   for (const auto& trace : traces) {
     for (const auto& line : trace.getTraces()) {
       result << trace.getName() << " -- " << line << "\n";
     }
   }
   return result.str();
 }

 void ControllerSocketProtocol::writeJstackResponse(io::BaseStream &stream) {
   io::BufferStream resp;
   auto op = static_cast<uint8_t>(Operation::describe);
   resp.write(&op, 1);
   std::string jstack_response;
   if (auto controller_socket_reporter = controller_socket_reporter_.lock()) {
     jstack_response = getJstack();
   }
   resp.write(jstack_response, true);
   stream.write(resp.getBuffer());
 }

 void ControllerSocketProtocol::writeFlowStatusResponse(io::BaseStream &stream) {
   std::string query;
   {
     const auto size = stream.read(query);
     if (io::isError(size)) {
       logger_->log_error("Connection broke");
       return;
     }
   }
   auto request_strings = utils::string::splitAndTrimRemovingEmpty(query, ";");
   std::vector<FlowStatusRequest> requests;
   for (const auto& request_string : request_strings) {
     try {
       requests.push_back(FlowStatusRequest(request_string));
     } catch (const std::exception& e) {
       logger_->log_error("Invalid flow status request: {}", e.what());
     }
   }
   io::BufferStream resp;
   auto op = static_cast<uint8_t>(Operation::describe);
   resp.write(&op, 1);
   std::string flowstatus_response;
   if (auto controller_socket_reporter = controller_socket_reporter_.lock()) {
     flowstatus_response = controller_socket_reporter->getFlowStatus(requests);
   }
   resp.write(flowstatus_response, true);
   stream.write(resp.getBuffer());
 }

 void ControllerSocketProtocol::handleDescribe(io::BaseStream &stream) {
   std::string what;
   const auto size = stream.read(what);
   if (io::isError(size)) {
     logger_->log_error("Connection broke");
     return;
   }
   if (what == "queue") {
     writeQueueSizesResponse(stream);
   } else if (what == "components") {
     writeComponentsResponse(stream);
   } else if (what == "connections") {
     writeConnectionsResponse(stream);
   } else if (what == "getfull") {
     writeGetFullResponse(stream);
   } else if (what == "manifest") {
     writeManifestResponse(stream);
   } else if (what == "jstack") {
     writeJstackResponse(stream);
   } else if (what == "flowstatus") {
     writeFlowStatusResponse(stream);
   } else {
     logger_->log_error("Unknown C2 describe parameter: {}", what);
   }
 }

 void ControllerSocketProtocol::writeDebugBundleResponse(io::BaseStream &stream) {
   auto files = update_sink_.getDebugInfo();
   auto bundle = createDebugBundleArchive(files);
   io::BufferStream resp;
   auto op = static_cast<uint8_t>(Operation::transfer);
   resp.write(&op, 1);
   if (!bundle) {
     logger_->log_error("Creating debug bundle failed: {}", bundle.error());
     resp.write(static_cast<size_t>(0));
     stream.write(resp.getBuffer());
     return;
   }

   size_t bundle_size = bundle.value()->size();
   resp.write(bundle_size);
   static constexpr auto BUFFER_SIZE = utils::configuration::DEFAULT_BUFFER_SIZE;
   std::array<std::byte, BUFFER_SIZE> out_buffer{};
   while (bundle_size > 0) {
     const auto next_write_size = (std::min)(bundle_size, BUFFER_SIZE);
     const auto size_read = bundle.value()->read(std::as_writable_bytes(std::span(out_buffer).subspan(0, next_write_size)));
     resp.write(reinterpret_cast<const uint8_t*>(out_buffer.data()), size_read);
     bundle_size -= size_read;
   }

   stream.write(resp.getBuffer());
 }

 void ControllerSocketProtocol::handleTransfer(io::BaseStream &stream) {
   std::string what;
   const auto size = stream.read(what);
   if (io::isError(size)) {
     logger_->log_error("Connection broke");
     return;
   }
   if (what == "debug") {
     writeDebugBundleResponse(stream);
   } else {
     logger_->log_error("Unknown C2 transfer parameter: {}", what);
   }
 }

 asio::awaitable<void> ControllerSocketProtocol::handleCommand(std::unique_ptr<io::BaseStream> stream) {
   uint8_t head = 0;
   if (stream->read(head) != 1) {
     logger_->log_error("Connection broke");
     co_return;
   }

   if (socket_restart_processor_.isSocketRestarting()) {
     logger_->log_debug("Socket restarting, dropping command");
     co_return;
   }

   auto op = static_cast<Operation>(head);
   switch (op) {
     case Operation::start:
       handleStart(*stream);
       break;
     case Operation::stop:
       handleStop(*stream);
       break;
     case Operation::clear:
       handleClear(*stream);
       break;
     case Operation::update:
       handleUpdate(*stream);
       break;
     case Operation::describe:
       handleDescribe(*stream);
       break;
     case Operation::transfer:
       handleTransfer(*stream);
       break;
     default:
       logger_->log_error("Unhandled C2 operation: {}", head);
   }
 }

 }  // namespace org::apache::nifi::minifi::c2
	/**
	* 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 "c2/ControllerSocketProtocol.h"

	#include <fstream>
	#include <sstream>
	#include <string>
	#include <utility>
	#include <vector>

	#include "asio/detached.hpp"
	#include "asio/ssl/stream.hpp"
	#include "c2/C2Payload.h"
	#include "c2/C2Utils.h"
	#include "controllers/SSLContextService.h"
	#include "io/AsioStream.h"
	#include "properties/Configuration.h"
	#include "utils/ConfigurationUtils.h"
	#include "utils/StringUtils.h"
	#include "minifi-cpp/utils/gsl.h"
	#include "utils/net/AsioSocketUtils.h"

	namespace org::apache::nifi::minifi::c2 {

	ControllerSocketProtocol::SocketRestartCommandProcessor::SocketRestartCommandProcessor(state::StateMonitor& update_sink, const std::shared_ptr<core::logging::Logger>& logger) :
	update_sink_(update_sink),
	logger_(logger) {
	command_queue_.start();
	command_processor_thread_ = std::thread([this] {
	while (running_) {
	CommandData command_data;
	if (command_queue_.dequeueWait(command_data)) {
	if (command_data.command == Command::FLOW_UPDATE) {
	auto result = update_sink_.applyUpdate("ControllerSocketProtocol", command_data.data, true);
	if (!result) {
	logger_->log_error("Failed to apply flow update: {}", result.error());
	}
	} else if (command_data.command == Command::START) {
	update_sink_.executeOnComponent(command_data.data, [](state::StateController& component) {
	component.start();
	});
	}
	}
	is_socket_restarting_ = false;
	}
	});
	}

	ControllerSocketProtocol::SocketRestartCommandProcessor::~SocketRestartCommandProcessor() {
	running_ = false;
	command_queue_.stop();
	if (command_processor_thread_.joinable()) {
	command_processor_thread_.join();
	}
	}

	ControllerSocketProtocol::ControllerSocketProtocol(state::StateMonitor& update_sink, std::shared_ptr<Configure> configuration,
	const std::shared_ptr<ControllerSocketReporter>& controller_socket_reporter)
	: update_sink_(update_sink),
	controller_socket_reporter_(controller_socket_reporter),
	configuration_(std::move(configuration)),
	socket_restart_processor_(update_sink_, logger_) {
	gsl_Expects(configuration_);
	}

	ControllerSocketProtocol::~ControllerSocketProtocol() {
	stopListener();
	}

	void ControllerSocketProtocol::stopListener() {
	if (acceptor_) {
	asio::post(io_context_, [this] {
	acceptor_->close();
	});
	}
	if (server_thread_.joinable()) {
	server_thread_.join();
	}
	io_context_.restart();
	}

	asio::awaitable<void> ControllerSocketProtocol::startAccept() {
	while (true) {
	auto [accept_error, socket] = co_await acceptor_->async_accept(utils::net::use_nothrow_awaitable);
	if (accept_error) {
	if (accept_error == asio::error::operation_aborted \|\| accept_error == asio::error::bad_descriptor) {
	logger_->log_debug("Controller socket accept aborted");
	co_return;
	}
	logger_->log_error("Controller socket accept failed with the following message: '{}'", accept_error.message());
	continue;
	}
	auto stream = std::make_unique<io::AsioStream<asio::ip::tcp::socket>>(std::move(socket));
	co_spawn(io_context_, handleCommand(std::move(stream)), asio::detached);
	}
	}

	asio::awaitable<void> ControllerSocketProtocol::handshakeAndHandleCommand(asio::ip::tcp::socket&& socket, std::shared_ptr<minifi::controllers::SSLContextServiceInterface> ssl_context_service) {
	asio::ssl::context ssl_context = utils::net::getSslContext(*ssl_context_service, asio::ssl::context::tls_server);
	ssl_context.set_options(utils::net::MINIFI_SSL_OPTIONS);
	asio::ssl::stream<asio::ip::tcp::socket> ssl_socket(std::move(socket), ssl_context);

	auto [handshake_error] = co_await ssl_socket.async_handshake(utils::net::HandshakeType::server, utils::net::use_nothrow_awaitable);
	if (handshake_error) {
	logger_->log_error("Controller socket handshake failed with the following message: '{}'", handshake_error.message());
	co_return;
	}

	auto stream = std::make_unique<io::AsioStream<asio::ssl::stream<asio::ip::tcp::socket>>>(std::move(ssl_socket));
	co_return co_await handleCommand(std::move(stream));
	}

	asio::awaitable<void> ControllerSocketProtocol::startAcceptSsl(std::shared_ptr<minifi::controllers::SSLContextServiceInterface> ssl_context_service) {
	while (true) { // NOLINT(clang-analyzer-core.NullDereference) suppressing asio library linter warning
	auto [accept_error, socket] = co_await acceptor_->async_accept(utils::net::use_nothrow_awaitable);
	if (accept_error) {
	if (accept_error == asio::error::operation_aborted \|\| accept_error == asio::error::bad_descriptor) {
	logger_->log_debug("Controller socket accept aborted");
	co_return;
	}
	logger_->log_error("Controller socket accept failed with the following message: '{}'", accept_error.message());
	continue;
	}

	co_spawn(io_context_, handshakeAndHandleCommand(std::move(socket), ssl_context_service), asio::detached);
	}
	}

	void ControllerSocketProtocol::initialize() {
	std::unique_lock<std::mutex> lock(initialization_mutex_);
	std::shared_ptr<minifi::controllers::SSLContextServiceInterface> secure_context;
	std::string secure_str;
	if (configuration_->get(Configure::nifi_remote_input_secure, secure_str) && org::apache::nifi::minifi::utils::string::toBool(secure_str).value_or(false)) {
	secure_context = std::make_shared<minifi::controllers::SSLContextService>("ControllerSocketProtocolSSL", configuration_);
	secure_context->onEnable();
	}

	std::string limit_str;
	const bool any_interface = configuration_->get(Configuration::controller_socket_local_any_interface, limit_str) && utils::string::toBool(limit_str).value_or(false);

	// if host name isn't defined we will use localhost
	std::string host = "localhost";
	configuration_->get(Configuration::controller_socket_host, host);

	std::string port;
	stopListener();
	if (configuration_->get(Configuration::controller_socket_port, port)) {
	// if we have a localhost hostname and we did not manually specify any.interface we will bind only to the loopback adapter
	if ((host == "localhost" \|\| host == "127.0.0.1" \|\| host == "::") && !any_interface) {
	acceptor_ = std::make_unique<asio::ip::tcp::acceptor>(io_context_, asio::ip::tcp::endpoint(asio::ip::address_v4::loopback(), std::stoi(port)));
	} else {
	acceptor_ = std::make_unique<asio::ip::tcp::acceptor>(io_context_, asio::ip::tcp::endpoint(asio::ip::tcp::v4(), std::stoi(port)));
	}

	if (secure_context) {
	co_spawn(io_context_, startAcceptSsl(std::move(secure_context)), asio::detached);
	} else {
	co_spawn(io_context_, startAccept(), asio::detached);
	}
	server_thread_ = std::thread([this] {
	io_context_.run();
	});
	}
	}

	void ControllerSocketProtocol::setRoot(core::ProcessGroup* root) {
	if (auto controller_socket_reporter = controller_socket_reporter_.lock()) {
	controller_socket_reporter->setRoot(root);
	}
	}

	void ControllerSocketProtocol::handleStart(io::BaseStream &stream) {
	std::string component_str;
	const auto size = stream.read(component_str);
	if (!io::isError(size)) {
	if (component_str == "FlowController") {
	// Starting flow controller resets socket
	socket_restart_processor_.enqueue({SocketRestartCommandProcessor::Command::START, component_str});
	} else {
	update_sink_.executeOnComponent(component_str, [](state::StateController& component) {
	component.start();
	});
	}
	} else {
	logger_->log_error("Connection broke");
	}
	}

	void ControllerSocketProtocol::handleStop(io::BaseStream &stream) {
	std::string component_str;
	const auto size = stream.read(component_str);
	if (!io::isError(size)) {
	update_sink_.executeOnComponent(component_str, [](state::StateController& component) {
	component.stop();
	});
	} else {
	logger_->log_error("Connection broke");
	}
	}

	void ControllerSocketProtocol::handleClear(io::BaseStream &stream) {
	std::string connection;
	const auto size = stream.read(connection);
	if (!io::isError(size)) {
	update_sink_.clearConnection(connection);
	}
	}

	void ControllerSocketProtocol::handleUpdate(io::BaseStream &stream) {
	std::string what;
	{
	const auto size = stream.read(what);
	if (io::isError(size)) {
	logger_->log_error("Connection broke");
	return;
	}
	}
	if (what == "flow") {
	std::string ff_loc;
	{
	const auto size = stream.read(ff_loc);
	if (io::isError(size)) {
	logger_->log_error("Connection broke");
	return;
	}
	}
	std::ifstream tf(ff_loc);
	std::string flow_configuration((std::istreambuf_iterator<char>(tf)),
	std::istreambuf_iterator<char>());
	socket_restart_processor_.enqueue({SocketRestartCommandProcessor::Command::FLOW_UPDATE, flow_configuration});
	}
	}

	void ControllerSocketProtocol::writeQueueSizesResponse(io::BaseStream &stream) {
	std::string connection;
	const auto size_ = stream.read(connection);
	if (io::isError(size_)) {
	logger_->log_error("Connection broke");
	return;
	}
	std::unordered_map<std::string, ControllerSocketReporter::QueueSize> sizes;
	if (auto controller_socket_reporter = controller_socket_reporter_.lock()) {
	sizes = controller_socket_reporter->getQueueSizes();
	}
	std::stringstream response;
	if (sizes.contains(connection)) {
	response << sizes[connection].queue_size << " / " << sizes[connection].queue_size_max;
	} else {
	response << "not found";
	}
	io::BufferStream resp;
	auto op = static_cast<uint8_t>(Operation::describe);
	resp.write(&op, 1);
	resp.write(response.str());
	stream.write(resp.getBuffer());
	}

	void ControllerSocketProtocol::writeComponentsResponse(io::BaseStream &stream) {
	std::vector<std::pair<std::string, bool>> components;
	update_sink_.executeOnAllComponents([&components](state::StateController& component) {
	components.emplace_back(component.getComponentName(), component.isRunning());
	});
	io::BufferStream resp;
	auto op = static_cast<uint8_t>(Operation::describe);
	resp.write(&op, 1);
	resp.write(gsl::narrow<uint16_t>(components.size()));
	for (const auto& [name, is_running] : components) {
	resp.write(name);
	resp.write(is_running ? "true" : "false");
	}

	stream.write(resp.getBuffer());
	}

	void ControllerSocketProtocol::writeConnectionsResponse(io::BaseStream &stream) {
	io::BufferStream resp;
	auto op = static_cast<uint8_t>(Operation::describe);
	resp.write(&op, 1);
	std::unordered_set<std::string> connections;
	if (auto controller_socket_reporter = controller_socket_reporter_.lock()) {
	connections = controller_socket_reporter->getConnections();
	}

	const auto size = gsl::narrow<uint16_t>(connections.size());
	resp.write(size);
	for (const auto &connection : connections) {
	resp.write(connection, false);
	}
	stream.write(resp.getBuffer());
	}

	void ControllerSocketProtocol::writeGetFullResponse(io::BaseStream &stream) {
	io::BufferStream resp;
	auto op = static_cast<uint8_t>(Operation::describe);
	resp.write(&op, 1);
	std::unordered_set<std::string> full_connections;
	if (auto controller_socket_reporter = controller_socket_reporter_.lock()) {
	full_connections = controller_socket_reporter->getFullConnections();
	}

	const auto size = gsl::narrow<uint16_t>(full_connections.size());
	resp.write(size);
	for (const auto &connection : full_connections) {
	resp.write(connection, false);
	}
	stream.write(resp.getBuffer());
	}

	void ControllerSocketProtocol::writeManifestResponse(io::BaseStream &stream) {
	io::BufferStream resp;
	auto op = static_cast<uint8_t>(Operation::describe);
	resp.write(&op, 1);
	std::string manifest;
	if (auto controller_socket_reporter = controller_socket_reporter_.lock()) {
	manifest = controller_socket_reporter->getAgentManifest();
	}
	resp.write(manifest, true);
	stream.write(resp.getBuffer());
	}

	std::string ControllerSocketProtocol::getJstack() {
	if (!update_sink_.isRunning()) {
	return {};
	}
	std::stringstream result;
	const auto traces = update_sink_.getTraces();
	for (const auto& trace : traces) {
	for (const auto& line : trace.getTraces()) {
	result << trace.getName() << " -- " << line << "\n";
	}
	}
	return result.str();
	}

	void ControllerSocketProtocol::writeJstackResponse(io::BaseStream &stream) {
	io::BufferStream resp;
	auto op = static_cast<uint8_t>(Operation::describe);
	resp.write(&op, 1);
	std::string jstack_response;
	if (auto controller_socket_reporter = controller_socket_reporter_.lock()) {
	jstack_response = getJstack();
	}
	resp.write(jstack_response, true);
	stream.write(resp.getBuffer());
	}

	void ControllerSocketProtocol::writeFlowStatusResponse(io::BaseStream &stream) {
	std::string query;
	{
	const auto size = stream.read(query);
	if (io::isError(size)) {
	logger_->log_error("Connection broke");
	return;
	}
	}
	auto request_strings = utils::string::splitAndTrimRemovingEmpty(query, ";");
	std::vector<FlowStatusRequest> requests;
	for (const auto& request_string : request_strings) {
	try {
	requests.push_back(FlowStatusRequest(request_string));
	} catch (const std::exception& e) {
	logger_->log_error("Invalid flow status request: {}", e.what());
	}
	}
	io::BufferStream resp;
	auto op = static_cast<uint8_t>(Operation::describe);
	resp.write(&op, 1);
	std::string flowstatus_response;
	if (auto controller_socket_reporter = controller_socket_reporter_.lock()) {
	flowstatus_response = controller_socket_reporter->getFlowStatus(requests);
	}
	resp.write(flowstatus_response, true);
	stream.write(resp.getBuffer());
	}

	void ControllerSocketProtocol::handleDescribe(io::BaseStream &stream) {
	std::string what;
	const auto size = stream.read(what);
	if (io::isError(size)) {
	logger_->log_error("Connection broke");
	return;
	}
	if (what == "queue") {
	writeQueueSizesResponse(stream);
	} else if (what == "components") {
	writeComponentsResponse(stream);
	} else if (what == "connections") {
	writeConnectionsResponse(stream);
	} else if (what == "getfull") {
	writeGetFullResponse(stream);
	} else if (what == "manifest") {
	writeManifestResponse(stream);
	} else if (what == "jstack") {
	writeJstackResponse(stream);
	} else if (what == "flowstatus") {
	writeFlowStatusResponse(stream);
	} else {
	logger_->log_error("Unknown C2 describe parameter: {}", what);
	}
	}

	void ControllerSocketProtocol::writeDebugBundleResponse(io::BaseStream &stream) {
	auto files = update_sink_.getDebugInfo();
	auto bundle = createDebugBundleArchive(files);
	io::BufferStream resp;
	auto op = static_cast<uint8_t>(Operation::transfer);
	resp.write(&op, 1);
	if (!bundle) {
	logger_->log_error("Creating debug bundle failed: {}", bundle.error());
	resp.write(static_cast<size_t>(0));
	stream.write(resp.getBuffer());
	return;
	}

	size_t bundle_size = bundle.value()->size();
	resp.write(bundle_size);
	static constexpr auto BUFFER_SIZE = utils::configuration::DEFAULT_BUFFER_SIZE;
	std::array<std::byte, BUFFER_SIZE> out_buffer{};
	while (bundle_size > 0) {
	const auto next_write_size = (std::min)(bundle_size, BUFFER_SIZE);
	const auto size_read = bundle.value()->read(std::as_writable_bytes(std::span(out_buffer).subspan(0, next_write_size)));
	resp.write(reinterpret_cast<const uint8_t*>(out_buffer.data()), size_read);
	bundle_size -= size_read;
	}

	stream.write(resp.getBuffer());
	}

	void ControllerSocketProtocol::handleTransfer(io::BaseStream &stream) {
	std::string what;
	const auto size = stream.read(what);
	if (io::isError(size)) {
	logger_->log_error("Connection broke");
	return;
	}
	if (what == "debug") {
	writeDebugBundleResponse(stream);
	} else {
	logger_->log_error("Unknown C2 transfer parameter: {}", what);
	}
	}

	asio::awaitable<void> ControllerSocketProtocol::handleCommand(std::unique_ptr<io::BaseStream> stream) {
	uint8_t head = 0;
	if (stream->read(head) != 1) {
	logger_->log_error("Connection broke");
	co_return;
	}

	if (socket_restart_processor_.isSocketRestarting()) {
	logger_->log_debug("Socket restarting, dropping command");
	co_return;
	}

	auto op = static_cast<Operation>(head);
	switch (op) {
	case Operation::start:
	handleStart(*stream);
	break;
	case Operation::stop:
	handleStop(*stream);
	break;
	case Operation::clear:
	handleClear(*stream);
	break;
	case Operation::update:
	handleUpdate(*stream);
	break;
	case Operation::describe:
	handleDescribe(*stream);
	break;
	case Operation::transfer:
	handleTransfer(*stream);
	break;
	default:
	logger_->log_error("Unhandled C2 operation: {}", head);
	}
	}

	} // namespace org::apache::nifi::minifi::c2