/*-
 * SPDX-License-Identifier: BSD-2-Clause
 *
 * Copyright (c) 2023-2025 Ruslan Bukin <br@bsdpad.com>
 *
 * This work was supported by Innovate UK project 105694, "Digital Security
 * by Design (DSbD) Technology Platform Prototype".
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
 * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
 * SUCH DAMAGE.
 */

#include <sys/param.h>
#include <sys/ioccom.h>
#include <sys/conf.h>
#include <sys/proc.h>
#include <sys/kernel.h>
#include <sys/malloc.h>
#include <sys/mman.h>
#include <sys/refcount.h>
#include <sys/rwlock.h>
#include <sys/hwt.h>
#include <sys/smp.h>

#include <vm/vm.h>
#include <vm/pmap.h>
#include <vm/vm_extern.h>
#include <vm/vm_param.h>
#include <vm/vm_kern.h>
#include <vm/vm_page.h>
#include <vm/vm_object.h>
#include <vm/vm_pager.h>
#include <vm/vm_pageout.h>
#include <vm/vm_phys.h>

#include <dev/hwt/hwt_hook.h>
#include <dev/hwt/hwt_context.h>
#include <dev/hwt/hwt_contexthash.h>
#include <dev/hwt/hwt_config.h>
#include <dev/hwt/hwt_cpu.h>
#include <dev/hwt/hwt_owner.h>
#include <dev/hwt/hwt_ownerhash.h>
#include <dev/hwt/hwt_thread.h>
#include <dev/hwt/hwt_backend.h>
#include <dev/hwt/hwt_vm.h>
#include <dev/hwt/hwt_record.h>

#define	HWT_THREAD_DEBUG
#undef	HWT_THREAD_DEBUG

#ifdef	HWT_THREAD_DEBUG
#define	dprintf(fmt, ...)	printf(fmt, ##__VA_ARGS__)
#else
#define	dprintf(fmt, ...)
#endif

static MALLOC_DEFINE(M_HWT_VM, "hwt_vm", "Hardware Trace");

static int
hwt_vm_fault(vm_object_t vm_obj, vm_ooffset_t offset,
    int prot, vm_page_t *mres)
{

	return (0);
}

static int
hwt_vm_ctor(void *handle, vm_ooffset_t size, vm_prot_t prot,
    vm_ooffset_t foff, struct ucred *cred, u_short *color)
{

	*color = 0;

	return (0);
}

static void
hwt_vm_dtor(void *handle)
{

}

static struct cdev_pager_ops hwt_vm_pager_ops = {
	.cdev_pg_fault = hwt_vm_fault,
	.cdev_pg_ctor = hwt_vm_ctor,
	.cdev_pg_dtor = hwt_vm_dtor
};

static int
hwt_vm_alloc_pages(struct hwt_vm *vm, int kva_req)
{
	vm_paddr_t low, high, boundary;
	vm_memattr_t memattr;
#ifdef  __aarch64__
	uintptr_t va;
#endif
	int alignment;
	vm_page_t m;
	int pflags;
	int tries;
	int i;

	alignment = PAGE_SIZE;
	low = 0;
	high = -1UL;
	boundary = 0;
	pflags = VM_ALLOC_NORMAL | VM_ALLOC_WIRED | VM_ALLOC_ZERO;
	memattr = VM_MEMATTR_DEVICE;

	if (kva_req) {
		vm->kvaddr = kva_alloc(vm->npages * PAGE_SIZE);
		if (!vm->kvaddr)
			return (ENOMEM);
	}

	vm->obj = cdev_pager_allocate(vm, OBJT_MGTDEVICE,
	    &hwt_vm_pager_ops, vm->npages * PAGE_SIZE, PROT_READ, 0,
	    curthread->td_ucred);

	for (i = 0; i < vm->npages; i++) {
		tries = 0;
retry:
		m = vm_page_alloc_noobj_contig(pflags, 1, low, high,
		    alignment, boundary, memattr);
		if (m == NULL) {
			if (tries < 3) {
				if (!vm_page_reclaim_contig(pflags, 1, low,
				    high, alignment, boundary))
					vm_wait(NULL);
				tries++;
				goto retry;
			}

			return (ENOMEM);
		}

#if 0
		/* TODO: could not clean device memory on arm64. */
		if ((m->flags & PG_ZERO) == 0)
			pmap_zero_page(m);
#endif

#ifdef __aarch64__
		va = PHYS_TO_DMAP(VM_PAGE_TO_PHYS(m));
		cpu_dcache_wb_range((void *)va, PAGE_SIZE);
#endif

		m->valid = VM_PAGE_BITS_ALL;
		m->oflags &= ~VPO_UNMANAGED;
		m->flags |= PG_FICTITIOUS;
		vm->pages[i] = m;

		VM_OBJECT_WLOCK(vm->obj);
		vm_page_insert(m, vm->obj, i);
		if (kva_req)
			pmap_qenter(vm->kvaddr + i * PAGE_SIZE, &m, 1);
		VM_OBJECT_WUNLOCK(vm->obj);
	}

	return (0);
}

