/* radare - LGPL - Copyright 2017-2025 - javierptd, pancake */

#include "dmh_output.h"

#if R_INCLUDE_BEGIN

#include "./i/heap/jemalloc/jemalloc.h"

#ifndef INCLUDE_HEAP_JEMALLOC_STD_C
#define INCLUDE_HEAP_JEMALLOC_STD_C
#define HEAP32 1
#include "dmh_jemalloc.inc.c"
#undef HEAP32
#endif

#undef GH
#undef GHT
#undef GHT_MAX
#undef PFMTx

// FIXME: It should be detected at runtime, not during the compilation stage
#if HEAP32
#define GH(x) x##_32
#define GHT ut32
#define GHT_MAX UT32_MAX
#define PFMTx PFMT32x
#else
#define GH(x) x##_64
#define GHT ut64
#define GHT_MAX UT64_MAX
#define PFMTx PFMT64x
#endif

#if __linux__
static GHT GH(je_get_va_symbol)(RCore *core, const char *path, const char *sym_name) {
	GHT vaddr = GHT_MAX;
	RBinFileOptions opt = {0};
	r_bin_file_options_init (&opt, -1, 0, 0, false);
	RBinSymbol *s;
	RBinFile *current_bf = r_bin_cur (core->bin);
	if (!r_bin_open (core->bin, path, &opt)) {
		return vaddr;
	}
	RBinFile *libc_bf = r_bin_cur (core->bin);
 	if (!libc_bf) {
 		return vaddr;
 	}
	RVecRBinSymbol *syms = r_bin_get_symbols_vec (core->bin);
	R_VEC_FOREACH (syms, s) {
		const char *sname = r_bin_name_tostring (s->name);
		if (!strcmp (sname, sym_name)) {
 			vaddr = s->vaddr;
 			break;
		}
	}
	r_bin_file_delete (core->bin, libc_bf->id);
 	r_bin_file_set_cur_binfile (core->bin, current_bf);
	return vaddr;
}

static int GH(je_matched)(const char *ptr, const char *str) {
	int ret = strncmp (ptr, str, strlen (str) - 1);
	return !ret;
}
#endif

static bool GH(r_resolve_jemalloc)(RCore *core, char *symname, ut64 *symbol) {
	RListIter *iter;
	RDebugMap *map;
	const char *path = NULL;
	ut64 jemalloc_addr = UT64_MAX;

	if (!core || !core->dbg || !core->dbg->maps) {
		return false;
	}
	r_debug_map_sync (core->dbg);
	r_list_foreach (core->dbg->maps, iter, map) {
		if (strstr (map->name, "libjemalloc.")) {
			jemalloc_addr = map->addr;
			path = map->name;
			break;
		}
	}
	if (!path) {
		R_LOG_WARN ("cannot find jemalloc mapped in memory (see `dm`)");
		return false;
	}
#if __linux__
	bool is_debug_file = GH(je_matched)(path, "/usr/local/lib");

	if (!is_debug_file) {
		R_LOG_WARN ("Cannot find libjemalloc.so.2 in /usr/local/lib");
		return false;
	}
	if (r_file_exists (path)) {
		ut64 vaddr = GH(je_get_va_symbol)(core, path, symname);
		if (jemalloc_addr != GHT_MAX && vaddr != 0) {
			*symbol = jemalloc_addr + vaddr;
			return true;
		}
	}
	return false;
#else
	eprintf ("[*] Resolving %s from libjemalloc.2... ", symname);
	// this is quite sloooow, we must optimize dmi
	char *va = r_core_cmd_strf (core, "dmi libjemalloc.2 %s$~[1]", symname);
	ut64 n = r_num_get (NULL, va);
	if (n && n != UT64_MAX) {
		*symbol = n;
		eprintf ("0x%08"PFMT64x" 0x%08" PFMT64x"\n", n, jemalloc_addr);
	} else {
		eprintf ("NOT FOUND\n");
	}
	free (va);
	return true;
#endif
}

