From 008687562c19ea56086a46f8d97300fa10253764 Mon Sep 17 00:00:00 2001
From: Patryk Kaminski <patryk.kaminski@intel.com>
Date: Fri, 5 Jul 2024 08:16:57 +0200
Subject: [PATCH] Fix integer overflows

---
 src/base_alloc/base_alloc_global.c |  9 +++++++++
 src/provider/provider_os_memory.c  | 17 +++++++++++++++++
 test/common/provider.hpp           |  4 ++++
 3 files changed, 30 insertions(+)

diff --git a/src/base_alloc/base_alloc_global.c b/src/base_alloc/base_alloc_global.c
index 003e43a031..d5c04e2fff 100644
--- a/src/base_alloc/base_alloc_global.c
+++ b/src/base_alloc/base_alloc_global.c
@@ -155,10 +155,19 @@ void *umf_ba_global_aligned_alloc(size_t size, size_t alignment) {
         return NULL;
     }
 
+    if (SIZE_MAX - size < ALLOC_METADATA_SIZE) {
+        LOG_ERR("base_alloc: allocation size (%zu) too large.", size);
+        return NULL;
+    }
+
     // for metadata
     size += ALLOC_METADATA_SIZE;
 
     if (alignment > ALLOC_METADATA_SIZE) {
+        if (SIZE_MAX - size < alignment) {
+            LOG_ERR("base_alloc: allocation size (%zu) too large.", size);
+            return NULL;
+        }
         size += alignment;
     }
 
diff --git a/src/provider/provider_os_memory.c b/src/provider/provider_os_memory.c
index 894a911de2..2cb45d77a5 100644
--- a/src/provider/provider_os_memory.c
+++ b/src/provider/provider_os_memory.c
@@ -360,6 +360,15 @@ validatePartitions(umf_os_memory_provider_params_t *params) {
     return UMF_RESULT_SUCCESS;
 }
 
+static umf_result_t validatePartSize(os_memory_provider_t *provider, umf_os_memory_provider_params_t *params) {
+    size_t page_size;
+    if (ALIGN_UP(params->part_size, os_get_min_page_size(provider, NULL, &page_size)) < params->part_size) {
+        LOG_ERR("partition size (%zu) is too big, cannot align with a page size (%zu)", params->part_size, page_size);
+        return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+    return UMF_RESULT_SUCCESS;
+}
+
 static void free_bitmaps(os_memory_provider_t *provider) {
     for (unsigned i = 0; i < provider->nodeset_len; i++) {
         hwloc_bitmap_free(provider->nodeset[i]);
@@ -443,6 +452,14 @@ static umf_result_t translate_params(umf_os_memory_provider_params_t *in_params,
         return result;
     }
 
+    if(in_params->numa_mode == UMF_NUMA_MODE_INTERLEAVE) {
+        result = validatePartSize(in_params);
+        if (result != UMF_RESULT_SUCCESS) {
+            LOG_ERR("incorrect partition size: %zu", in_params->part_size);
+            return result;
+        }
+    }
+
     int is_dedicated_node_bind = dedicated_node_bind(in_params);
     provider->numa_policy =
         translate_numa_mode(in_params->numa_mode, is_dedicated_node_bind);
diff --git a/test/common/provider.hpp b/test/common/provider.hpp
index 88918da7a7..c2c55f71de 100644
--- a/test/common/provider.hpp
+++ b/test/common/provider.hpp
@@ -106,6 +106,10 @@ 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