static int
hwt_vm_open(struct cdev *cdev, int oflags, int devtype, struct thread *td)
{

	dprintf("%s\n", __func__);

	return (0);
}

static int
hwt_vm_mmap_single(struct cdev *cdev, vm_ooffset_t *offset,
    vm_size_t mapsize, struct vm_object **objp, int nprot)
{
	struct hwt_vm *vm;

	vm = cdev->si_drv1;

	if (nprot != PROT_READ || *offset != 0)
		return (ENXIO);

	vm_object_reference(vm->obj);
	*objp = vm->obj;

	return (0);
}

static void
hwt_vm_start_cpu_mode(struct hwt_context *ctx)
{
	cpuset_t enable_cpus;
	int cpu_id;

	CPU_ZERO(&enable_cpus);

	CPU_FOREACH_ISSET(cpu_id, &ctx->cpu_map) {
#ifdef SMP
		/* Ensure CPU is not halted. */
		if (CPU_ISSET(cpu_id, &hlt_cpus_mask))
			continue;
#endif

		hwt_backend_configure(ctx, cpu_id, cpu_id);

		CPU_SET(cpu_id, &enable_cpus);
	}

	if (ctx->hwt_backend->ops->hwt_backend_enable_smp == NULL) {
		CPU_FOREACH_ISSET(cpu_id, &enable_cpus)
			hwt_backend_enable(ctx, cpu_id);
	} else {
		/* Some backends require enabling all CPUs at once. */
		hwt_backend_enable_smp(ctx);
	}
}

static int
hwt_vm_ioctl(struct cdev *dev, u_long cmd, caddr_t addr, int flags,
    struct thread *td)
{
	struct hwt_record_get *rget;
	struct hwt_set_config *sconf;
	struct hwt_bufptr_get *ptr_get;
	struct hwt_svc_buf *sbuf;

	struct hwt_context *ctx;
	struct hwt_vm *vm;
	struct hwt_owner *ho;

	vm_offset_t offset;
	int ident;
	int error;
	uint64_t data = 0;
	void *data2;
	size_t data_size;
	int data_version;

	vm = dev->si_drv1;
	KASSERT(vm != NULL, ("si_drv1 is NULL"));

	ctx = vm->ctx;

	/* Ensure process is registered owner of this HWT. */
	ho = hwt_ownerhash_lookup(td->td_proc);
	if (ho == NULL)
		return (ENXIO);

	if (ctx->hwt_owner != ho)
		return (EPERM);

	switch (cmd) {
	case HWT_IOC_START:
		dprintf("%s: start tracing\n", __func__);

		HWT_CTX_LOCK(ctx);
		if (ctx->state == CTX_STATE_RUNNING) {
			/* Already running ? */
			HWT_CTX_UNLOCK(ctx);
			return (ENXIO);
		}
		ctx->state = CTX_STATE_RUNNING;
		HWT_CTX_UNLOCK(ctx);

		if (ctx->mode == HWT_MODE_CPU)
			hwt_vm_start_cpu_mode(ctx);
		else {
			/*
			 * Tracing backend will be configured and enabled
			 * during hook invocation. See hwt_hook.c.
			 */
		}

		break;

	case HWT_IOC_STOP:
		if (ctx->state == CTX_STATE_STOPPED)
			return (ENXIO);
		hwt_backend_stop(ctx);
		ctx->state = CTX_STATE_STOPPED;
		break;

	case HWT_IOC_RECORD_GET:
		rget = (struct hwt_record_get *)addr;
		error = hwt_record_send(ctx, rget);
		if (error)
			return (error);
		break;

	case HWT_IOC_SET_CONFIG:
		if (ctx->state == CTX_STATE_RUNNING) {
			return (ENXIO);
		}
		sconf = (struct hwt_set_config *)addr;
		error = hwt_config_set(td, ctx, sconf);
		if (error)
			return (error);
		ctx->pause_on_mmap = sconf->pause_on_mmap ? 1 : 0;
		break;

	case HWT_IOC_WAKEUP:

		if (ctx->mode == HWT_MODE_CPU)
			return (ENXIO);

		KASSERT(vm->thr != NULL, ("thr is NULL"));

		wakeup(vm->thr);

		break;

	case HWT_IOC_BUFPTR_GET:
		ptr_get = (struct hwt_bufptr_get *)addr;

		error = hwt_backend_read(ctx, vm, &ident, &offset, &data);
		if (error)
			return (error);

		if (ptr_get->ident)
			error = copyout(&ident, ptr_get->ident, sizeof(int));
		if (error)
			return (error);

		if (ptr_get->offset)
			error = copyout(&offset, ptr_get->offset,
			    sizeof(vm_offset_t));
		if (error)
			return (error);

		if (ptr_get->data)
			error = copyout(&data, ptr_get->data, sizeof(uint64_t));
		if (error)
			return (error);

		break;

	case HWT_IOC_SVC_BUF:
		if (ctx->state == CTX_STATE_STOPPED) {
			return (ENXIO);
		}

		sbuf = (struct hwt_svc_buf *)addr;
		data_size = sbuf->data_size;
		data_version = sbuf->data_version;

		if (data_size == 0 || data_size > PAGE_SIZE)
			return (EINVAL);

		data2 = malloc(data_size, M_HWT_VM, M_WAITOK | M_ZERO);
		error = copyin(sbuf->data, data2, data_size);
		if (error) {
			free(data2, M_HWT_VM);
			return (error);
		}

		error = hwt_backend_svc_buf(ctx, data2, data_size, data_version);
		if (error) {
			free(data2, M_HWT_VM);
			return (error);
		}

		free(data2, M_HWT_VM);
		break;

	default:
		break;
	}

