Skip to content

Commit 882ae93

Browse files
t00sht00sh
committed
Working on 'heap' command...
1 parent f78376e commit 882ae93

File tree

3 files changed

+149
-100
lines changed

3 files changed

+149
-100
lines changed

‎Makefile

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,5 @@
1-
.PHONY: clean release test
21
include Makefile.inc
32

4-
53
CC = gcc
64

75
CFLAGS = -O2 -Wall -Wextra -Wwrite-strings -Wstrict-prototypes -Wuninitialized
@@ -34,6 +32,12 @@ EXE_STATIC = $(EXE)-static
3432

3533
LIB_HEAP = libheap-$(ARCH).so
3634

35+
36+
37+
38+
39+
.PHONY: clean release test $(LIB_HEAP)
40+
3741
all: $(EXE) $(LIB_HEAP)
3842
static: $(EXE_STATIC)
3943

‎lib/heap/src/heap.c

Lines changed: 126 additions & 93 deletions
Original file line numberDiff line numberDiff line change
@@ -27,162 +27,189 @@
2727
#include <stdio.h>
2828
#include <string.h>
2929

30-
void* (*heap_fun_malloc)(size_t);
31-
void* (*heap_fun_realloc)(void*, size_t);
32-
void* (*heap_fun_calloc)(size_t, size_t);
33-
void* (*heap_fun_free)(void*);
34-
35-
struct malloc_chunk_header {
30+
typedef struct {
3631
size_t prev_size;
3732
size_t size;
38-
};
33+
}libheap_chunk_s;
3934

40-
static int libc_initialized = 0;
41-
static int initialize = 0;
4235

43-
#define TMP_HEAP_SIZE 0x1000
36+
void* (*libheap_malloc)(size_t);
37+
void* (*libheap_realloc)(void*, size_t);
38+
void* (*libheap_calloc)(size_t, size_t);
39+
void* (*libheap_free)(void*);
4440

45-
static u8 tmp_heap[TMP_HEAP_SIZE];
46-
static size_t tmp_heap_allocated = 0;
41+
static int libheap_initialized = 0;
42+
static int libheap_under_initialization = 0;
4743

48-
#define PREV_INUSE 0x1
49-
#define IS_MMAPED 0x2
50-
#define NON_MAIN_ARENA 0x4
44+
#define LIBHEAP_HEAP_SIZE 0x1000
5145

52-
#define CHUNK_CHECK_FLAG(c,f) (c->size & f)
53-
#define CHUNK_SIZE(c) (c->size & ~(PREV_INUSE|IS_MMAPED|NON_MAIN_ARENA))
54-
#define GET_CHUNK(ptr) ((struct malloc_chunk_header*)(ptr - 2*sizeof(size_t)))
46+
static u8 libheap_heap[LIBHEAP_HEAP_SIZE];
47+
static size_t libheap_heap_size = 0;
5548

56-
#define DUMP_HEAP(...) do {dump_heap();printf(__VA_ARGS__);}while(0)
49+
#define LIBHEAP_PREV_INUSE 0x1
50+
#define LIBHEAP_IS_MMAPED 0x2
51+
#define LIBHEAP_NON_MAIN_ARENA 0x4
5752

53+
#define LIBHEAP_CHUNK_FLAG(c,f) (c->size & f)
54+
#define LIBHEAP_CHUNK_SIZE(c) (c->size & ~(LIBHEAP_PREV_INUSE|LIBHEAP_IS_MMAPED|LIBHEAP_NON_MAIN_ARENA))
55+
#define LIBHEAP_NEXT_CHUNK(c) ((void*)(((u8*)c)+LIBHEAP_CHUNK_SIZE(c)))
56+
#define LIBHEAP_GET_CHUNK(ptr) ((libheap_chunk_s*)(ptr - 2*sizeof(size_t)))
57+
#define LIBHEAP_USER_ADDR(c) (((u8*)c) + 2*sizeof(size_t))
5858

59-
static struct malloc_chunk_header *first_chunk = NULL;
60-
static struct malloc_chunk_header *last_chunk = NULL;
59+
#define LIBHEAP_DUMP(...) do { \
60+
R_UTILS_PRINT_RED_BG_BLACK(libheap_options_color, __VA_ARGS__); \
61+
libheap_dump(); \
62+
}while(0)
6163

