From 832f3e58c05887b76deb523eee96541eea5e8c97 Mon Sep 17 00:00:00 2001
From: Aster Seker <git.electroyar@gmail.com>
Date: Sat, 6 Sep 2025 07:14:14 +0300
Subject: [PATCH] feat(secure-buffer): add optional page locking

Add optional memory locking to secure_buffer with copy/move support and automatic unlocking.
---
 include/hmac_cpp/secure_buffer.hpp | 74 ++++++++++++++++++++++++++++--
 1 file changed, 69 insertions(+), 5 deletions(-)

diff --git a/include/hmac_cpp/secure_buffer.hpp b/include/hmac_cpp/secure_buffer.hpp
index c6d04e3..bafb468 100644
--- a/include/hmac_cpp/secure_buffer.hpp
+++ b/include/hmac_cpp/secure_buffer.hpp
@@ -6,6 +6,7 @@
 #include <vector>
 #include <type_traits>
 #include <string>
+#include "hmac_cpp/memlock.hpp"
 
 // Macro to mark deprecated APIs in a compiler-portable way
 #ifndef HMACCPP_DEPRECATED
@@ -32,31 +33,93 @@ inline void secure_zero(void* ptr, size_t len) {
 
 /// \brief Vector-like buffer that zeroizes its contents on destruction.
 /// \tparam T Trivial value type stored in the buffer (defaults to uint8_t).
-template<class T = uint8_t>
+/// \tparam LockOnAlloc Lock pages in memory on allocation.
+template<class T = uint8_t, bool LockOnAlloc = false>
 struct secure_buffer {
     static_assert(std::is_trivial<T>::value, "secure_buffer requires trivial type");
 
-    secure_buffer() = default;
+    secure_buffer() {
+        if (LockOnAlloc && !buf.empty()) {
+            locked_ = lock_pages(buf.data(), buf.size() * sizeof(T));
+        }
+    }
 
     /// \brief Construct with n default-initialized elements.
     /// \param n Element count.
-    explicit secure_buffer(size_t n) : buf(n) {}
+    explicit secure_buffer(size_t n) : buf(n) {
+        if (LockOnAlloc && !buf.empty()) {
+            locked_ = lock_pages(buf.data(), buf.size() * sizeof(T));
+        }
+    }
 
     /// \brief Construct from vector, moving its contents.
     /// \param v Source vector.
-    explicit secure_buffer(std::vector<T>&& v) : buf(std::move(v)) {}
+    explicit secure_buffer(std::vector<T>&& v) : buf(std::move(v)) {
+        if (LockOnAlloc && !buf.empty()) {
+            locked_ = lock_pages(buf.data(), buf.size() * sizeof(T));
+        }
+    }
 
     /// \brief Construct from std::string rvalue and zeroize the source.
     /// \param s Source string.
     template<class U = T, typename std::enable_if<std::is_same<U, uint8_t>::value, int>::type = 0>
     explicit secure_buffer(std::string&& s) : buf(s.begin(), s.end()) {
+        if (LockOnAlloc && !buf.empty()) {
+            locked_ = lock_pages(buf.data(), buf.size() * sizeof(T));
+        }
         if (!s.empty()) {
             secure_zero(&s[0], s.size());
             s.clear();
         }
     }
+
+    secure_buffer(const secure_buffer& other) : buf(other.buf) {
+        if (LockOnAlloc && !buf.empty()) {
+            locked_ = lock_pages(buf.data(), buf.size() * sizeof(T));
+        }
+    }
+
+    secure_buffer& operator=(const secure_buffer& other) {
+        if (this != &other) {
+            secure_zero(buf.data(), buf.size() * sizeof(T));
+            if (locked_) {
+                unlock_pages(buf.data(), buf.size() * sizeof(T));
+            }
+            buf = other.buf;
+            if (LockOnAlloc && !buf.empty()) {
+                locked_ = lock_pages(buf.data(), buf.size() * sizeof(T));
+            } else {
+                locked_ = false;
+            }
+        }
+        return *this;
+    }
+
+    secure_buffer(secure_buffer&& other) noexcept
+        : buf(std::move(other.buf)), locked_(other.locked_) {
+        other.locked_ = false;
+    }
+
+    secure_buffer& operator=(secure_buffer&& other) noexcept {
+        if (this != &other) {
+            secure_zero(buf.data(), buf.size() * sizeof(T));
+            if (locked_) {
+                unlock_pages(buf.data(), buf.size() * sizeof(T));
+            }
+            buf = std::move(other.buf);
+            locked_ = other.locked_;
+            other.locked_ = false;
+        }
+        return *this;
+    }
+
     /// \brief Zeroize contents on destruction.
-    ~secure_buffer() { secure_zero(buf.data(), buf.size() * sizeof(T)); }
+    ~secure_buffer() {
+        secure_zero(buf.data(), buf.size() * sizeof(T));
+        if (locked_) {
+            unlock_pages(buf.data(), buf.size() * sizeof(T));
+        }
+    }
 
     T* data() { return buf.data(); }
     const T* data() const { return buf.data(); }
@@ -72,6 +135,7 @@ struct secure_buffer {
 
 private:
     std::vector<T> buf;
+    bool locked_{};
 };
 
 } // namespace hmac_cpp