[Project] Allow mempool allocated mime strings

2 files changed, 42 insertions, 0 deletions

diff --git a/src/libmime/mime_string.hxx b/src/libmime/mime_string.hxx
index 259788cd5..7474c6381 100644
--- a/src/libmime/mime_string.hxx
+++ b/src/libmime/mime_string.hxx
@@ -24,6 +24,7 @@
 #include <cstdlib>
 #include <cstring>
 #include <iosfwd>
+#include "libutil/mem_pool.h"
 #include "function2/function2.hpp"
 #include "unicode/utf8.h"
 #include "contrib/fastutf8/fastutf8.h"
@@ -42,6 +43,7 @@ template<class T=char, class Allocator = std::allocator<T>,
 		class Functor = fu2::function_view<UChar32(UChar32)>> class basic_mime_string;
 
 using mime_string = basic_mime_string<char>;
+using mime_pool_string = basic_mime_string<char, mempool_allocator<char>>;
 
 /* Helpers for type safe flags */
 enum class mime_string_flags : std::uint8_t {
@@ -361,6 +363,16 @@ public:
 			append_c_string_unfiltered(other.data(), other.size());
 		}
 	}
+	auto assign_copy(const view_type &other) {
+		storage.clear();
+
+		if (filter_func) {
+			append_c_string_filtered(other.data(), other.size());
+		}
+		else {
+			append_c_string_unfiltered(other.data(), other.size());
+		}
+	}
 
 	/* Mutators */
 	auto append(const CharT* str, std::size_t size) -> std::size_t {
diff --git a/src/libutil/mem_pool.h b/src/libutil/mem_pool.h
index d4488ac55..d3927a1bc 100644
--- a/src/libutil/mem_pool.h
+++ b/src/libutil/mem_pool.h
@@ -397,4 +397,34 @@ GList *rspamd_mempool_glist_append (rspamd_mempool_t *pool,
 }
 #endif
 
+#ifdef  __cplusplus
+#include <memory> /* For allocator */
+namespace rspamd {
+
+template<class T>
+class mempool_allocator {
+public:
+	typedef T value_type;
+
+	mempool_allocator() = delete;
+	template<class U>
+	mempool_allocator(const mempool_allocator<U> &other) : pool(other.pool) {}
+	mempool_allocator(rspamd_mempool_t *_pool) : pool(_pool) {}
+	[[nodiscard]] constexpr T* allocate(std::size_t n)
+	{
+		if (G_MAXSIZE / 2 / sizeof(T) > n) {
+			throw std::runtime_error("integer overflow");
+		}
+		return reinterpret_cast<T*>(rspamd_mempool_alloc(pool, n * sizeof(T)));
+	}
+	constexpr void deallocate(T* p, std::size_t n) {
+		/* Do nothing */
+	}
+private:
+	rspamd_mempool_t *pool;
+};
+
+}
+#endif
+
 #endif