/* $OpenBSD$ */

/*
 * Copyright (c) 2024 The University of Queensland
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

#include <sys/types.h>
#include <sys/ktrace.h>
#include <sys/cdefs.h>

#include <stdint.h>
#include <string.h>
#include <stdlib.h>
#include <time.h>

#include <stdio.h>
#include <dlfcn.h>

__dso_hidden void __cyg_profile_func_enter(void *, void *);
__dso_hidden void __cyg_profile_func_exit(void *, void *);

int _thread_atfork(void (*)(void), void (*)(void), void (*)(void), void *);

struct omalloc_event {
	uint64_t	 ts;
	uintptr_t	 fn;
};

#define EVTLEN (KTR_USER_MAXLEN / sizeof(struct omalloc_event))

static struct omalloc_event omalloc_events[EVTLEN];
static unsigned int omalloc_idx;

static uintptr_t omalloc_base;

static void
record(int op, void *fn)
{
	struct omalloc_event *e = &omalloc_events[omalloc_idx++];
	struct timespec now;

	clock_gettime(CLOCK_REALTIME, &now);

	if (omalloc_base == 0) {
	}

	e->ts = now.tv_sec * 1000000000 + now.tv_nsec;
	e->fn = (uintptr_t)fn - omalloc_base;
	e->fn |= op;

	if (omalloc_idx >= EVTLEN) {
		utrace("omalloc:tr", omalloc_events, sizeof(omalloc_events));
		omalloc_idx = 0;
	}
}

struct omalloc_meta {
	uintptr_t	base;
};

static void
omalloc_utrace_meta(void)
{
	Dl_info info;
	char buf[KTR_USER_MAXLEN];
	struct omalloc_meta *meta = (struct omalloc_meta *)buf;
	size_t len;

	if (dladdr(free, &info) == 0)
		return;

	len = strlen(info.dli_fname);
	if (len > sizeof(buf) - sizeof(*meta))
		len = sizeof(buf) - sizeof(*meta);

	meta->base = (uintptr_t)info.dli_fbase;
	memcpy(meta + 1, info.dli_fname, len);

	utrace("omalloc:md", buf, sizeof(*meta) + len);
}

static void
omalloc_atfork(void)
{
	omalloc_idx = 0;
	omalloc_utrace_meta();
}

static void
__attribute__((constructor))
omalloc_ctor(void)
{
	_thread_atfork(NULL, NULL, omalloc_atfork, NULL);
	omalloc_utrace_meta();
}

static void
__attribute__((destructor))
omalloc_dtor(void)
{
	unsigned int i, c;
	utrace("omalloc:tr",
	    omalloc_events, omalloc_idx * sizeof(struct omalloc_event));
	omalloc_idx = 0;
}

void
__cyg_profile_func_enter(void *fn, void *pc)
{
	record(0, fn);
}

void
__cyg_profile_func_exit(void *fn, void *pc)
{
	record(1, fn);
}