Support for dlmalloc with KASan

.gitignore

@@ -0,0 +1,4 @@
+*.o
+*.su
+.vscode/
+kasan_test.lds

Makefile

@@ -35,6 +35,7 @@ CFLAGS += -DPRINTF_DISABLE_SUPPORT_FLOAT
 CFLAGS += -DPRINTF_DISABLE_SUPPORT_EXPONENTIAL
 CFLAGS += -DPRINTF_DISABLE_SUPPORT_PTRDIFF_T
 CFLAGS += -DPRINTF_DISABLE_SUPPORT_LONG_LONG
+CFLAGS += -DDLMALLOC_ENABLED
 
 CFLAGS += -Wno-incompatible-library-redeclaration
 
@@ -54,8 +55,10 @@ KASAN_CC_FLAGS += -mllvm -asan-stack=$(KASAN_SANITIZE_STACK)
 KASAN_CC_FLAGS += -mllvm -asan-globals=$(KASAN_SANITIZE_GLOBALS)
 KASAN_CC_FLAGS += -DKASAN_ENABLED
 
-SRCS := kasan.c \
-        heap.c \
+SRCS := heap.c \
+        kasan_common.c \
+		kasan_simple_malloc.c \
+		kasan_dlmalloc.c \
         kasan_test.c \
         sanitized_lib.c \
         rt_utils.c \

README.md

