From b638f77b158b9d097658859d2ead521183c39c47 Mon Sep 17 00:00:00 2001
From: m3g4d1v3r <99227953+m3g4d1v3r@users.noreply.github.com>
Date: Sat, 2 Nov 2024 02:59:36 +0200
Subject: [PATCH] add realloc_array function

Adds a new memory reallocation function for arrays which
includes an overflow check on the size multiplication.
---
 include/private/memory.h  |  5 +++++
 src/cache.c               |  2 +-
 src/element.c             |  7 +------
 src/memory.c              | 11 ++++++++++-
 src/target/journald.c     |  2 +-
 test/function/element.cpp | 17 +++++++++++++++++
 6 files changed, 35 insertions(+), 9 deletions(-)

diff --git a/include/private/memory.h b/include/private/memory.h
index 860cbd5fc..1599e274a 100644
--- a/include/private/memory.h
+++ b/include/private/memory.h
@@ -114,6 +114,11 @@ size_t get_paged_size( size_t size );
  *         allocation failure error.
  */
 void *realloc_mem( const void *mem, size_t size );
+
+// A wrapper for alloc_mem that checks for overflow before proper allocation
 void *alloc_array( size_t item_count, size_t item_size );
 
+// A wrapper for realloc_mem that checks for overflow before proper allocation
+void *realloc_array( const void *mem, size_t item_count, size_t item_size );
+
 #endif /* __STUMPLESS_PRIVATE_MEMORY_H */
diff --git a/src/cache.c b/src/cache.c
index d39770b64..166e15b50 100644
--- a/src/cache.c
+++ b/src/cache.c
@@ -67,7 +67,7 @@ add_page( struct cache *c ) {
   char **new_pages;
   char *new_page;
 
-  new_pages = realloc_mem( c->pages, sizeof( char * ) * ( c->page_count + 1 ) );
+  new_pages = realloc_array( c->pages, ( c->page_count + 1 ), sizeof( char * ) );
   if( !new_pages ) {
     return -1;
   }
diff --git a/src/element.c b/src/element.c
index 0abb513eb..37f41a603 100644
--- a/src/element.c
+++ b/src/element.c
@@ -56,18 +56,13 @@ struct stumpless_element *
 stumpless_add_param( struct stumpless_element *element,
                      struct stumpless_param *param ) {
   struct stumpless_param **new_params;
-  size_t old_params_size;
-  size_t new_params_size;
 
   VALIDATE_ARG_NOT_NULL( element );
   VALIDATE_ARG_NOT_NULL( param );
 
   lock_element( element );
 
-  old_params_size = sizeof( param ) * element->param_count;
-  new_params_size = old_params_size + sizeof( param );
-
-  new_params = realloc_mem( element->params, new_params_size );
+  new_params = realloc_array( element->params, element->param_count + 1, sizeof( struct stumpless_param * ) );
   if( !new_params ) {
     unlock_element( element );
     return NULL;
diff --git a/src/memory.c b/src/memory.c
index 51bacf94c..9e984559f 100644
--- a/src/memory.c
+++ b/src/memory.c
@@ -136,6 +136,15 @@ alloc_array( size_t item_count, size_t item_size ) {
   return alloc_mem(item_count * item_size);
 }
 
+void *
+realloc_array( const void *mem, size_t item_count, size_t item_size ) {
+  if (item_count && item_count >= (size_t)-1/item_size) {
+    raise_memory_allocation_failure();
+    return NULL;
+  }
+  return realloc_mem(mem, item_count * item_size);
+}
+
 malloc_func_t stumpless_get_malloc(void) {
     return stumpless_malloc ? stumpless_malloc : NULL;
 }
@@ -146,4 +155,4 @@ free_func_t stumpless_get_free(void) {
 
 realloc_func_t stumpless_get_realloc(void) {
     return stumpless_realloc ? stumpless_realloc : NULL;
-}
\ No newline at end of file
+}
diff --git a/src/target/journald.c b/src/target/journald.c
index dc860c1b9..2f90413ef 100644
--- a/src/target/journald.c
+++ b/src/target/journald.c
@@ -191,7 +191,7 @@ void
 init_fields( size_t field_count ) {
   struct iovec *new_fields;
 
-  new_fields = realloc_mem( fields, sizeof( *fields ) * field_count );
+  new_fields = realloc_array( fields, field_count, sizeof( *fields ) );
   if( !new_fields ) {
     return;
   }
diff --git a/test/function/element.cpp b/test/function/element.cpp
index 4391027e9..28f447cb9 100644
--- a/test/function/element.cpp
+++ b/test/function/element.cpp
@@ -607,6 +607,23 @@ namespace {
     EXPECT_ERROR_ID_EQ( STUMPLESS_INVALID_ENCODING );
   }
 
+  TEST_F( ElementTest, HighParamCount ) {
+    size_t original_param_count;
+    const struct stumpless_element *result;
+
+    original_param_count = stumpless_get_param_count( basic_element );
+    basic_element->param_count = SIZE_MAX - 1;
+
+    result = stumpless_add_new_param( basic_element, "high-param-count-name", "high-param-count-value" );
+    EXPECT_NULL( result );
+
+    EXPECT_ERROR_ID_EQ( STUMPLESS_MEMORY_ALLOCATION_FAILURE );
+
+    basic_element->param_count = original_param_count;
+    EXPECT_EQ( stumpless_get_param_count( basic_element ),
+               original_param_count );
+  }
+
   TEST_F( ElementTest, SetName ) {
     const char *new_name = "awesome-new-name";
     const struct stumpless_element *result;