static void GH(jemalloc_get_chunks)(RCore *core, RStrBuf *sb, const char *input) {
	ut64 cnksz;

	if (!GH(r_resolve_jemalloc)(core, "je_chunksize", &cnksz)) {
		R_LOG_ERROR ("Cannot find 'je_chunksize' symbol");
		return;
	}
	r_io_read_at (core->io, cnksz, (ut8 *)&cnksz, sizeof (GHT));

	switch (input[0]) {
	case '\0':
		R_LOG_ERROR ("need an arena_t to associate chunks");
		break;
	case ' ':
		{
			GHT arena = GHT_MAX;
			arena_t *ar = R_NEW0 (arena_t);
			extent_node_t *node = R_NEW0 (extent_node_t), *head = R_NEW0 (extent_node_t);
			input++;
			arena = r_num_math (core->num, input);

			if (arena) {
				r_io_read_at (core->io, arena, (ut8 *)ar, sizeof (arena_t));
				r_io_read_at (core->io, (GHT)(size_t)ar->achunks.qlh_first, (ut8 *)head, sizeof (extent_node_t));
				if (head->en_addr) {
					hprint (sb, "   Chunk - start: 0x%08"PFMT64x", end: 0x%08"PFMT64x", size: 0x%08"PFMT64x"\n",
						(ut64)(size_t)head->en_addr, (ut64)(size_t)((char *)head->en_addr + cnksz), (ut64)cnksz);
					r_io_read_at (core->io, (ut64)(size_t)head->ql_link.qre_next, (ut8 *)node, sizeof (extent_node_t));
					while (node && node->en_addr != head->en_addr) {
						hprint (sb, "   Chunk - start: 0x%08"PFMT64x", end: 0x%"PFMT64x", size: 0x%08"PFMT64x"\n",
							(ut64)(size_t)node->en_addr, (ut64)(size_t)((char *)node->en_addr + cnksz), cnksz);
						r_io_read_at (core->io, (ut64)(size_t)node->ql_link.qre_next, (ut8 *)node, sizeof (extent_node_t));
					}
				}
			}
			free (ar);
			free (head);
			free (node);
		break;
		}
	case '*':
		{
			int i = 0;
			ut64 sym;
			GHT arenas = GHT_MAX, arena = GHT_MAX;
			arena_t *ar = R_NEW0 (arena_t);
			extent_node_t *node = R_NEW0 (extent_node_t);
			extent_node_t *head = R_NEW0 (extent_node_t);

			if (GH(r_resolve_jemalloc) (core, "je_arenas", &sym)) {
				r_io_read_at (core->io, sym, (ut8 *)&arenas, sizeof (GHT));
				for (;;) {
					r_io_read_at (core->io, arenas + i * sizeof (GHT), (ut8 *)&arena, sizeof (GHT));
					if (!arena) {
						break;
					}
					hprint (sb, "arenas[%d]: @ 0x%"PFMTx" { \n", i++, (GHT)arena);
					r_io_read_at (core->io, arena, (ut8 *)ar, sizeof (arena_t));
					r_io_read_at (core->io, (GHT)(size_t)ar->achunks.qlh_first, (ut8 *)head, sizeof (extent_node_t));
					if (head->en_addr != 0) {
						hprint (sb, "   Chunk - start: 0x%08"PFMT64x", end: 0x%"PFMT64x", size: 0x%08"PFMT64x"\n",
							(ut64)(size_t)head->en_addr, (ut64)(size_t)((char *)head->en_addr + cnksz), (ut64)cnksz);
						ut64 addr = (ut64) (size_t)head->ql_link.qre_next;
						r_io_read_at (core->io, addr, (ut8 *)node, sizeof (extent_node_t));
						while (node && head && node->en_addr != head->en_addr) {
							hprint (sb, "   Chunk - start: 0x%08"PFMT64x", end: 0x%"PFMT64x", size: 0x%"PFMT64x"\n",
								(ut64)(size_t)node->en_addr, (ut64)(size_t)((char *)node->en_addr + cnksz), cnksz);
							r_io_read_at (core->io, (GHT)(size_t)node->ql_link.qre_next, (ut8 *)node, sizeof (extent_node_t));
						}
					}
					hprint (sb, " }\n");
				}
			}
			free (ar);
			free (head);
			free (node);
		}
		break;
	}
}