62-
void dump_chunk(struct malloc_chunk_header *chunk) {
63-
printf("+++++++++++++++++++++++++++++++++++++++++++++\n");
64-
printf("+ ADDR: %-36p+\n", chunk);
65-
printf("+ USER_ADDR: %-31p+\n", ((u8*)(chunk)) + 2*sizeof(size_t));
64+
#define LIBHEAP_DUMP_FIELD(f,...) do { \
65+
R_UTILS_PRINT_YELLOW_BG_BLACK(libheap_options_color,f); \
66+
R_UTILS_PRINT_GREEN_BG_BLACK(libheap_options_color,__VA_ARGS__); \
67+
}while(0)
6668

67-
if(!CHUNK_CHECK_FLAG(chunk, IS_MMAPED) &&
68-
!CHUNK_CHECK_FLAG(chunk, PREV_INUSE))
69-
printf("+ PREV_SIZE: %-31" SIZE_T_FMT_X "+\n", chunk->prev_size);
70-
else
71-
printf("+ PREV_SIZE: UNUSED +\n");
72-
printf("+ SIZE: %-36" SIZE_T_FMT_X "+\n", CHUNK_SIZE(chunk));
73-
printf("+ FLAGS: %c%c +\n",
74-
CHUNK_CHECK_FLAG(chunk, IS_MMAPED) ? 'M' : '-',
75-
CHUNK_CHECK_FLAG(chunk, PREV_INUSE) ? 'P' : '-');
69+
static libheap_chunk_s *libheap_first_chunk = NULL;
70+
static libheap_chunk_s *libheap_last_chunk = NULL;
7671

77-
printf("+++++++++++++++++++++++++++++++++++++++++++++\n");
78-
}
72+
static int libheap_options_color = 1;
7973

