Skip to content

Commit a7ea84c

Browse files
committed
Harden malloc implementation
The original malloc was development-only with critical flaws: - No alignment - Overflow vulnerability (calloc multiplication) - Null pointer crashes (cleanup functions) - Complex allocation logic New implementation adds: - 8-byte alignment for cache optimization - Bounds checking preventing integer overflow - Null-safe memory management operations
1 parent 396a595 commit a7ea84c

File tree

5 files changed

+38
-37
lines changed

5 files changed

+38
-37
lines changed

lib/c.c

Lines changed: 34 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -578,6 +578,10 @@ int fputc(int c, FILE *stream)
578578
#define CHUNK_GET_SIZE(size) (size & CHUNK_SIZE_SZ_MASK)
579579
#define IS_CHUNK_GET_FREED(size) (size & CHUNK_SIZE_FREED_MASK)
580580

581+
/* Minimum alignment for all memory allocations. */
582+
#define MIN_ALIGNMENT 8
583+
#define ALIGN_UP(val, align) (((val) + (align) - 1) & ~((align) - 1))
584+
581585
typedef struct chunk {
582586
struct chunk *next, *prev;
583587
int size;
@@ -595,8 +599,7 @@ void chunk_clear_freed(chunk_t *chunk)
595599

596600
int __align_up(int size)
597601
{
598-
int mask = PAGESIZE - 1;
599-
return ((size - 1) | mask) + 1;
602+
return ALIGN_UP(size, PAGESIZE);
600603
}
601604

602605
chunk_t *__alloc_head;
@@ -611,6 +614,9 @@ void *malloc(int size)
611614
int flags = 34; /* MAP_PRIVATE (0x02) | MAP_ANONYMOUS (0x20) */
612615
int prot = 3; /* PROT_READ (0x01) | PROT_WRITE (0x02) */
613616

617+
/* Align size to MIN_ALIGNMENT */
618+
size = ALIGN_UP(size, MIN_ALIGNMENT);
619+
614620
if (!__alloc_head) {
615621
chunk_t *tmp =
616622
__syscall(__syscall_mmap2, NULL, __align_up(sizeof(chunk_t)), prot,
@@ -632,45 +638,31 @@ void *malloc(int size)
632638
__freelist_head->size = -1;
633639
}
634640

635-
/* to search the best chunk */
641+
/* Search for the best fit chunk in the free list */
636642
chunk_t *best_fit_chunk = NULL;
637643
chunk_t *allocated;
644+
int best_size = 0;
638645

639646
if (!__freelist_head->next) {
640-
/* If no more chunks in the free chunk list, allocate best_fit_chunk
641-
* as NULL.
642-
*/
643-
allocated = best_fit_chunk;
647+
allocated = NULL;
644648
} else {
645-
/* record the size of the chunk */
646-
int bsize = 0;
647-
648649
for (chunk_t *fh = __freelist_head; fh->next; fh = fh->next) {
649650
int fh_size = CHUNK_GET_SIZE(fh->size);
650-
if (fh_size >= size && !best_fit_chunk) {
651-
/* first time setting fh as best_fit_chunk */
651+
if (fh_size >= size && (!best_fit_chunk || fh_size < best_size)) {
652652
best_fit_chunk = fh;
653-
bsize = fh_size;
654-
} else if ((fh_size >= size) && best_fit_chunk &&
655-
(fh_size < bsize)) {
656-
/* If there is a smaller chunk available, replace it. */
657-
best_fit_chunk = fh;
658-
bsize = fh_size;
653+
best_size = fh_size;
659654
}
660655
}
661656

662-
/* a suitable chunk has been found */
663657
if (best_fit_chunk) {
664-
/* remove the chunk from the freelist */
658+
/* Remove from freelist */
665659
if (best_fit_chunk->prev) {
666-
chunk_t *tmp = best_fit_chunk->prev;
667-
tmp->next = best_fit_chunk->next;
668-
} else
660+
best_fit_chunk->prev->next = best_fit_chunk->next;
661+
} else {
669662
__freelist_head = best_fit_chunk->next;
670-
663+
}
671664
if (best_fit_chunk->next) {
672-
chunk_t *tmp = best_fit_chunk->next;
673-
tmp->prev = best_fit_chunk->prev;
665+
best_fit_chunk->next->prev = best_fit_chunk->prev;
674666
}
675667
}
676668
allocated = best_fit_chunk;
@@ -683,19 +675,26 @@ void *malloc(int size)
683675
allocated->size = __align_up(sizeof(chunk_t) + size);
684676
}
685677

678+
/* Add to allocation list */
686679
__alloc_tail->next = allocated;
687680
allocated->prev = __alloc_tail;
688-
689681
__alloc_tail = allocated;
690682
__alloc_tail->next = NULL;
691683
__alloc_tail->size = allocated->size;
692684
chunk_clear_freed(__alloc_tail);
685+
693686
void *ptr = __alloc_tail + 1;
694687
return ptr;
695688
}
696689

697690
void *calloc(int n, int size)
698691
{
692+
/* Check for overflow before multiplication */
693+
if (!n || !size)
694+
return NULL;
695+
if (size != 0 && n > INT_MAX / size)
696+
return NULL; /* Overflow protection */
697+
699698
int total = n * size;
700699
char *p = malloc(total);
701700

@@ -722,7 +721,7 @@ int __free_all(void)
722721
int size;
723722

724723
/* release freelist */
725-
while (cur->next) {
724+
while (cur && cur->next) {
726725
rel = cur;
727726
cur = cur->next;
728727
rel->next = NULL;
@@ -731,7 +730,7 @@ int __free_all(void)
731730
__rfree(rel, size);
732731
}
733732

734-
if (__alloc_head->next) {
733+
if (__alloc_head && __alloc_head->next) {
735734
cur = __alloc_head->next;
736735
/* release chunks which not be free */
737736
while (cur) {
@@ -758,17 +757,18 @@ void free(void *ptr)
758757
abort();
759758
}
760759

761-
chunk_t *prev;
760+
chunk_t *prev = NULL;
762761
if (cur->prev) {
763762
prev = cur->prev;
764763
prev->next = cur->next;
765-
} else
764+
} else {
766765
__alloc_head = cur->next;
766+
}
767767

768768
if (cur->next) {
769769
chunk_t *next = cur->next;
770770
next->prev = cur->prev;
771-
} else {
771+
} else if (prev) {
772772
prev->next = NULL;
773773
__alloc_tail = prev;
774774
}
@@ -777,6 +777,7 @@ void free(void *ptr)
777777
cur->next = __freelist_head;
778778
cur->prev = NULL;
779779
chunk_set_freed(cur);
780-
__freelist_head->prev = cur;
780+
if (__freelist_head)
781+
__freelist_head->prev = cur;
781782
__freelist_head = cur;
782783
}

tests/snapshots/fib-arm.json

Lines changed: 1 addition & 1 deletion
Large diffs are not rendered by default.

tests/snapshots/fib-riscv.json

Lines changed: 1 addition & 1 deletion
Large diffs are not rendered by default.

tests/snapshots/hello-arm.json

Lines changed: 1 addition & 1 deletion
Large diffs are not rendered by default.

tests/snapshots/hello-riscv.json

Lines changed: 1 addition & 1 deletion
Large diffs are not rendered by default.

0 commit comments

Comments
 (0)