static void GH(jemalloc_print_narenas)(RCore *core, RStrBuf *sb, const char *input) {
	ut64 symaddr;
	ut64 arenas;
	GHT arena = GHT_MAX;
	arena_t *ar = R_NEW0 (arena_t);
	int i = 0;
	GHT narenas = 0;

	switch (input[0]) {
	case '\0':
		if (GH(r_resolve_jemalloc)(core, "narenas_total", &symaddr)) {
			r_io_read_at (core->io, symaddr, (ut8 *)&narenas, sizeof (GHT));
			hprint (sb, "narenas : %"PFMT64d"\n", (ut64)narenas);
		}
		if (narenas == 0) {
			R_LOG_ERROR ("No arenas allocated");
			free (ar);
			return;
		}
		if (narenas == GHT_MAX) {
			R_LOG_ERROR ("Cannot find narenas_total");
			free (ar);
			return;
		}

		if (GH(r_resolve_jemalloc)(core, "je_arenas", &arenas)) {
			r_io_read_at (core->io, arenas, (ut8 *)&arenas, sizeof (GHT));
			hprint (sb, "arenas[%"PFMT64d"] @ 0x%"PFMT64x" {\n", (ut64)narenas, (ut64)arenas);
			for (i = 0; i < narenas; i++) {
				ut64 at = arenas + (i * sizeof (GHT));
				r_io_read_at (core->io, at, (ut8 *)&arena, sizeof (GHT));
				if (!arena) {
					hprint (sb, "  arenas[%d]: (empty)\n", i);
					continue;
				}
				hprint (sb, "  arenas[%d]: @ 0x%"PFMT64x"\n", i, at);
			}
		}
		hprint (sb, " }\n");
		break;
	case ' ':
		arena = r_num_math (core->num, input + 1);
		r_io_read_at (core->io, (GHT)arena, (ut8 *)ar, sizeof (arena_t));

		hprint (sb, "struct arena_s {\n");
#define OO(x) (ut64)(arena + r_offsetof (arena_t, x))
		hprint (sb, "  ind = 0x%x\n", ar->ind);
		hprint (sb, "  nthreads: application allocation = 0x%"PFMT64x"\n", (ut64)ar->nthreads[0]);
		hprint (sb, "  nthreads: internal metadata allocation = 0x%"PFMT64x"\n", (ut64)ar->nthreads[1]);
		hprint (sb, "  lock = 0x%"PFMT64x"\n", OO(lock));
		hprint (sb, "  stats = 0x%"PFMT64x"\n", OO(stats));
		hprint (sb, "  tcache_ql = 0x%"PFMT64x"\n", OO(tcache_ql));
		hprint (sb, "  prof_accumbytes = 0x%"PFMT64x"x\n", (ut64)ar->prof_accumbytes);
		hprint (sb, "  offset_state = 0x%"PFMT64x"\n", (ut64)ar->offset_state);
		hprint (sb, "  dss_prec_t = 0x%"PFMT64x"\n", OO(dss_prec));
		hprint (sb, "  achunks = 0x%"PFMT64x"\n", OO(achunks));
		hprint (sb, "  extent_sn_next = 0x%"PFMT64x"\n", (ut64)(size_t)ar->extent_sn_next);
		hprint (sb, "  spare = 0x%"PFMT64x"\n", (ut64)(size_t)ar->spare);
		hprint (sb, "  lg_dirty_mult = 0x%"PFMT64x"\n", (ut64)(ssize_t)ar->lg_dirty_mult);
		hprint (sb, "  purging = %s\n", r_str_bool (ar->purging));
		hprint (sb, "  nactive = 0x%"PFMT64x"\n", (ut64)(size_t)ar->nactive);
		hprint (sb, "  ndirty = 0x%"PFMT64x"\n", (ut64)(size_t)ar->ndirty);

		hprint (sb, "  runs_dirty = 0x%"PFMT64x"\n", OO(runs_dirty));
		hprint (sb, "  chunks_cache = 0x%"PFMT64x"\n", OO(chunks_cache));
		hprint (sb, "  huge = 0x%"PFMT64x"\n", OO(huge));
		hprint (sb, "  huge_mtx = 0x%"PFMT64x"\n", OO(huge_mtx));

		hprint (sb, "  chunks_szsnad_cached = 0x%"PFMT64x"\n", OO(chunks_szsnad_cached));
		hprint (sb, "  chunks_ad_cached = 0x%"PFMT64x"\n", OO(chunks_ad_cached));
		hprint (sb, "  chunks_szsnad_retained = 0x%"PFMT64x"\n", OO(chunks_szsnad_retained));
		hprint (sb, "  chunks_ad_cached = 0x%"PFMT64x"\n", OO(chunks_ad_retained));

		hprint (sb, "  chunks_mtx = 0x%"PFMT64x"\n", OO(chunks_mtx));
		hprint (sb, "  node_cache = 0x%"PFMT64x"\n", OO(node_cache));
		hprint (sb, "  node_cache_mtx = 0x%"PFMT64x"\n", OO(node_cache_mtx));
		hprint (sb, "  chunks_hooks = 0x%"PFMT64x"\n", OO(chunk_hooks));
		hprint (sb, "  bins = %d 0x%"PFMT64x"\n", JM_NBINS, OO(bins));
		hprint (sb, "  runs_avail = %d 0x%"PFMT64x"\n", NPSIZES, OO(runs_avail));
		hprint (sb, " }\n");
		break;
	}
	free (ar);
}

