How can this code be changed to have no vulnerability to arithmetic overflow?

Question

Computer Systems: a Programmer's Perspective says:

1 /* Illustration of code vulnerability similar to that found in
2 * Sun’s XDR library.
3 */
4 void* copy_elements(void *ele_src[], int ele_cnt, size_t ele_size) {
5 /*
6 * Allocate buffer for ele_cnt objects, each of ele_size bytes
7 * and copy from locations designated by ele_src
8 */
9 void *result = malloc(ele_cnt * ele_size);
10 if (result == NULL)
11 /* malloc failed */
12 return NULL;
13 void *next = result;
14 int i;
15 for (i = 0; i < ele_cnt; i++) {
16 /* Copy object i to destination */
17 memcpy(next, ele_src[i], ele_size);
18 /* Move pointer to next memory region */
19 next += ele_size;
20 }
21 return result;
22 }

The function copy_elements is designed to copy ele_cnt data structures, each consisting of ele_ size bytes into a buffer allocated by the function on line 9. The number of bytes required is computed as ele_cnt * ele_size.

Imagine, however, that a malicious programmer calls this function with ele_cnt being 1,048,577 (2^20 + 1) and ele_size being 4,096 (2^12) with the program compiled for 32 bits. Then the multiplication on line 9 will overflow, causing only 4,096 bytes to be allocated, rather than the 4,294,971,392 bytes required to hold that much data. The loop starting at line 15 will attempt to copy all of those bytes, overrunning the end of the allocated buffer, and therefore corrupting other data structures. This could cause the program to crash or otherwise misbehave.

I was wondering how to change the code to have no vulnerabilities due to arithmetic overflow?

Thanks.

Answer 1

From a mathematical viewpoint you want to check (size_t)-1 < ele_cnt * ele_size. However you can't do that in code because of the overflow. You can instead apply some algebra to avoid the overflow. You also want to first check that both values are positive:

if ((ele_size == 0) || (ele_cnt <= 0) || ((size_t)-1 / ele_size < ele_cnt)) {
    return NULL;
}

Regarding the cast (size_t)-1, because size_t is an unsigned type the conversion is well defined and evaluates to the largest value that can be stored in a size_t.