80-
void dump_heap(void) {
81-
struct malloc_chunk_header *chunk;
74+
static void libheap_dump_chunk(libheap_chunk_s *chunk) {
75+
LIBHEAP_DUMP_FIELD("addr: ", "0x%p, ", chunk);
76+
LIBHEAP_DUMP_FIELD("usr_addr: ", "0x%p, ", LIBHEAP_USER_ADDR(chunk));
77+
78+
if(!LIBHEAP_CHUNK_FLAG(chunk, LIBHEAP_IS_MMAPED) &&
79+
!LIBHEAP_CHUNK_FLAG(chunk, LIBHEAP_PREV_INUSE))
80+
LIBHEAP_DUMP_FIELD("prev_size: ", "0x%"SIZE_T_FMT_X", ", chunk->prev_size);
81+
LIBHEAP_DUMP_FIELD("size: ", "0x%"SIZE_T_FMT_X", ", LIBHEAP_CHUNK_SIZE(chunk));
82+
LIBHEAP_DUMP_FIELD("flags: ", "%c%c%c\n",
83+
LIBHEAP_CHUNK_FLAG(chunk, LIBHEAP_IS_MMAPED) ? 'M' : '-',
84+
LIBHEAP_CHUNK_FLAG(chunk, LIBHEAP_PREV_INUSE) ? 'P' : '-',
85+
LIBHEAP_CHUNK_FLAG(chunk, LIBHEAP_NON_MAIN_ARENA) ? 'A' : '-');
86+
}
8287

83-
printf("\n\n######################\n");
84-
printf("# HEAP #\n");
85-
printf("######################\n");
88+
static void libheap_dump(void) {
89+
libheap_chunk_s *chunk;
8690

87-
chunk = first_chunk;
91+
chunk = libheap_first_chunk;
8892

8993
while(chunk != NULL) {
90-
dump_chunk(chunk);
94+
libheap_dump_chunk(chunk);
95+
chunk = LIBHEAP_NEXT_CHUNK(chunk);
9196

92-
if(((u8*)(chunk))+CHUNK_SIZE(chunk) > (u8*)last_chunk) {
97+
if(chunk > libheap_last_chunk) {
9398
chunk = NULL;
94-
} else {
95-
chunk = (struct malloc_chunk_header*)(((u8*)(chunk))+CHUNK_SIZE(chunk));
9699
}
97100
}
98101
}
99102

100-
void update_chunk(u8 *ptr) {
101-
if(!CHUNK_CHECK_FLAG(GET_CHUNK(ptr), IS_MMAPED)) {
102-
if(first_chunk == NULL) {
103-
first_chunk = last_chunk = GET_CHUNK(ptr);
104-
} else {
105-
if(GET_CHUNK(ptr) > last_chunk)
106-
last_chunk = GET_CHUNK(ptr);
107-
}
103+
static void libheap_update_chunk(u8 *ptr) {
104+
if(libheap_first_chunk == NULL) {
105+
libheap_first_chunk = libheap_last_chunk = LIBHEAP_GET_CHUNK(ptr);
106+
} else {
107+
if(LIBHEAP_GET_CHUNK(ptr) > libheap_last_chunk)
108+
libheap_last_chunk = LIBHEAP_GET_CHUNK(ptr);
109+
}
110+
}
111+
112+
static void libheap_get_options(void) {
113+
char *env;
114+
115+
if((env = getenv("LIBHEAP_COLOR")) == NULL) {
116+
R_UTILS_ERR("Can't get LIBHEAP_COLOR environment variable");
108117
}
118+
libheap_options_color = atoi(env);
119+
109120
}
110121

111-
void initialize_libc(void) {
122+
static void libheap_initialize(void) {
123+
124+
libheap_get_options();
112125
dlerror();
113126

114-
if((heap_fun_malloc = dlsym(RTLD_NEXT, "malloc")) == NULL) {
115-
fprintf(stderr, "Can't resolve malloc: %s\n", dlerror());
127+
if((libheap_malloc = dlsym(RTLD_NEXT, "malloc")) == NULL) {
128+
fprintf(stderr, "[-] Can't resolve malloc: %s\n", dlerror());
116129
exit(EXIT_FAILURE);
117130
}
118131

119132
dlerror();
120133

121-
if((heap_fun_realloc = dlsym(RTLD_NEXT, "realloc")) == NULL) {
122-
fprintf(stderr, "Can't resolve realloc: %s\n", dlerror());
134+
135+
136+
if((libheap_realloc = dlsym(RTLD_NEXT, "realloc")) == NULL) {
137+
fprintf(stderr, "[-] Can't resolve realloc: %s\n", dlerror());
123138
exit(EXIT_FAILURE);
124139
}
125140

126141
dlerror();
127142

128-
if((heap_fun_calloc = dlsym(RTLD_NEXT, "calloc")) == NULL) {
129-
fprintf(stderr, "Can't resolve calloc: %s\n", dlerror());
143+
if((libheap_calloc = dlsym(RTLD_NEXT, "calloc")) == NULL) {
144+
fprintf(stderr, "[-] Can't resolve calloc: %s\n", dlerror());
130145
exit(EXIT_FAILURE);
131146
}
132147

133148
dlerror();
134149

135-
if((heap_fun_free = dlsym(RTLD_NEXT, "free")) == NULL) {
136-
fprintf(stderr, "Can't resolve free: %s\n", dlerror());
150+
if((libheap_free = dlsym(RTLD_NEXT, "free")) == NULL) {
151+
fprintf(stderr, "[-] Can't resolve free: %s\n", dlerror());
137152
exit(EXIT_FAILURE);
138153
}
139154

140-
libc_initialized = 1;
155+
libheap_initialized = 1;
141156
}
142157

158+
159+
160+
161+
162+
/**************************************************************************************/
163+
/* Modified libc functions : malloc, calloc, realloc, free */
164+
/**************************************************************************************/
165+
143166
void* malloc(size_t s) {
144167
void *p;
145168

146-
if(!libc_initialized) {
147-
if(!initialize) {
148-
initialize = 1;
149-
initialize_libc();
150-
initialize = 0;
151-
p = heap_fun_malloc(s);
169+
if(!libheap_initialized) {
170+
if(!libheap_under_initialization) {
171+
libheap_under_initialization = 1;
172+
libheap_initialize();
173+
libheap_under_initialization = 0;
174+
p = libheap_malloc(s);
152175
} else {
153-
if(s > TMP_HEAP_SIZE || s+tmp_heap_allocated > TMP_HEAP_SIZE) {
154-
fprintf(stderr, "Temporary heap too small for initialization !\n");
176+
if(s > LIBHEAP_HEAP_SIZE || s+libheap_heap_size > LIBHEAP_HEAP_SIZE) {
177+
fprintf(stderr, "[-] Temporary heap too small for initialization !\n");
155178
exit(EXIT_FAILURE);
156179
}
157180

158-
p = tmp_heap + tmp_heap_allocated;
159-
tmp_heap_allocated += s;
181+
p = libheap_heap + libheap_heap_size;
182+
libheap_heap_size += s;
160183
}
161184
} else {
162-
p = heap_fun_malloc(s);
185+
p = libheap_malloc(s);
163186
}
164187

165-
if(!initialize) {
166-
update_chunk(p);
167-
DUMP_HEAP("malloc(%" SIZE_T_FMT_X ") = %p\n", s, p);
188+
if(!libheap_under_initialization) {
189+
if(!LIBHEAP_CHUNK_FLAG(LIBHEAP_GET_CHUNK(p), LIBHEAP_IS_MMAPED)) {
190+
libheap_update_chunk(p);
191+
LIBHEAP_DUMP("malloc(%" SIZE_T_FMT_X ") = %p\n", s, p);
192+
}
168193
}
169194
return p;
170195
}
171196

172197
void* realloc(void *ptr, size_t s) {
173198
void *p;
174199

175-
if(!libc_initialized) {
200+
if(!libheap_initialized) {
176201
p = malloc(s);
177202
if(p && ptr)
178203
memmove(p, ptr, s);
179204
} else {
180-
p = heap_fun_realloc(ptr, s);
205+
p = libheap_realloc(ptr, s);
181206
}
182207

183-
if(!initialize) {
184-
update_chunk(p);
185-
DUMP_HEAP("realloc(%p,%" SIZE_T_FMT_X ") = %p\n", ptr, s, p);
208+
if(!libheap_under_initialization) {
209+
if(!LIBHEAP_CHUNK_FLAG(LIBHEAP_GET_CHUNK(p), LIBHEAP_IS_MMAPED)) {
210+
libheap_update_chunk(p);
211+
LIBHEAP_DUMP("realloc(%p,%" SIZE_T_FMT_X ") = %p\n", ptr, s, p);
212+
}
186213
}
187214

188215
return p;
@@ -191,29 +218,35 @@ void* realloc(void *ptr, size_t s) {
191218
void* calloc(size_t nmemb, size_t size) {
192219
void *p;
193220

194-
if(!libc_initialized) {
221+
if(!libheap_initialized) {
195222
p = malloc(nmemb*size);
196223
if(p)
197224
memset(p, 0, nmemb*size);
198225
} else {
199-
p = heap_fun_calloc(nmemb, size);
226+
p = libheap_calloc(nmemb, size);
200227
}
201228

202-
if(!initialize) {
203-
update_chunk(p);
204-
DUMP_HEAP("calloc(%" SIZE_T_FMT_X ",%" SIZE_T_FMT_X ") = %p\n", nmemb, size, p);
229+
if(!libheap_under_initialization) {
230+
if(!LIBHEAP_CHUNK_FLAG(LIBHEAP_GET_CHUNK(p), LIBHEAP_IS_MMAPED)) {
231+
libheap_update_chunk(p);
232+
LIBHEAP_DUMP("calloc(%" SIZE_T_FMT_X ",%" SIZE_T_FMT_X ") = %p\n", nmemb, size, p);
233+
}
205234
}
206235
return p;
207236
}
208237

209238
void free(void *ptr) {
210239

211-
if(initialize)
240+
if(libheap_under_initialization)
212241
return;
213242

214-
if(!libc_initialized)
215-
initialize_libc();
243+
if(!libheap_initialized)
244+
libheap_initialize();
216245

217-
heap_fun_free(ptr);
218-
DUMP_HEAP("free(%p)\n", ptr);
246+
if(ptr) {
247+
if(!LIBHEAP_CHUNK_FLAG(LIBHEAP_GET_CHUNK(ptr), LIBHEAP_IS_MMAPED)) {
248+
libheap_free(ptr);
249+
LIBHEAP_DUMP("free(%p)\n", ptr);
250+
}
251+
}
219252
}

0 commit comments

Comments
 (0)