static void GH(jemalloc_get_bins)(RCore *core, RStrBuf *sb, const char *input) {
	int i = 0, j;
	ut64 bin_info;
	ut64 arenas;
	GHT arena = GHT_MAX; //, bin = GHT_MAX;
	arena_t *ar = NULL;
	arena_bin_info_t *b = NULL;

	switch (input[0]) {
	case ' ':
		ar = R_NEW0 (arena_t);
		b = R_NEW0 (arena_bin_info_t);
		if (!GH(r_resolve_jemalloc)(core, "je_arena_bin_info", &bin_info)) {
			R_LOG_ERROR ("resolving je_arena_bin_info");
			R_FREE (b);
			break;
		}
		if (GH(r_resolve_jemalloc)(core, "je_arenas", &arenas)) {
			r_io_read_at (core->io, arenas, (ut8 *)&arenas, sizeof (GHT));
			hprint (sb, "arenas @ 0x%"PFMTx" {\n", (GHT)arenas);
			for (;;) {
				r_io_read_at (core->io, arenas + i * sizeof (GHT), (ut8 *)&arena, sizeof (GHT));
				if (!arena) {
					R_FREE (b);
					break;
				}
				hprint (sb, "   arenas[%d]: @ 0x%"PFMTx" {\n", i++, (GHT)arena);
				r_io_read_at (core->io, arena, (ut8 *)ar, sizeof (arena_t));
				for (j = 0; j < JM_NBINS; j++) {
					r_io_read_at (core->io, (GHT)(bin_info + j * sizeof (arena_bin_info_t)),
						(ut8*)b, sizeof (arena_bin_info_t));
					hprint (sb, "    {\n"
						"       regsize : 0x%zx\n"
						"       redzone size 0x%zx\n"
						"       reg_interval : 0x%zx\n"
						"       run_size : 0x%zx\n"
						"       nregs : 0x%x\n",
						b->reg_size, b->redzone_size, b->reg_interval, b->run_size, b->nregs);
					// FIXME: It's a structure of bitmap_info_t
					// hprint (sb, "       bitmap_info : ");
					// hprint (sb, "0x%"PFMT64x"\n", b->bitmap_info);
					hprint (sb, "       reg0_offset : 0x%"PFMT64x"\n\n", (ut64)b->reg0_offset);
					// FIXME: It's a structure of malloc_mutex_t
					// hprint (sb, "       bins[%d]->lock ", j);
					// hprint (sb, "= 0x%"PFMT64x"\n", ar->bins[j].lock);
					// FIXME: It's a structure of arena_run_t*
					// hprint (sb, "       bins[%d]->runcur ", j);
					// hprint (sb, "@ 0x%"PFMT64x"\n", ar->bins[j].runcur);
					// FIXME: It's a structure of arena_run_heap_t*
					// hprint (sb, "       bins[%d]->runs ", j);
					// hprint (sb, "@ 0x%"PFMTx"\n", ar->bins[j].runs);
					// FIXME: It's a structure of malloc_bin_stats_t
					// hprint (sb, "       bins[%d]->stats ", j);
					// hprint (sb, "= 0x%"PFMTx"\n", ar->bins[j].stats);
					hprint (sb, "    }\n");
				}
				hprint (sb, "  }\n");
			}
		}
		hprint (sb, " }\n");
		break;
	}
	free (ar);
	free (b);
}

