From 56d58112830d69a259464973f641b8403de552b0 Mon Sep 17 00:00:00 2001
From: Julian Brost <julian.brost@icinga.com>
Date: Mon, 12 Feb 2024 16:28:37 +0100
Subject: [PATCH] AsioTlsStream: add GracefulDisconnect() and ForceDisconnect()

Calling `AsioTlsStream::async_shutdown()` performs a TLS shutdown which
exchanges messages (that's why it takes a `yield_context`) and thus has the
potential to block the coroutine. Therefore, it should be protected with a
timeout. As `async_shutdown()` doesn't simply take a timeout, this has to be
implemented using a timer. So far, these timers are scattered throughout the
codebase with some places missing them entirely. This commit adds helper
functions to properly shutdown a TLS connection with a single function call.
---
 lib/base/tlsstream.cpp | 64 ++++++++++++++++++++++++++++++++++++++++++
 lib/base/tlsstream.hpp |  3 ++
 2 files changed, 67 insertions(+)

diff --git a/lib/base/tlsstream.cpp b/lib/base/tlsstream.cpp
index a71451a53..ed8005837 100644
--- a/lib/base/tlsstream.cpp
+++ b/lib/base/tlsstream.cpp
@@ -7,6 +7,8 @@
 #include "base/logger.hpp"
 #include "base/configuration.hpp"
 #include "base/convert.hpp"
+#include "base/defer.hpp"
+#include "base/io-engine.hpp"
 #include <boost/asio/ssl/context.hpp>
 #include <boost/asio/ssl/verify_context.hpp>
 #include <boost/asio/ssl/verify_mode.hpp>
@@ -103,3 +105,65 @@ void UnbufferedAsioTlsStream::BeforeHandshake(handshake_type type)
 	}
 #endif /* SSL_CTRL_SET_TLSEXT_HOSTNAME */
 }
+
+/**
+ * Forcefully close the connection, typically (details are up to the operating system) using a TCP RST.
+ */
+void AsioTlsStream::ForceDisconnect()
+{
+	if (!lowest_layer().is_open()) {
+		// Already disconnected, nothing to do.
+		return;
+	}
+
+	boost::system::error_code ec;
+
+	// Close the socket. In case the connection wasn't shut down cleanly by GracefulDisconnect(), the operating system
+	// will typically terminate the connection with a TCP RST. Otherwise, this just releases the file descriptor.
+	lowest_layer().close(ec);
+}
+
+/**
+ * Try to cleanly shut down the connection. This involves sending a TLS close_notify shutdown alert and terminating the
+ * underlying TCP connection. Sending these additional messages can block, hence the method takes a yield context and
+ * internally implements a timeout of 10 seconds for the operation after which the connection is forcefully terminated
+ * using ForceDisconnect().
+ *
+ * @param strand Asio strand used for other operations on this connection.
+ * @param yc Yield context for Asio coroutines
+ */
+void AsioTlsStream::GracefulDisconnect(boost::asio::io_context::strand& strand, boost::asio::yield_context& yc)
+{
+	if (!lowest_layer().is_open()) {
+		// Already disconnected, nothing to do.
+		return;
+	}
+
+	{
+		Timeout::Ptr shutdownTimeout(new Timeout(strand.context(), strand, boost::posix_time::seconds(10),
+			[this](boost::asio::yield_context yc) {
+				// Forcefully terminate the connection if async_shutdown() blocked more than 10 seconds.
+				ForceDisconnect();
+			}
+		));
+		Defer cancelTimeout ([&shutdownTimeout]() {
+			shutdownTimeout->Cancel();
+		});
+
+		// Close the TLS connection, effectively uses SSL_shutdown() to send a close_notify shutdown alert to the peer.
+		boost::system::error_code ec;
+		next_layer().async_shutdown(yc[ec]);
+	}
+
+	if (!lowest_layer().is_open()) {
+		// Connection got closed in the meantime, most likely by the timeout, so nothing more to do.
+		return;
+	}
+
+	// Shut down the TCP connection.
+	boost::system::error_code ec;
+	lowest_layer().shutdown(lowest_layer_type::shutdown_both, ec);
+
+	// Clean up the connection (closes the file descriptor).
+	ForceDisconnect();
+}
diff --git a/lib/base/tlsstream.hpp b/lib/base/tlsstream.hpp
index 9a6340baf..9eed8d3b1 100644
--- a/lib/base/tlsstream.hpp
+++ b/lib/base/tlsstream.hpp
@@ -111,6 +111,9 @@ public:
 	{
 	}
 
+	void ForceDisconnect();
+	void GracefulDisconnect(boost::asio::io_context::strand& strand, boost::asio::yield_context& yc);
+
 private:
 	inline
 	AsioTlsStream(UnbufferedAsioTlsStreamParams init)