	return (0);
}

static struct cdevsw hwt_vm_cdevsw = {
	.d_version	= D_VERSION,
	.d_name		= "hwt",
	.d_open		= hwt_vm_open,
	.d_mmap_single	= hwt_vm_mmap_single,
	.d_ioctl	= hwt_vm_ioctl,
};

static int
hwt_vm_create_cdev(struct hwt_vm *vm, char *path)
{
	struct make_dev_args args;
	int error;

	dprintf("%s: path %s\n", __func__, path);

	make_dev_args_init(&args);
	args.mda_devsw = &hwt_vm_cdevsw;
	args.mda_flags = MAKEDEV_CHECKNAME | MAKEDEV_WAITOK;
	args.mda_uid = UID_ROOT;
	args.mda_gid = GID_WHEEL;
	args.mda_mode = 0660;
	args.mda_si_drv1 = vm;

	error = make_dev_s(&args, &vm->cdev, "%s", path);
	if (error != 0)
		return (error);

	return (0);
}

static int
hwt_vm_alloc_buffers(struct hwt_vm *vm, int kva_req)
{
	int error;

	vm->pages = malloc(sizeof(struct vm_page *) * vm->npages,
	    M_HWT_VM, M_WAITOK | M_ZERO);

	error = hwt_vm_alloc_pages(vm, kva_req);
	if (error) {
		printf("%s: could not alloc pages\n", __func__);
		return (error);
	}

	return (0);
}

static void
hwt_vm_destroy_buffers(struct hwt_vm *vm)
{
	vm_page_t m;
	int i;

	if (vm->ctx->hwt_backend->kva_req && vm->kvaddr != 0) {
		pmap_qremove(vm->kvaddr, vm->npages);
		kva_free(vm->kvaddr, vm->npages * PAGE_SIZE);
	}
	VM_OBJECT_WLOCK(vm->obj);
	for (i = 0; i < vm->npages; i++) {
		m = vm->pages[i];
		if (m == NULL)
			break;

		vm_page_busy_acquire(m, 0);
		cdev_pager_free_page(vm->obj, m);
		m->flags &= ~PG_FICTITIOUS;
		vm_page_unwire_noq(m);
		vm_page_free(m);

	}
	vm_pager_deallocate(vm->obj);
	VM_OBJECT_WUNLOCK(vm->obj);

	free(vm->pages, M_HWT_VM);
}

void
hwt_vm_free(struct hwt_vm *vm)
{

	dprintf("%s\n", __func__);

	if (vm->cdev)
		destroy_dev_sched(vm->cdev);
	hwt_vm_destroy_buffers(vm);
	free(vm, M_HWT_VM);
}

int
hwt_vm_alloc(size_t bufsize, int kva_req, char *path, struct hwt_vm **vm0)
{
	struct hwt_vm *vm;
	int error;

	vm = malloc(sizeof(struct hwt_vm), M_HWT_VM, M_WAITOK | M_ZERO);
	vm->npages = bufsize / PAGE_SIZE;

	error = hwt_vm_alloc_buffers(vm, kva_req);
	if (error) {
		free(vm, M_HWT_VM);
		return (error);
	}

	error = hwt_vm_create_cdev(vm, path);
	if (error) {
		hwt_vm_free(vm);
		return (error);
	}

	*vm0 = vm;

	return (0);
}