#if 0
static void GH(jemalloc_get_runs)(RCore *core, const char *input) {
	switch (input[0]) {
	case ' ':
		{
			int pageind;
			ut64 npages, chunksize_mask, map_bias, map_misc_offset, chunk, mapbits;
			arena_chunk_t *c = R_NEW0 (arena_chunk_t);

			if (!c) {
				return;
			}

			input += 1;
			chunk = r_num_math (core->num, input);

			if (!GH(r_resolve_jemalloc)(core, "je_chunk_npages", &npages)) {
				R_LOG_ERROR ("resolving je_chunk_npages");
				return;
			}
			if (!GH(r_resolve_jemalloc)(core, "je_chunksize_mask", &chunksize_mask)) {
				R_LOG_ERROR ("resolving je_chunksize_mask");
				return;
			}
			if (!GH(r_resolve_jemalloc)(core, "je_map_bias", &map_bias)) {
				eprintf ("resolving je_map_bias");
				return;
			}
			if (!GH(r_resolve_jemalloc)(core, "je_map_misc_offset", &map_misc_offset)) {
				eprintf ("resolving je_map_misc_offset");
				return;
			}

			r_io_read_at (core->io, npages, (ut8*)&npages, sizeof (GHT));
			r_io_read_at (core->io, chunksize_mask, (ut8*)&chunksize_mask, sizeof (GHT));
			r_io_read_at (core->io, map_bias, (ut8*)&map_bias, sizeof (GHT));
			r_io_read_at (core->io, map_misc_offset, (ut8*)&map_misc_offset, sizeof (GHT));

			eprintf ("map_misc_offset 0x%08"PFMT64x"\n", (ut64)map_misc_offset);

			r_io_read_at (core->io, chunk, (ut8 *)c, sizeof (arena_chunk_t));
			mapbits = *(GHT *)&c->map_bits;
			eprintf ("map_bits: 0x%08"PFMT64x"\n", (ut64)mapbits);

			uint32_t offset = r_offsetof (arena_chunk_t, map_bits);

			arena_chunk_map_bits_t *dwords = (void *)calloc (sizeof (arena_chunk_map_bits_t), npages);
			r_io_read_at (core->io, chunk + offset, (ut8*)dwords, sizeof (arena_chunk_map_bits_t) * npages);
			eprintf ("map_bits @ 0x%08"PFMT64x"\n", (ut64)(chunk + offset));

			arena_run_t *r = R_NEW0 (arena_run_t);
			if (!r) {
				R_LOG_ERROR ("calling calloc");
				return;
			}
			for (pageind = map_bias; pageind < npages; pageind++) {
				arena_chunk_map_bits_t mapelm = dwords[pageind-map_bias];
				if (mapelm.bits & CHUNK_MAP_ALLOCATED) {
					// ut64 elm = ((arena_chunk_map_misc_t *)((uintptr_t)chunk + (uintptr_t)map_misc_offset) + pageind-map_bias);
					ut64 elm = chunk + map_misc_offset + pageind-map_bias;
					eprintf ("\nelm: 0x%"PFMT64x"\n", elm);
					arena_chunk_map_misc_t *m = R_NEW0 (arena_chunk_map_misc_t);
					if (m) {
						ut64 run = elm + r_offsetof (arena_chunk_map_misc_t, run);
						r_io_read_at (core->io, elm, (ut8*)m, sizeof (arena_chunk_map_misc_t));
						eprintf ("Small run @ 0x%08"PFMT64x"\n", (ut64)elm);
						r_io_read_at (core->io, run, (ut8*)r, sizeof (arena_run_t));
						eprintf ("binind: 0x%08"PFMT64x"\n", (ut64)r->binind);
						eprintf ("nfree: 0x%08"PFMT64x"\n", (ut64)r->nfree);
						eprintf ("bitmap: 0x%08"PFMT64x"\n\n", (ut64)*(GHT*)r->bitmap);
						free (m);
					}
				} else if (mapelm.bits & CHUNK_MAP_LARGE) {
					ut64 run = (ut64) (size_t) chunk + (pageind << LG_PAGE);
					eprintf ("Large run @ 0x%08"PFMT64x"\n", run);
					r_io_read_at (core->io, run, (ut8*)r, sizeof (arena_run_t));
					eprintf ("binind: 0x%08"PFMT64x"\n", (ut64)r->binind);
					eprintf ("nfree: 0x%08"PFMT64x"\n", (ut64)r->nfree);
					eprintf ("bitmap: 0x%08"PFMT64x"\n\n", (ut64)*(GHT*)r->bitmap);
				}
			}
			free (c);
			free (r);
		}
	break;
	}
}
#endif

static char *GH(dmh_jemalloc)(RCore *core, const char *input) {
	RStrBuf sb;
	r_strbuf_init (&sb);
	switch (input[0]) {
	case 'a': //dmha
		GH(jemalloc_print_narenas) (core, &sb, input + 1);
		break;
	case 'b': //dmhb
		GH(jemalloc_get_bins) (core, &sb, input + 1);
		break;
	case 'c': //dmhc
		GH(jemalloc_get_chunks) (core, &sb, input + 1);
		break;
	default:
		R_LOG_ERROR ("Unknown subcommand. See dmh[abc]");
		break;
	/*
	case 'r': //dmhr
		GH(jemalloc_get_runs) (core, input + 1);
		break;
	*/
	}
	return r_strbuf_drain_nofree (&sb);
}

#endif