From 507a5d8a07f69984df6bd20d8fb519e6336d1078 Mon Sep 17 00:00:00 2001
From: Timothy Flynn <trflynn89@pm.me>
Date: Wed, 27 Dec 2023 10:11:20 -0500
Subject: [PATCH] AK: Add an option to zero-fill ByteBuffer data upon growth

This is to avoid UB in cases where we need to be able to read from the
buffer immediately after resizing it.
---
 AK/ByteBuffer.h             | 16 +++++++++++++---
 Tests/AK/TestByteBuffer.cpp | 17 +++++++++++++++++
 2 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/AK/ByteBuffer.h b/AK/ByteBuffer.h
index 316c113709..d2972714b6 100644
--- a/AK/ByteBuffer.h
+++ b/AK/ByteBuffer.h
@@ -163,9 +163,14 @@ public:
         m_size = 0;
     }
 
-    ALWAYS_INLINE void resize(size_t new_size)
+    enum class ZeroFillNewElements {
+        No,
+        Yes,
+    };
+
+    ALWAYS_INLINE void resize(size_t new_size, ZeroFillNewElements zero_fill_new_elements = ZeroFillNewElements::No)
     {
-        MUST(try_resize(new_size));
+        MUST(try_resize(new_size, zero_fill_new_elements));
     }
 
     void trim(size_t size, bool may_discard_existing_data)
@@ -181,13 +186,18 @@ public:
         MUST(try_ensure_capacity(new_capacity));
     }
 
-    ErrorOr<void> try_resize(size_t new_size)
+    ErrorOr<void> try_resize(size_t new_size, ZeroFillNewElements zero_fill_new_elements = ZeroFillNewElements::No)
     {
         if (new_size <= m_size) {
             trim(new_size, false);
             return {};
         }
         TRY(try_ensure_capacity(new_size));
+
+        if (zero_fill_new_elements == ZeroFillNewElements::Yes) {
+            __builtin_memset(data() + m_size, 0, new_size - m_size);
+        }
+
         m_size = new_size;
         return {};
     }
diff --git a/Tests/AK/TestByteBuffer.cpp b/Tests/AK/TestByteBuffer.cpp
index 7122b3264a..12f9a19a44 100644
--- a/Tests/AK/TestByteBuffer.cpp
+++ b/Tests/AK/TestByteBuffer.cpp
@@ -46,6 +46,23 @@ TEST_CASE(byte_buffer_vector_contains_slow_bytes)
     EXPECT_EQ(vector.contains_slow(c), true);
 }
 
+TEST_CASE(zero_fill_new_elements_on_growth)
+{
+    auto buffer = MUST(ByteBuffer::create_uninitialized(5));
+
+    buffer.span().fill(1);
+    EXPECT_EQ(buffer.span(), (Array<u8, 5> { 1, 1, 1, 1, 1 }));
+
+    buffer.resize(8, ByteBuffer::ZeroFillNewElements::Yes);
+    EXPECT_EQ(buffer.span(), (Array<u8, 8> { 1, 1, 1, 1, 1, 0, 0, 0 }));
+
+    buffer.span().fill(2);
+    EXPECT_EQ(buffer.span(), (Array<u8, 8> { 2, 2, 2, 2, 2, 2, 2, 2 }));
+
+    buffer.resize(10, ByteBuffer::ZeroFillNewElements::Yes);
+    EXPECT_EQ(buffer.span(), (Array<u8, 10> { 2, 2, 2, 2, 2, 2, 2, 2, 0, 0 }));
+}
+
 BENCHMARK_CASE(append)
 {
     ByteBuffer bb;