// SPDX-License-Identifier: GPL-2.0
/*
 * Copyright (C) 2025 - Google LLC
 * Author: Vincent Donnefort <vdonnefort@google.com>
 */

#include <linux/module.h>
#include <linux/simple_ring_buffer.h>
#include <linux/trace_remote.h>
#include <linux/tracefs.h>
#include <linux/types.h>

#define REMOTE_EVENT_INCLUDE_FILE kernel/trace/remote_test_events.h
#include <trace/define_remote_events.h>

static DEFINE_PER_CPU(struct simple_rb_per_cpu *, simple_rbs);
static struct trace_buffer_desc *remote_test_buffer_desc;

/*
 * The trace_remote lock already serializes accesses from the trace_remote_callbacks.
 * However write_event can still race with load/unload.
 */
static DEFINE_MUTEX(simple_rbs_lock);

static int remote_test_load_simple_rb(int cpu, struct ring_buffer_desc *rb_desc)
{
	struct simple_rb_per_cpu *cpu_buffer;
	struct simple_buffer_page *bpages;
	int ret = -ENOMEM;

	cpu_buffer = kmalloc_obj(*cpu_buffer);
	if (!cpu_buffer)
		return ret;

	bpages = kmalloc_objs(*bpages, rb_desc->nr_page_va);
	if (!bpages)
		goto err_free_cpu_buffer;

	ret = simple_ring_buffer_init(cpu_buffer, bpages, rb_desc);
	if (ret)
		goto err_free_bpages;

	scoped_guard(mutex, &simple_rbs_lock) {
		WARN_ON(*per_cpu_ptr(&simple_rbs, cpu));
		*per_cpu_ptr(&simple_rbs, cpu) = cpu_buffer;
	}

	return 0;

err_free_bpages:
	kfree(bpages);

err_free_cpu_buffer:
	kfree(cpu_buffer);

	return ret;
}

static void remote_test_unload_simple_rb(int cpu)
{
	struct simple_rb_per_cpu *cpu_buffer = *per_cpu_ptr(&simple_rbs, cpu);
	struct simple_buffer_page *bpages;

	if (!cpu_buffer)
		return;

	guard(mutex)(&simple_rbs_lock);

	bpages = cpu_buffer->bpages;
	simple_ring_buffer_unload(cpu_buffer);
	kfree(bpages);
	kfree(cpu_buffer);
	*per_cpu_ptr(&simple_rbs, cpu) = NULL;
}

static struct trace_buffer_desc *remote_test_load(unsigned long size, void *unused)
{
	struct ring_buffer_desc *rb_desc;
	struct trace_buffer_desc *desc;
	size_t desc_size;
	int cpu, ret;

	if (WARN_ON(remote_test_buffer_desc))
		return ERR_PTR(-EINVAL);

	desc_size = trace_buffer_desc_size(size, num_possible_cpus());
	if (desc_size == SIZE_MAX) {
		ret = -E2BIG;
		goto err;
	}

	desc = kmalloc(desc_size, GFP_KERNEL);
	if (!desc) {
		ret = -ENOMEM;
		goto err;
	}

	ret = trace_remote_alloc_buffer(desc, desc_size, size, cpu_possible_mask);
	if (ret)
		goto err_free_desc;

	for_each_ring_buffer_desc(rb_desc, cpu, desc) {
		ret = remote_test_load_simple_rb(rb_desc->cpu, rb_desc);
		if (ret)
			goto err_unload;
	}

	remote_test_buffer_desc = desc;

	return remote_test_buffer_desc;

err_unload:
	for_each_ring_buffer_desc(rb_desc, cpu, desc)
		remote_test_unload_simple_rb(rb_desc->cpu);
	trace_remote_free_buffer(desc);

err_free_desc:
	kfree(desc);

err:
	return ERR_PTR(ret);
}