@@ -26,14 +26,13 @@ sudo apt-get install build-essential gcc-multilib llvm clang lld \
                      qemu-system-arm qemu-system-misc qemu-system-x86
 ```
 
-## Project layout
+## Project layout (with dlmalloc)
 
 The project constists of the following components:
 
 * `kasan_test.c` -- main test driver which runs KASan test cases
 * `sanitized_lib.c` -- this module implements the test cases and is built with
                     the KASan instrumentation
-* `kasan.c` -- implementation of runtime routines needed for KASan sanitizer
 * `heap.c` -- simple implementation of heap management routines for testing
               KASan
 * `third_party/printf.c` -- a compact implementation of `printf` function
@@ -48,6 +47,11 @@ The project constists of the following components:
 * `Makefile.arch` -- Makefile fragments whith architecture-specific parameters
                      for building and running the project in the emulator
 
+Files for `malloc`:
+* `kasan_common.c` -- implementation of common runtime routines needed for KASan sanitizer
+* `kasan_simple_malloc.c` -- implementation of KASan runtime routines for simple malloc 
+* `kasan_common.c` -- implementation of KASan runtime routines for dlmalloc 
+
 
 ## Running
 

heap.c

@@ -10,34 +10,46 @@
  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  * GNU General Public License for more details.
  */
+#pragma
 
 #include "common.h"
-#include "kasan.h"
+#include "kasan_common.h"
+#include "kasan_simple_malloc.h"
+#include "kasan_dlmalloc.h"
+#include "heap.h"
+#include "printf.h"
 
 // These symbols are defined in the linker script.
 extern char __heap_start;
 extern char __heap_end;
 
-static void *heap_head;
-static size_t heap_size;
+size_t mem_malloc_size = 0;
+unsigned long mem_malloc_start = 0;
+unsigned long mem_malloc_end = 0;
+unsigned long mem_malloc_brk = 0;
 
 void initialize_heap(void) {
-  heap_head = (void *)&__heap_start;
-  heap_size = (void *)&__heap_end - (void *)&__heap_start;
-}
-
-void *allocate_chunk(unsigned long size) {
-  void *result = heap_head;
-  if (size > heap_size) return NULL;
-
-  size = (size + 7) & (~7UL);
-  heap_head += size;
-  heap_size -= size;
-  return result;
+  mem_malloc_start = (unsigned long)&__heap_start;
+  mem_malloc_end = (unsigned long)&__heap_end;
+  mem_malloc_size = mem_malloc_end - mem_malloc_start;
+  mem_malloc_brk = mem_malloc_start;
+  initialize_kasan();
 }
 
 void free_chunk(void *ptr) { (void)ptr; }
 
-void *malloc(unsigned long size) { return kasan_malloc_hook(size); }
+void *malloc(unsigned long size) { 
+  #ifdef DLMALLOC_ENABLED
+  return kasan_dlmalloc_hook(size); 
+  #else
+  return kasan_malloc_hook(size); 
+  #endif // DLMALLOC_ENABLED
+}
 
-void free(void *ptr) { return kasan_free_hook(ptr); }
+void free(void *ptr) {
+  #ifdef DLMALLOC_ENABLED
+  return kasan_dlfree_hook(ptr); 
+  #else
+  return kasan_free_hook(ptr); 
+  #endif // DLMALLOC_ENABLED
+}

heap.h

@@ -10,15 +10,18 @@
  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  * GNU General Public License for more details.
  */
+#include "common.h"
 
 #ifndef __KASAN_HEAP_H__
 #define __KASAN_HEAP_H__
 
-void initialize_heap(void);
+extern size_t mem_malloc_size;
+extern unsigned long mem_malloc_start;
+extern unsigned long mem_malloc_end;
+extern unsigned long mem_malloc_brk;
 
-void *allocate_chunk(unsigned long size);
+void initialize_heap(void);
 void free_chunk(void *ptr);
-
 void *malloc(unsigned long size);
 void free(void *ptr);
 

kasan.c → kasan_common.c

@@ -12,27 +12,9 @@
  */
 
 #include "common.h"
-#include "heap.h"
+// #include "heap.h"
 #include "printf.h"
-
-#define KASAN_SHADOW_SHIFT 3
-#define KASAN_SHADOW_GRANULE_SIZE (1UL << KASAN_SHADOW_SHIFT)
-#define KASAN_SHADOW_MASK (KASAN_SHADOW_GRANULE_SIZE - 1)
-
-#define ASAN_SHADOW_UNPOISONED_MAGIC 0x00
-#define ASAN_SHADOW_RESERVED_MAGIC 0xff
-#define ASAN_SHADOW_GLOBAL_REDZONE_MAGIC 0xf9
-#define ASAN_SHADOW_HEAP_HEAD_REDZONE_MAGIC 0xfa
-#define ASAN_SHADOW_HEAP_TAIL_REDZONE_MAGIC 0xfb
-#define ASAN_SHADOW_HEAP_FREE_MAGIC 0xfd
-
-#define KASAN_HEAP_HEAD_REDZONE_SIZE 0x20
-#define KASAN_HEAP_TAIL_REDZONE_SIZE 0x20
-
-#define KASAN_MEM_TO_SHADOW(addr) \
-  (((addr) >> KASAN_SHADOW_SHIFT) + KASAN_SHADOW_MAPPING_OFFSET)
-#define KASAN_SHADOW_TO_MEM(shadow) \
-  (((shadow) - KASAN_SHADOW_MAPPING_OFFSET) << KASAN_SHADOW_SHIFT)
+#include "kasan_common.h"
 
 void kasan_bug_report(unsigned long addr, size_t size,
                       unsigned long buggy_shadow_address, uint8_t is_write,
@@ -75,7 +57,7 @@ static inline unsigned long get_poisoned_shadow_address(unsigned long addr,
 }
 
 // Both `address` and `size` must be 8-byte aligned.
-static void poison_shadow(unsigned long address, size_t size, uint8_t value) {
+void poison_shadow(unsigned long address, size_t size, uint8_t value) {
   unsigned long shadow_start, shadow_end;
   size_t shadow_length = 0;
 
@@ -87,7 +69,7 @@ static void poison_shadow(unsigned long address, size_t size, uint8_t value) {
 }
 
 // `address` must be 8-byte aligned
-static void unpoison_shadow(unsigned long address, size_t size) {
+void unpoison_shadow(unsigned long address, size_t size) {
   poison_shadow(address, size & (~KASAN_SHADOW_MASK),
                 ASAN_SHADOW_UNPOISONED_MAGIC);
 
@@ -165,50 +147,6 @@ void *__kasan_memset(void *buf, int c, unsigned int size, unsigned long pc) {
   return memset(buf, c, size);
 }
 
-// Implement KASan heap management hooks.
-
-struct KASAN_HEAP_HEADER {
-  unsigned int aligned_size;
-};
-
-void *kasan_malloc_hook(unsigned int size) {
-  struct KASAN_HEAP_HEADER *kasan_heap_hdr = NULL;
-  unsigned int algined_size = (size + KASAN_SHADOW_MASK) & (~KASAN_SHADOW_MASK);
-  unsigned int total_size = algined_size + KASAN_HEAP_HEAD_REDZONE_SIZE +
-                            KASAN_HEAP_TAIL_REDZONE_SIZE;
-
-  void *ptr = allocate_chunk(total_size);
-  if (ptr == NULL) return NULL;
-
-  kasan_heap_hdr = (struct KASAN_HEAP_HEADER *)ptr;
-  kasan_heap_hdr->aligned_size = algined_size;
-
-  unpoison_shadow((unsigned long)(ptr + KASAN_HEAP_HEAD_REDZONE_SIZE), size);
-  poison_shadow((unsigned long)ptr, KASAN_HEAP_HEAD_REDZONE_SIZE,
-                ASAN_SHADOW_HEAP_HEAD_REDZONE_MAGIC);
-  poison_shadow(
-      (unsigned long)(ptr + KASAN_HEAP_HEAD_REDZONE_SIZE + algined_size),
-      KASAN_HEAP_TAIL_REDZONE_SIZE, ASAN_SHADOW_HEAP_TAIL_REDZONE_MAGIC);
-
-  return ptr + KASAN_HEAP_HEAD_REDZONE_SIZE;
-}
-
-void kasan_free_hook(void *ptr) {
-  struct KASAN_HEAP_HEADER *kasan_heap_hdr = NULL;
-  unsigned int aligned_size = 0;
-
-  if (ptr == NULL) return;
-
-  kasan_heap_hdr =
-      (struct KASAN_HEAP_HEADER *)(ptr - KASAN_HEAP_HEAD_REDZONE_SIZE);
-  aligned_size = kasan_heap_hdr->aligned_size;
-
-  free_chunk(kasan_heap_hdr);
-  poison_shadow((unsigned long)ptr, aligned_size, ASAN_SHADOW_HEAP_FREE_MAGIC);
-
-  return;
-}
-
 // Implement KAsan error reporting routines.
 
 static void kasan_print_16_bytes_no_bug(const char *prefix,

kasan_common.h

@@ -0,0 +1,41 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * This program is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU General Public License
+ * version 2 as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * GNU General Public License for more details.
+ */
+
+#include "common.h"
+#ifndef __KASAN_COMMON__
+#define __KASAN_COMMON__
+
+#define KASAN_SHADOW_SHIFT 3
+#define KASAN_SHADOW_GRANULE_SIZE (1UL << KASAN_SHADOW_SHIFT)
+#define KASAN_SHADOW_MASK (KASAN_SHADOW_GRANULE_SIZE - 1)
+
+#define ASAN_SHADOW_UNPOISONED_MAGIC 0x00
+#define ASAN_SHADOW_RESERVED_MAGIC 0xff
+#define ASAN_SHADOW_GLOBAL_REDZONE_MAGIC 0xf9
+#define ASAN_SHADOW_HEAP_HEAD_REDZONE_MAGIC 0xfa
+#define ASAN_SHADOW_HEAP_TAIL_REDZONE_MAGIC 0xfb
+#define ASAN_SHADOW_HEAP_FREE_MAGIC 0xfd
+
+#define KASAN_HEAP_HEAD_REDZONE_SIZE 0x20
+#define KASAN_HEAP_TAIL_REDZONE_SIZE 0x20
+
+#define KASAN_MEM_TO_SHADOW(addr) \
+    (((addr) >> KASAN_SHADOW_SHIFT) + KASAN_SHADOW_MAPPING_OFFSET)
+#define KASAN_SHADOW_TO_MEM(shadow) \
+    (((shadow) - KASAN_SHADOW_MAPPING_OFFSET) << KASAN_SHADOW_SHIFT)
+
+void initialize_kasan(void);
+void poison_shadow(unsigned long address, size_t size, uint8_t value);
+void unpoison_shadow(unsigned long address, size_t size);
+
+#endif