Skip to content

Commit 3c30a17

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 - Streamlined best-fit allocation algorithm
1 parent ea1b6ab commit 3c30a17

File tree

1 file changed

+33
-31
lines changed

1 file changed

+33
-31
lines changed

lib/c.c

Lines changed: 33 additions & 31 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;
@@ -611,6 +615,9 @@ void *malloc(int size)
611615
int flags = 34; /* MAP_PRIVATE (0x02) | MAP_ANONYMOUS (0x20) */
612616
int prot = 3; /* PROT_READ (0x01) | PROT_WRITE (0x02) */
613617

618+
/* Align size to MIN_ALIGNMENT */
619+
size = ALIGN_UP(size, MIN_ALIGNMENT);
620+
614621
if (!__alloc_head) {
615622
chunk_t *tmp =
616623
__syscall(__syscall_mmap2, NULL, __align_up(sizeof(chunk_t)), prot,
@@ -632,45 +639,31 @@ void *malloc(int size)
632639
__freelist_head->size = -1;
633640
}
634641

635-
/* to search the best chunk */
642+
/* Search for the best fit chunk in the free list */
636643
chunk_t *best_fit_chunk = NULL;
637644
chunk_t *allocated;
645+
int best_size = 0;
638646

639647
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;
648+
allocated = NULL;
644649
} else {
645-
/* record the size of the chunk */
646-
int bsize = 0;
647-
648650
for (chunk_t *fh = __freelist_head; fh->next; fh = fh->next) {
649651
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 */
652+
if (fh_size >= size && (!best_fit_chunk || fh_size < best_size)) {
652653
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;
654+
best_size = fh_size;
659655
}
660656
}
661657

662-
/* a suitable chunk has been found */
663658
if (best_fit_chunk) {
664-
/* remove the chunk from the freelist */
659+
/* Remove from freelist */
665660
if (best_fit_chunk->prev) {
666-
chunk_t *tmp = best_fit_chunk->prev;
667-
tmp->next = best_fit_chunk->next;
668-
} else
661+
best_fit_chunk->prev->next = best_fit_chunk->next;
662+
} else {
669663
__freelist_head = best_fit_chunk->next;
670-
664+
}
671665
if (best_fit_chunk->next) {
672-
chunk_t *tmp = best_fit_chunk->next;
673-
tmp->prev = best_fit_chunk->prev;
666+
best_fit_chunk->next->prev = best_fit_chunk->prev;
674667
}
675668
}
676669
allocated = best_fit_chunk;
@@ -683,19 +676,26 @@ void *malloc(int size)
683676
allocated->size = __align_up(sizeof(chunk_t) + size);
684677
}
685678

679+
/* Add to allocation list */
686680
__alloc_tail->next = allocated;
687681
allocated->prev = __alloc_tail;
688-
689682
__alloc_tail = allocated;
690683
__alloc_tail->next = NULL;
691684
__alloc_tail->size = allocated->size;
692685
chunk_clear_freed(__alloc_tail);
686+
693687
void *ptr = __alloc_tail + 1;
694688
return ptr;
695689
}
696690

697691
void *calloc(int n, int size)
698692
{
693+
/* Check for overflow before multiplication */
694+
if (!n || !size)
695+
return NULL;
696+
if (size != 0 && n > INT_MAX / size)
697+
return NULL; /* Overflow protection */
698+
699699
int total = n * size;
700700
char *p = malloc(total);
701701

@@ -722,7 +722,7 @@ int __free_all(void)
722722
int size;
723723

724724
/* release freelist */
725-
while (cur->next) {
725+
while (cur && cur->next) {
726726
rel = cur;
727727
cur = cur->next;
728728
rel->next = NULL;
@@ -731,7 +731,7 @@ int __free_all(void)
731731
__rfree(rel, size);
732732
}
733733

734-
if (__alloc_head->next) {
734+
if (__alloc_head && __alloc_head->next) {
735735
cur = __alloc_head->next;
736736
/* release chunks which not be free */
737737
while (cur) {
@@ -758,17 +758,18 @@ void free(void *ptr)
758758
abort();
759759
}
760760

761-
chunk_t *prev;
761+
chunk_t *prev = NULL;
762762
if (cur->prev) {
763763
prev = cur->prev;
764764
prev->next = cur->next;
765-
} else
765+
} else {
766766
__alloc_head = cur->next;
767+
}
767768

768769
if (cur->next) {
769770
chunk_t *next = cur->next;
770771
next->prev = cur->prev;
771-
} else {
772+
} else if (prev) {
772773
prev->next = NULL;
773774
__alloc_tail = prev;
774775
}
@@ -777,6 +778,7 @@ void free(void *ptr)
777778
cur->next = __freelist_head;
778779
cur->prev = NULL;
779780
chunk_set_freed(cur);
780-
__freelist_head->prev = cur;
781+
if (__freelist_head)
782+
__freelist_head->prev = cur;
781783
__freelist_head = cur;
782784
}

0 commit comments

Comments
 (0)