Better error handling in array allocation.

socketry · Oct 16, 2024 · 6697370 · 6697370
1 parent 92a7db1
commit 6697370
Show file tree

Hide file tree

Showing 4 changed files with 46 additions and 10 deletions.
diff --git a/ext/io/event/selector/array.h b/ext/io/event/selector/array.h
@@ -8,6 +8,9 @@
 #include <errno.h>
 #include <assert.h>
 
+const size_t IO_EVENT_ARRAY_MAXIMUM_COUNT = SIZE_MAX / sizeof(void*);
+const size_t IO_EVENT_ARRAY_DEFAULT_COUNT = 128;
+
 struct IO_Event_Array {
 	// The array of pointers to elements:
 	void **base;
@@ -25,20 +28,27 @@ struct IO_Event_Array {
 	void (*element_free)(void*);
 };
 
-inline static void IO_Event_Array_allocate(struct IO_Event_Array *array, size_t count, size_t element_size)
+inline static int IO_Event_Array_allocate(struct IO_Event_Array *array, size_t count, size_t element_size)
 {
+	array->limit = 0;
+	array->element_size = element_size;
+
 	if (count) {
 		array->base = (void**)calloc(count, sizeof(void*));
-		assert(array->base);
+
+		if (array->base == NULL) {
+			return -1;
+		}
 
 		array->count = count;
+
+		return 1;
 	} else {
 		array->base = NULL;
 		array->count = 0;
+
+		return 0;
 	}
-
-	array->limit = 0;
-	array->element_size = element_size;
 }
 
 inline static size_t IO_Event_Array_memory_size(const struct IO_Event_Array *array)
@@ -69,12 +79,27 @@ inline static void IO_Event_Array_free(struct IO_Event_Array *array)
 inline static int IO_Event_Array_resize(struct IO_Event_Array *array, size_t count)
 {
 	if (count <= array->count) {
+		// Already big enough:
 		return 0;
 	}
 
-	// Compute the next multiple (ideally a power of 2):
+	if (count > IO_EVENT_ARRAY_MAXIMUM_COUNT) {
+		errno = ENOMEM;
+		return -1;
+	}
+
 	size_t new_count = array->count;
-	while (new_count < count) {
+
+	// If the array is empty, we need to set the initial size:
+	if (new_count == 0) new_count = IO_EVENT_ARRAY_DEFAULT_COUNT;
+	else while (new_count < count) {
+		// Ensure we don't overflow:
+		if (new_count > (IO_EVENT_ARRAY_MAXIMUM_COUNT / 2)) {
+			new_count = IO_EVENT_ARRAY_MAXIMUM_COUNT;
+			break;
+		}
+
+		// Compute the next multiple (ideally a power of 2):
 		new_count *= 2;
 	}
 
@@ -90,6 +115,7 @@ inline static int IO_Event_Array_resize(struct IO_Event_Array *array, size_t cou
 	array->base = (void**)new_base;
 	array->count = new_count;
 
+	// Resizing sucessful:
 	return 1;
 }
 

diff --git a/ext/io/event/selector/epoll.c b/ext/io/event/selector/epoll.c
@@ -337,7 +337,10 @@ VALUE IO_Event_Selector_EPoll_allocate(VALUE self) {
 
 	selector->descriptors.element_initialize = IO_Event_Selector_EPoll_Descriptor_initialize;
 	selector->descriptors.element_free = IO_Event_Selector_EPoll_Descriptor_free;
-	IO_Event_Array_allocate(&selector->descriptors, 1024, sizeof(struct IO_Event_Selector_EPoll_Descriptor));
+	int result = IO_Event_Array_allocate(&selector->descriptors, IO_EVENT_ARRAY_DEFAULT_COUNT, sizeof(struct IO_Event_Selector_EPoll_Descriptor));
+	if (result < 0) {
+		rb_sys_fail("IO_Event_Selector_EPoll_allocate:IO_Event_Array_allocate");
+	}
 
 	return instance;
 }

diff --git a/ext/io/event/selector/kqueue.c b/ext/io/event/selector/kqueue.c
@@ -311,7 +311,11 @@ VALUE IO_Event_Selector_KQueue_allocate(VALUE self) {
 
 	selector->descriptors.element_initialize = IO_Event_Selector_KQueue_Descriptor_initialize;
 	selector->descriptors.element_free = IO_Event_Selector_KQueue_Descriptor_free;
-	IO_Event_Array_allocate(&selector->descriptors, 1024, sizeof(struct IO_Event_Selector_KQueue_Descriptor));
+
+	int result = IO_Event_Array_allocate(&selector->descriptors, IO_EVENT_ARRAY_DEFAULT_COUNT, sizeof(struct IO_Event_Selector_KQueue_Descriptor));
+	if (result < 0) {
+		rb_sys_fail("IO_Event_Selector_KQueue_allocate:IO_Event_Array_allocate");
+	}
 
 	return instance;
 }

diff --git a/ext/io/event/selector/uring.c b/ext/io/event/selector/uring.c
@@ -238,7 +238,10 @@ VALUE IO_Event_Selector_URing_allocate(VALUE self) {
 
 	selector->completions.element_initialize = IO_Event_Selector_URing_Completion_initialize;
 	selector->completions.element_free = IO_Event_Selector_URing_Completion_free;
-	IO_Event_Array_allocate(&selector->completions, 1024, sizeof(struct IO_Event_Selector_URing_Completion));
+	int result = IO_Event_Array_allocate(&selector->completions, IO_EVENT_ARRAY_DEFAULT_COUNT, sizeof(struct IO_Event_Selector_URing_Completion));
+	if (result < 0) {
+		rb_sys_fail("IO_Event_Selector_URing_allocate:IO_Event_Array_allocate");
+	}
 
 	return instance;
 }