static void remote_test_unload(struct trace_buffer_desc *desc, void *unused)
{
	struct ring_buffer_desc *rb_desc;
	int cpu;

	if (WARN_ON(desc != remote_test_buffer_desc))
		return;

	for_each_ring_buffer_desc(rb_desc, cpu, desc)
		remote_test_unload_simple_rb(rb_desc->cpu);

	remote_test_buffer_desc = NULL;
	trace_remote_free_buffer(desc);
	kfree(desc);
}

static int remote_test_enable_tracing(bool enable, void *unused)
{
	struct ring_buffer_desc *rb_desc;
	int cpu;

	if (!remote_test_buffer_desc)
		return -ENODEV;

	for_each_ring_buffer_desc(rb_desc, cpu, remote_test_buffer_desc)
		WARN_ON(simple_ring_buffer_enable_tracing(*per_cpu_ptr(&simple_rbs, rb_desc->cpu),
							  enable));
	return 0;
}

static int remote_test_swap_reader_page(unsigned int cpu, void *unused)
{
	struct simple_rb_per_cpu *cpu_buffer;

	if (cpu >= NR_CPUS)
		return -EINVAL;

	cpu_buffer = *per_cpu_ptr(&simple_rbs, cpu);
	if (!cpu_buffer)
		return -EINVAL;

	return simple_ring_buffer_swap_reader_page(cpu_buffer);
}

static int remote_test_reset(unsigned int cpu, void *unused)
{
	struct simple_rb_per_cpu *cpu_buffer;

	if (cpu >= NR_CPUS)
		return -EINVAL;

	cpu_buffer = *per_cpu_ptr(&simple_rbs, cpu);
	if (!cpu_buffer)
		return -EINVAL;

	return simple_ring_buffer_reset(cpu_buffer);
}

static int remote_test_enable_event(unsigned short id, bool enable, void *unused)
{
	if (id != REMOTE_TEST_EVENT_ID)
		return -EINVAL;

	/*
	 * Let's just use the struct remote_event enabled field that is turned on and off by
	 * trace_remote. This is a bit racy but good enough for a simple test module.
	 */
	return 0;
}

static ssize_t
write_event_write(struct file *filp, const char __user *ubuf, size_t cnt, loff_t *pos)
{
	struct remote_event_format_selftest *evt_test;
	struct simple_rb_per_cpu *cpu_buffer;
	unsigned long val;
	int ret;

	ret = kstrtoul_from_user(ubuf, cnt, 10, &val);
	if (ret)
		return ret;

	guard(mutex)(&simple_rbs_lock);

	if (!remote_event_selftest.enabled)
		return -ENODEV;

	guard(preempt)();

	cpu_buffer = *this_cpu_ptr(&simple_rbs);
	if (!cpu_buffer)
		return -ENODEV;

	evt_test = simple_ring_buffer_reserve(cpu_buffer,
					      sizeof(struct remote_event_format_selftest),
					      trace_clock_global());
	if (!evt_test)
		return -ENODEV;

	evt_test->hdr.id = REMOTE_TEST_EVENT_ID;
	evt_test->id = val;

	simple_ring_buffer_commit(cpu_buffer);

	return cnt;
}

static const struct file_operations write_event_fops = {
	.write	= write_event_write,
};

static int remote_test_init_tracefs(struct dentry *d, void *unused)
{
	return tracefs_create_file("write_event", 0200, d, NULL, &write_event_fops) ?
		0 : -ENOMEM;
}

static struct trace_remote_callbacks trace_remote_callbacks = {
	.init			= remote_test_init_tracefs,
	.load_trace_buffer	= remote_test_load,
	.unload_trace_buffer	= remote_test_unload,
	.enable_tracing		= remote_test_enable_tracing,
	.swap_reader_page	= remote_test_swap_reader_page,
	.reset			= remote_test_reset,
	.enable_event		= remote_test_enable_event,
};

static int __init remote_test_init(void)
{
	return trace_remote_register("test", &trace_remote_callbacks, NULL,
				     &remote_event_selftest, 1);
}

module_init(remote_test_init);

MODULE_DESCRIPTION("Test module for the trace remote interface");
MODULE_AUTHOR("Vincent Donnefort");
MODULE_LICENSE("GPL");