wrapper trieu (done)

 avatar
Huongnm
c_cpp
a month ago
3.6 kB
6
Indexable
#include <tee_internal_api.h>
#include <stdio.h>
#include <stdlib.h>
#include <fcntl.h>

extern void *__real_malloc(size_t size);
extern void *__real_calloc(size_t nmemb, size_t size);
extern void *__real_realloc(void *ptr, size_t size);
extern void __real_free(void *ptr);
extern void *__real_mmap(void *addr, size_t len, int prot, int flags, int fd, off_t offset);
extern int __real_munmap(void *addr, size_t length);
extern void *__real_TEE_Malloc(size_t size, uint32_t hint);
extern void *__real_TEE_Realloc(void *buffer, size_t newSize);
extern void __real_TEE_Free(void *buffer);

static size_t total_mem, alloc_cnt, free_cnt;

#define dbg_printf(format, args...)

void print_memory()
{
	dbg_printf("%s [%zu][%zu][%zu][%zu]\n", __func__, total_mem, alloc_cnt, free_cnt, alloc_cnt - free_cnt);
}

static void print_footprint(void *ptr, int type)
{
	size_t *size = (size_t *)((unsigned long)ptr - sizeof(size_t));
	if (type)
		size = (size_t *)((unsigned long)size - sizeof(long));

	dbg_printf("%s [%p][%zu]\n", __func__, ptr, (*size) & ~0b11);
}

void *__wrap_malloc(size_t size)
{
	dbg_printf("%s [%zu]\n", __func__, size);
	void *ptr =  __real_malloc(size);
	if (ptr) {
		print_footprint(ptr, 0);
		total_mem += size;
		alloc_cnt++;
	}
	return ptr;
}

void *__wrap_calloc(size_t nmemb, size_t size)
{
	dbg_printf("%s [%zu][%zu]\n", __func__, nmemb, size);
	void *ptr =  __real_calloc(nmemb, size);
	if (ptr) {
		print_footprint(ptr, 0);
		total_mem += size;
		alloc_cnt++;
	}
	return ptr;
}

void *__wrap_realloc(void *ptr, size_t size)
{
	dbg_printf("%s [%p][%zu]\n", __func__, ptr, size);
	size_t size_before = *(size_t *)((unsigned long)ptr - sizeof(size_t));
	print_footprint(ptr, 0);
	void *new_ptr = __real_realloc(ptr, size);
	if (new_ptr) {
		print_footprint(new_ptr, 0);
		total_mem -= size_before & ~0b11;
		total_mem += size;
	}
	return new_ptr;
}

void __wrap_free(void *ptr)
{
	dbg_printf("%s [%p]\n", __func__, ptr);
	size_t size = *(size_t *)((unsigned long)ptr - sizeof(size_t));
	total_mem -= size & ~0b11;
	free_cnt++;
	print_footprint(ptr, 0);
	__real_free(ptr);
}

void *__wrap_mmap(void *addr, size_t len, int prot, int flags, int fd, off_t offset)
{
	dbg_printf("%s [%p][%d][%d][%d][%d][%d]\n", __func__, addr, len, prot, flags, fd, offset);
	return __real_mmap(addr, len, prot, flags, fd, offset);
}

int __wrap_munmap(void *addr, size_t length)
{
	dbg_printf("%s [%p][%zu]\n", __func__, addr, length);
	return __real_munmap(addr, length);
}

void *__wrap_TEE_Malloc(size_t size, uint32_t hint)
{
	dbg_printf("%s [%zu][%u]\n", __func__, size, hint);
	void *ptr = __real_TEE_Malloc(size, hint);
	if (ptr) {
		print_footprint(ptr, 1);
		total_mem += size;
		alloc_cnt++;
	}
	return ptr;
}

void *__wrap_TEE_Realloc(void *buffer, size_t newSize)
{
	dbg_printf("%s [%p][%zu]\n", __func__, buffer, newSize);
	size_t size_before = *(size_t *)((unsigned long)buffer - sizeof(size_t) - sizeof(long));
	print_footprint(buffer, 1);
	void *new_ptr = __real_TEE_Realloc(buffer, newSize);
	if (new_ptr) {
		print_footprint(new_ptr, 1);
		dbg_printf("%s [%zu]\n", __func__, size_before);
		total_mem -= size_before;
		total_mem += newSize;
	}
	return new_ptr;
}

void __wrap_TEE_Free(void *buffer)
{
	dbg_printf("%s [%p]\n", __func__, buffer);
	size_t size = *(size_t *)((unsigned long)buffer - sizeof(size_t) - sizeof(long));
	total_mem -= size;
	free_cnt++;
	print_footprint(buffer, 1);
	__real_TEE_Free(buffer);
}
Editor is loading...
Leave a Comment