From d29b5f08b3785319af411486b11948c8611beeef Mon Sep 17 00:00:00 2001
From: Charlie Vigue <charlie.vigue@openvpn.com>
Date: Fri, 22 Nov 2024 12:14:35 +0000
Subject: [PATCH] buffer: Add realign() member function

This function is used to realign the buffer so that the data starts at
a given offest from the beginning of the buffer. This is useful when
a struct is streamed into the buffer but it is not aligned to the
proper memory boundary.

Might also be useful for other use cases.

ConstBufferType::realign() is a const member function that does a
similar thing but since ConstBufferType doesn't own its buffer
it can't grow if this is required. This overload of realign() will
grow the buffer if that is needed, and if not it calls the const
version of realign().

Unit tests included.

Signed-off-by: Charlie Vigue <charlie.vigue@openvpn.com>
---
 openvpn/buffer/buffer.hpp      | 35 ++++++++++++++++----
 test/unittests/test_buffer.cpp | 59 ++++++++++++++++++++++++++++++++++
 2 files changed, 88 insertions(+), 6 deletions(-)

diff --git a/openvpn/buffer/buffer.hpp b/openvpn/buffer/buffer.hpp
index 64f64ac9f..9d8b4a946 100644
--- a/openvpn/buffer/buffer.hpp
+++ b/openvpn/buffer/buffer.hpp
@@ -823,7 +823,7 @@ class BufferType : public ConstBufferType<T>
     /**
      * @brief Default constructor for BufferType.
      */
-    BufferType(){};
+    BufferType() = default;
 
     /**
      * @brief Constructor for BufferType that takes a void pointer, size, and a flag indicating if the buffer is filled.
@@ -953,6 +953,15 @@ class BufferAllocatedType : public BufferType<T>
      */
     void realloc(const size_t newcap);
 
+    /**
+        @brief Realign the buffer with the specified headroom.
+        @param headroom The amount of headroom to reserve in the buffer.
+        @return BufferAllocatedType& A reference to the realigned buffer.
+        @note May reallocate or throw an exception if the reallocation fails.
+        @throws if the buffer is full and the reallocation fails.
+    */
+    BufferAllocatedType &realign(const size_t headroom);
+
     /**
      * @brief Resets the buffer with the specified minimum capacity and flags.
      * @param min_capacity The minimum capacity for the buffer.
@@ -1052,8 +1061,9 @@ class BufferAllocatedType : public BufferType<T>
     /**
      * @brief Reallocates the buffer to the specified new capacity.
      * @param newcap The new capacity for the buffer.
+     * @param new_offset The new offset for the buffer.
      */
-    void realloc_(const size_t newcap);
+    void realloc_(const size_t newcap, size_t new_offset);
 
     /**
      * @brief Frees the data associated with the buffer.
@@ -1709,7 +1719,20 @@ template <typename T>
 void BufferAllocatedType<T>::realloc(const size_t newcap)
 {
     if (newcap > capacity())
-        realloc_(newcap);
+        realloc_(newcap, offset());
+}
+
+template <typename T>
+BufferAllocatedType<T> &BufferAllocatedType<T>::realign(const size_t headroom)
+{
+    if (headroom != offset())
+    {
+        if (headroom + size() > capacity())
+            realloc_(headroom + size(), headroom);
+        else
+            BufferType<T>::realign(headroom);
+    }
+    return *this;
 }
 
 template <typename T>
@@ -1800,16 +1823,16 @@ void BufferAllocatedType<T>::resize(const size_t new_capacity)
     if (newcap > capacity())
     {
         if (flags_ & BufAllocFlags::GROW)
-            realloc_(newcap);
+            realloc_(newcap, offset());
         else
             buffer_full_error(newcap, true);
     }
 }
 
 template <typename T>
-void BufferAllocatedType<T>::realloc_(const size_t newcap)
+void BufferAllocatedType<T>::realloc_(const size_t newcap, size_t new_offset)
 {
-    auto tempBuffer = BufferAllocatedType(offset(), size(), newcap, flags_);
+    auto tempBuffer = BufferAllocatedType(new_offset, size(), newcap, flags_);
     if (size())
         std::memcpy(tempBuffer.data(), c_data(), size() * sizeof(T));
     swap(tempBuffer);
diff --git a/test/unittests/test_buffer.cpp b/test/unittests/test_buffer.cpp
index 0a30db1e7..315774edc 100644
--- a/test/unittests/test_buffer.cpp
+++ b/test/unittests/test_buffer.cpp
@@ -343,3 +343,62 @@ TEST(buffer, alloc_buffer_read1)
 
     EXPECT_EQ(memcmp(raw, data, sizeof(raw)), 0);
 }
+
+TEST(buffer, realign)
+{
+    BufferAllocated buf(64, 0);
+    buf_append_string(buf, "hello world");
+
+    buf.advance(5);
+    EXPECT_EQ(buf.c_data_raw()[0], 'h');
+
+    buf.realign(0);
+
+    EXPECT_EQ(buf[0], ' ');
+    EXPECT_EQ(buf[5], 'd');
+    EXPECT_THROW(buf[6], BufferException);
+    EXPECT_EQ(buf.size(), 6u);
+    EXPECT_EQ(buf.c_data_raw()[0], ' ');
+}
+
+TEST(buffer, realign2)
+{
+    BufferAllocated buf(64, 0);
+    buf_append_string(buf, "hello world");
+
+    EXPECT_EQ(buf.c_data_raw()[0], 'h');
+
+    buf.realign(5);
+
+    EXPECT_EQ(buf.c_data_raw()[5], 'h');
+    EXPECT_EQ(buf[0], 'h');
+    EXPECT_EQ(buf.size(), 11u);
+}
+
+TEST(buffer, realign3)
+{
+    BufferAllocated buf(11, 0);
+    buf_append_string(buf, "hello world");
+
+    EXPECT_EQ(buf.c_data_raw()[0], 'h');
+
+    buf.realign(5);
+
+    EXPECT_EQ(buf.c_data_raw()[5], 'h');
+    EXPECT_EQ(buf[0], 'h');
+    EXPECT_EQ(buf.size(), 11u);
+    EXPECT_EQ(buf.offset(), 5u);
+}
+
+TEST(buffer, realign4)
+{
+    BufferAllocated buf(32, 0);
+    buf.realign(7u);
+    buf_append_string(buf, "hello world");
+    EXPECT_EQ(buf.offset(), 7u);
+    buf.realign(0);
+
+    EXPECT_EQ(buf.c_data_raw()[0], 'h');
+    EXPECT_EQ(buf[0], 'h');
+    EXPECT_EQ(buf.offset(), 0);
+}