From b6e216d8fe77e7e30a0b57f12efb5a9903b4ed61 Mon Sep 17 00:00:00 2001
From: Rene Meusel <rene.meusel@rohde-schwarz.com>
Date: Thu, 4 Apr 2024 14:07:24 +0200
Subject: [PATCH] Add scoped_cleanup helper to stl_util.h

---
 src/lib/utils/stl_util.h | 33 ++++++++++++++++++++
 src/tests/test_utils.cpp | 67 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 100 insertions(+)

diff --git a/src/lib/utils/stl_util.h b/src/lib/utils/stl_util.h
index 32196ec3f46..26f624370d3 100644
--- a/src/lib/utils/stl_util.h
+++ b/src/lib/utils/stl_util.h
@@ -12,6 +12,7 @@
 #include <algorithm>
 #include <functional>
 #include <map>
+#include <optional>
 #include <set>
 #include <span>
 #include <string>
@@ -328,6 +329,38 @@ struct overloaded : Ts... {
 template <class... Ts>
 overloaded(Ts...) -> overloaded<Ts...>;
 
+/**
+ * @brief Helper class to create a RAII-style cleanup callback
+ *
+ * Ensures that the cleanup callback given in the object's constructor is called
+ * when the object is destroyed. Use this to ensure some cleanup code runs when
+ * leaving the current scope.
+ */
+template <std::invocable FunT>
+class scoped_cleanup {
+   public:
+      explicit scoped_cleanup(FunT cleanup) : m_cleanup(std::move(cleanup)) {}
+
+      scoped_cleanup(const scoped_cleanup&) = delete;
+      scoped_cleanup& operator=(const scoped_cleanup&) = delete;
+      scoped_cleanup(scoped_cleanup&&) = delete;
+      scoped_cleanup& operator=(scoped_cleanup&&) = delete;
+
+      ~scoped_cleanup() {
+         if(m_cleanup.has_value()) {
+            m_cleanup.value()();
+         }
+      }
+
+      /**
+       * Disengage the cleanup callback, i.e., prevent it from being called
+       */
+      void disengage() { m_cleanup.reset(); }
+
+   private:
+      std::optional<FunT> m_cleanup;
+};
+
 }  // namespace Botan
 
 #endif
diff --git a/src/tests/test_utils.cpp b/src/tests/test_utils.cpp
index ae6781fd183..16a52c1374f 100644
--- a/src/tests/test_utils.cpp
+++ b/src/tests/test_utils.cpp
@@ -1179,6 +1179,73 @@ class Formatter_Tests : public Test {
 
 BOTAN_REGISTER_TEST("utils", "fmt", Formatter_Tests);
 
+class ScopedCleanup_Tests : public Test {
+   public:
+      std::vector<Test::Result> run() override {
+         return {
+            Botan_Tests::CHECK("leaving a scope results in cleanup",
+                               [](Test::Result& result) {
+                                  bool ran = false;
+                                  {
+                                     auto clean = Botan::scoped_cleanup([&] { ran = true; });
+                                  }
+                                  result.confirm("cleanup ran", ran);
+                               }),
+
+            Botan_Tests::CHECK("leaving a function, results in cleanup",
+                               [](Test::Result& result) {
+                                  bool ran = false;
+                                  bool fn_called = false;
+                                  auto fn = [&] {
+                                     auto clean = Botan::scoped_cleanup([&] { ran = true; });
+                                     fn_called = true;
+                                  };
+
+                                  result.confirm("cleanup not yet ran", !ran);
+                                  fn();
+                                  result.confirm("fn called", fn_called);
+                                  result.confirm("cleanup ran", ran);
+                               }),
+
+            Botan_Tests::CHECK("stack unwinding results in cleanup",
+                               [](Test::Result& result) {
+                                  bool ran = false;
+                                  bool fn_called = false;
+                                  bool exception_caught = false;
+                                  auto fn = [&] {
+                                     auto clean = Botan::scoped_cleanup([&] { ran = true; });
+                                     fn_called = true;
+                                     throw std::runtime_error("test");
+                                  };
+
+                                  result.confirm("cleanup not yet ran", !ran);
+                                  try {
+                                     fn();
+                                  } catch(const std::exception&) {
+                                     exception_caught = true;
+                                  }
+
+                                  result.confirm("fn called", fn_called);
+                                  result.confirm("cleanup ran", ran);
+                                  result.confirm("exception caught", exception_caught);
+                               }),
+
+            Botan_Tests::CHECK("cleanup isn't called after disengaging",
+                               [](Test::Result& result) {
+                                  bool ran = false;
+                                  {
+                                     auto clean = Botan::scoped_cleanup([&] { ran = true; });
+                                     clean.disengage();
+                                  }
+                                  result.confirm("cleanup not ran", !ran);
+                               }),
+
+         };
+      }
+};
+
+BOTAN_REGISTER_TEST("utils", "scoped_cleanup", ScopedCleanup_Tests);
+
 #endif
 
 }  // namespace