Align values up safely

oneapi-src · Sep 4, 2024 · adf2055 · adf2055
1 parent 8c39bb0
commit adf2055
Show file tree

Hide file tree

Showing 8 changed files with 72 additions and 22 deletions.
diff --git a/src/base_alloc/base_alloc.c b/src/base_alloc/base_alloc.c
@@ -134,7 +134,10 @@ static void *ba_os_alloc_annotated(size_t pool_size) {
 }
 
 umf_ba_pool_t *umf_ba_create(size_t size) {
-    size_t chunk_size = ALIGN_UP(size, MEMORY_ALIGNMENT);
+    size_t chunk_size = ALIGN_UP_SAFE(size, MEMORY_ALIGNMENT);
+    if (!chunk_size) {
+        return NULL;
+    }
     size_t mutex_size = ALIGN_UP(util_mutex_get_size(), MEMORY_ALIGNMENT);
 
     size_t metadata_size = sizeof(struct umf_ba_main_pool_meta_t);
@@ -144,7 +147,10 @@ umf_ba_pool_t *umf_ba_create(size_t size) {
         pool_size = MINIMUM_POOL_SIZE;
     }
 
-    pool_size = ALIGN_UP(pool_size, ba_os_get_page_size());
+    pool_size = ALIGN_UP_SAFE(pool_size, ba_os_get_page_size());
+    if (!pool_size) {
+        return NULL;
+    }
 
     umf_ba_pool_t *pool = (umf_ba_pool_t *)ba_os_alloc_annotated(pool_size);
     if (!pool) {

diff --git a/src/base_alloc/base_alloc_global.c b/src/base_alloc/base_alloc_global.c
@@ -96,8 +96,12 @@ static void *add_metadata_and_align(void *ptr, size_t size, size_t alignment) {
     if (alignment <= ALLOC_METADATA_SIZE) {
         user_ptr = (void *)((uintptr_t)ptr + ALLOC_METADATA_SIZE);
     } else {
-        user_ptr =
-            (void *)ALIGN_UP((uintptr_t)ptr + ALLOC_METADATA_SIZE, alignment);
+        user_ptr = (void *)ALIGN_UP_SAFE((uintptr_t)ptr + ALLOC_METADATA_SIZE,
+                                         alignment);
+        if (!user_ptr) {
+            LOG_ERR("base_alloc: ptr alignment overflow");
+            return NULL;
+        }
     }
 
     size_t ptr_offset_from_original = (uintptr_t)user_ptr - (uintptr_t)ptr;

diff --git a/src/base_alloc/base_alloc_linear.c b/src/base_alloc/base_alloc_linear.c
@@ -88,7 +88,11 @@ umf_ba_linear_pool_t *umf_ba_linear_create(size_t pool_size) {
         pool_size = MINIMUM_LINEAR_POOL_SIZE;
     }
 
-    pool_size = ALIGN_UP(pool_size, ba_os_get_page_size());
+    pool_size = ALIGN_UP_SAFE(pool_size, ba_os_get_page_size());
+    if (!pool_size) {
+        LOG_ERR("umf_ba_linear_create: pool_size page alignment overflow");
+        return NULL;
+    }
 
     umf_ba_linear_pool_t *pool = (umf_ba_linear_pool_t *)ba_os_alloc(pool_size);
     if (!pool) {
@@ -122,15 +126,25 @@ void *umf_ba_linear_alloc(umf_ba_linear_pool_t *pool, size_t size) {
     if (size == 0) {
         return NULL;
     }
-    size_t aligned_size = ALIGN_UP(size, MEMORY_ALIGNMENT);
+    size_t aligned_size = ALIGN_UP_SAFE(size, MEMORY_ALIGNMENT);
+    if (!aligned_size) {
+        LOG_ERR("umf_ba_linear_alloc: size alignment overflow");
+        return NULL;
+    }
     util_mutex_lock(&pool->metadata.lock);
     if (pool->metadata.size_left < aligned_size) {
         size_t pool_size = MINIMUM_LINEAR_POOL_SIZE;
         size_t usable_size =
             pool_size - offsetof(umf_ba_next_linear_pool_t, data);
         if (usable_size < aligned_size) {
             pool_size += aligned_size - usable_size;
-            pool_size = ALIGN_UP(pool_size, ba_os_get_page_size());
+            pool_size = ALIGN_UP_SAFE(pool_size, ba_os_get_page_size());
+            if (!pool_size) {
+                util_mutex_unlock(&pool->metadata.lock);
+                LOG_ERR(
+                    "umf_ba_linear_alloc: pool_size page alignment overflow");
+                return NULL;
+            }
         }
 
         assert(pool_size - offsetof(umf_ba_next_linear_pool_t, data) >=

diff --git a/src/provider/provider_os_memory.c b/src/provider/provider_os_memory.c
@@ -366,7 +366,7 @@ static umf_result_t validatePartSize(os_memory_provider_t *provider,
                                      umf_os_memory_provider_params_t *params) {
     size_t page_size;
     os_get_min_page_size(provider, NULL, &page_size);
-    if (ALIGN_UP(params->part_size, page_size) < params->part_size) {
+    if (!ALIGN_UP(params->part_size, page_size)) {
         LOG_ERR("partition size (%zu) is too big, cannot align with a page "
                 "size (%zu)",
                 params->part_size, page_size);
@@ -851,12 +851,20 @@ static membind_t membindFirst(os_memory_provider_t *provider, void *addr,
     membind_t membind;
     memset(&membind, 0, sizeof(membind));
 
-    membind.alloc_size = ALIGN_UP(size, page_size);
+    membind.alloc_size = ALIGN_UP_SAFE(size, page_size);
+    if (membind.alloc_size == 0) {
+        LOG_ERR("size is too big, page align failed");
+        return membind;
+    }
     membind.page_size = page_size;
     membind.addr = addr;
     membind.pages = membind.alloc_size / membind.page_size;
     if (provider->nodeset_len == 1) {
-        membind.bind_size = ALIGN_UP(size, membind.page_size);
+        membind.bind_size = ALIGN_UP_SAFE(size, membind.page_size);
+        if (membind.bind_size == 0) {
+            LOG_ERR("size is too big, page align failed");
+            return membind;
+        }
         membind.bitmap = provider->nodeset[0];
         return membind;
     }
@@ -866,7 +874,12 @@ static membind_t membindFirst(os_memory_provider_t *provider, void *addr,
         size_t s = util_fetch_and_add64(&provider->alloc_sum, size);
         membind.node = (s / provider->part_size) % provider->nodeset_len;
         membind.bitmap = provider->nodeset[membind.node];
-        membind.bind_size = ALIGN_UP(provider->part_size, membind.page_size);
+        membind.bind_size =
+            ALIGN_UP_SAFE(provider->part_size, membind.page_size);
+        if (membind.bind_size == 0) {
+            LOG_ERR("size is too big, page align failed");
+            return membind;
+        }
         if (membind.bind_size > membind.alloc_size) {
             membind.bind_size = membind.alloc_size;
         }
@@ -902,7 +915,12 @@ static membind_t membindNext(os_memory_provider_t *provider,
         membind.node++;
         membind.node %= provider->nodeset_len;
         membind.bitmap = provider->nodeset[membind.node];
-        membind.bind_size = ALIGN_UP(provider->part_size, membind.page_size);
+        membind.bind_size =
+            ALIGN_UP_SAFE(provider->part_size, membind.page_size);
+        if (membind.bind_size == 0) {
+            LOG_ERR("part_size is too big, page align failed");
+            return membind;
+        }
         if (membind.bind_size > membind.alloc_size) {
             membind.bind_size = membind.alloc_size;
         }

diff --git a/src/proxy_lib/proxy_lib.c b/src/proxy_lib/proxy_lib.c
@@ -232,7 +232,7 @@ static inline void *ba_leak_realloc(void *ptr, size_t size, size_t max_size) {
 static inline void *ba_leak_aligned_alloc(size_t alignment, size_t size) {
     ba_leak_init_once();
     void *ptr = umf_ba_linear_alloc(Base_alloc_leak, size + alignment);
-    return (void *)ALIGN_UP((uintptr_t)ptr, alignment);
+    return (void *)ALIGN_UP_SAFE((uintptr_t)ptr, alignment);
 }
 
 static inline int ba_leak_free(void *ptr) {

diff --git a/src/utils/utils_common.h b/src/utils/utils_common.h
@@ -29,8 +29,13 @@ extern "C" {
         expression;                                                            \
     } while (0)
 
-#define ALIGN_UP(value, align) (((value) + (align)-1) & ~((align)-1))
-#define ALIGN_DOWN(value, align) ((value) & ~((align)-1))
+#define ALIGN_UP(value, align) ((value + align - 1) & ~(align - 1))
+#define ALIGN_UP_SAFE(value, alignment)                                        \
+    ((alignment == 0) ? (value)                                                \
+                      : ((value + alignment - 1) < value                       \
+                             ? 0                                               \
+                             : ((value + alignment - 1) & ~(alignment - 1))))
+#define ALIGN_DOWN(value, align) (value & ~(align - 1))
 
 #define VALGRIND_ANNOTATE_NEW_MEMORY(p, s) DO_WHILE_EMPTY
 #define VALGRIND_HG_DRD_DISABLE_CHECKING(p, s) DO_WHILE_EMPTY

diff --git a/test/common/provider.hpp b/test/common/provider.hpp
@@ -104,15 +104,14 @@ struct provider_malloc : public provider_base_t {
             align = 8;
         }
 
-        if (SIZE_MAX - size < align) {
-            return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;
-        }
-
         // aligned_malloc returns a valid pointer despite not meeting the
         // requirement of 'size' being multiple of 'align' even though the
         // documentation says that it has to. AddressSanitizer returns an
         // error because of this issue.
-        size_t aligned_size = ALIGN_UP(size, align);
+        size_t aligned_size = ALIGN_UP_SAFE(size, align);
+        if (!aligned_size) {
+            return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;
+        }
 
 #ifdef _WIN32
         *ptr = _aligned_malloc(aligned_size, align);

diff --git a/test/common/test_helpers.h b/test/common/test_helpers.h
@@ -75,8 +75,12 @@ static inline void UT_OUT(const char *format, ...) {
                       (unsigned long long)(rhs)),                              \
              0)))
 
-#ifndef ALIGN_UP
-#define ALIGN_UP(value, align) (((value) + (align)-1) & ~((align)-1))
+#ifndef ALIGN_UP_SAFE
+#define ALIGN_UP_SAFE(value, alignment)                                        \
+    ((alignment == 0) ? (value)                                                \
+                      : ((value + alignment - 1) < value                       \
+                             ? 0                                               \
+                             : ((value + alignment - 1) & ~(alignment - 1))))
 #endif
 
 int bufferIsFilledWithChar(void *ptr, size_t size, char c);