// SPDX-License-Identifier: GPL-2.0
/*
 * Test module for in-kernel kprobe event creation and generation.
 *
 * Copyright (C) 2019 Tom Zanussi <zanussi@kernel.org>
 */

#include <linux/module.h>
#include <linux/trace_events.h>

/*
 * This module is a simple test of basic functionality for in-kernel
 * kprobe/kretprobe event creation.  The first test uses
 * kprobe_event_gen_cmd_start(), kprobe_event_add_fields() and
 * kprobe_event_gen_cmd_end() to create a kprobe event, which is then
 * enabled in order to generate trace output.  The second creates a
 * kretprobe event using kretprobe_event_gen_cmd_start() and
 * kretprobe_event_gen_cmd_end(), and is also then enabled.
 *
 * To test, select CONFIG_KPROBE_EVENT_GEN_TEST and build the module.
 * Then:
 *
 * # insmod kernel/trace/kprobe_event_gen_test.ko
 * # cat /sys/kernel/tracing/trace
 *
 * You should see many instances of the "gen_kprobe_test" and
 * "gen_kretprobe_test" events in the trace buffer.
 *
 * To remove the events, remove the module:
 *
 * # rmmod kprobe_event_gen_test
 *
 */

static struct trace_event_file *gen_kprobe_test;
static struct trace_event_file *gen_kretprobe_test;

#define KPROBE_GEN_TEST_FUNC	"do_sys_open"

/* X86 */
#if defined(CONFIG_X86_64) || defined(CONFIG_X86_32)
#define KPROBE_GEN_TEST_ARG0	"dfd=%ax"
#define KPROBE_GEN_TEST_ARG1	"filename=%dx"
#define KPROBE_GEN_TEST_ARG2	"flags=%cx"
#define KPROBE_GEN_TEST_ARG3	"mode=+4($stack)"

/* ARM64 */
#elif defined(CONFIG_ARM64)
#define KPROBE_GEN_TEST_ARG0	"dfd=%x0"
#define KPROBE_GEN_TEST_ARG1	"filename=%x1"
#define KPROBE_GEN_TEST_ARG2	"flags=%x2"
#define KPROBE_GEN_TEST_ARG3	"mode=%x3"

/* ARM */
#elif defined(CONFIG_ARM)
#define KPROBE_GEN_TEST_ARG0	"dfd=%r0"
#define KPROBE_GEN_TEST_ARG1	"filename=%r1"
#define KPROBE_GEN_TEST_ARG2	"flags=%r2"
#define KPROBE_GEN_TEST_ARG3	"mode=%r3"

/* RISCV */
#elif defined(CONFIG_RISCV)
#define KPROBE_GEN_TEST_ARG0	"dfd=%a0"
#define KPROBE_GEN_TEST_ARG1	"filename=%a1"
#define KPROBE_GEN_TEST_ARG2	"flags=%a2"
#define KPROBE_GEN_TEST_ARG3	"mode=%a3"

/* others */
#else
#define KPROBE_GEN_TEST_ARG0	NULL
#define KPROBE_GEN_TEST_ARG1	NULL
#define KPROBE_GEN_TEST_ARG2	NULL
#define KPROBE_GEN_TEST_ARG3	NULL
#endif

static bool trace_event_file_is_valid(struct trace_event_file *input)
{
	return input && !IS_ERR(input);
}

/*
 * Test to make sure we can create a kprobe event, then add more
 * fields.
 */
static int __init test_gen_kprobe_cmd(void)
{
	struct dynevent_cmd cmd;
	char *buf;
	int ret;

	/* Create a buffer to hold the generated command */
	buf = kzalloc(MAX_DYNEVENT_CMD_LEN, GFP_KERNEL);
	if (!buf)
		return -ENOMEM;

	/* Before generating the command, initialize the cmd object */
	kprobe_event_cmd_init(&cmd, buf, MAX_DYNEVENT_CMD_LEN);

	/*
	 * Define the gen_kprobe_test event with the first 2 kprobe
	 * fields.
	 */
	ret = kprobe_event_gen_cmd_start(&cmd, "gen_kprobe_test",
					 KPROBE_GEN_TEST_FUNC,
					 KPROBE_GEN_TEST_ARG0, KPROBE_GEN_TEST_ARG1);
	if (ret)
		goto out;

	/* Use kprobe_event_add_fields to add the rest of the fields */

	ret = kprobe_event_add_fields(&cmd, KPROBE_GEN_TEST_ARG2, KPROBE_GEN_TEST_ARG3);
	if (ret)
		goto out;

	/*
	 * This actually creates the event.
	 */
	ret = kprobe_event_gen_cmd_end(&cmd);
	if (ret)
		goto out;

	/*
	 * Now get the gen_kprobe_test event file.  We need to prevent
	 * the instance and event from disappearing from underneath
	 * us, which trace_get_event_file() does (though in this case
	 * we're using the top-level instance which never goes away).
	 */
	gen_kprobe_test = trace_get_event_file(NULL, "kprobes",
					       "gen_kprobe_test");
	if (IS_ERR(gen_kprobe_test)) {
		ret = PTR_ERR(gen_kprobe_test);
		goto delete;
	}

	/* Enable the event or you won't see anything */
	ret = trace_array_set_clr_event(gen_kprobe_test->tr,
					"kprobes", "gen_kprobe_test", true);
	if (ret) {
		trace_put_event_file(gen_kprobe_test);
		goto delete;
	}
 out:
	kfree(buf);
	return ret;
 delete:
	if (trace_event_file_is_valid(gen_kprobe_test))
		gen_kprobe_test = NULL;
	/* We got an error after creating the event, delete it */
	kprobe_event_delete("gen_kprobe_test");
	goto out;
}

/*
 * Test to make sure we can create a kretprobe event.
 */
static int __init test_gen_kretprobe_cmd(void)
{
	struct dynevent_cmd cmd;
	char *buf;
	int ret;

	/* Create a buffer to hold the generated command */
	buf = kzalloc(MAX_DYNEVENT_CMD_LEN, GFP_KERNEL);
	if (!buf)
		return -ENOMEM;

	/* Before generating the command, initialize the cmd object */
	kprobe_event_cmd_init(&cmd, buf, MAX_DYNEVENT_CMD_LEN);

	/*
	 * Define the kretprobe event.
	 */
	ret = kretprobe_event_gen_cmd_start(&cmd, "gen_kretprobe_test",
					    KPROBE_GEN_TEST_FUNC,
					    "$retval");
	if (ret)
		goto out;

	/*
	 * This actually creates the event.
	 */
	ret = kretprobe_event_gen_cmd_end(&cmd);
	if (ret)
		goto out;

	/*
	 * Now get the gen_kretprobe_test event file.  We need to
	 * prevent the instance and event from disappearing from
	 * underneath us, which trace_get_event_file() does (though in
	 * this case we're using the top-level instance which never
	 * goes away).
	 */
	gen_kretprobe_test = trace_get_event_file(NULL, "kprobes",
						  "gen_kretprobe_test");
	if (IS_ERR(gen_kretprobe_test)) {
		ret = PTR_ERR(gen_kretprobe_test);
		goto delete;
	}

	/* Enable the event or you won't see anything */
	ret = trace_array_set_clr_event(gen_kretprobe_test->tr,
					"kprobes", "gen_kretprobe_test", true);
	if (ret) {
		trace_put_event_file(gen_kretprobe_test);
		goto delete;
	}
 out:
	kfree(buf);
	return ret;
 delete:
	if (trace_event_file_is_valid(gen_kretprobe_test))
		gen_kretprobe_test = NULL;
	/* We got an error after creating the event, delete it */
	kprobe_event_delete("gen_kretprobe_test");
	goto out;
}

static int __init kprobe_event_gen_test_init(void)
{
	int ret;

	ret = test_gen_kprobe_cmd();
	if (ret)
		return ret;

	ret = test_gen_kretprobe_cmd();
	if (ret) {
		if (trace_event_file_is_valid(gen_kretprobe_test)) {
			WARN_ON(trace_array_set_clr_event(gen_kretprobe_test->tr,
							  "kprobes",
							  "gen_kretprobe_test", false));
			trace_put_event_file(gen_kretprobe_test);
		}
		WARN_ON(kprobe_event_delete("gen_kretprobe_test"));
	}

	return ret;
}

static void __exit kprobe_event_gen_test_exit(void)
{
	if (trace_event_file_is_valid(gen_kprobe_test)) {
		/* Disable the event or you can't remove it */
		WARN_ON(trace_array_set_clr_event(gen_kprobe_test->tr,
						  "kprobes",
						  "gen_kprobe_test", false));

		/* Now give the file and instance back */
		trace_put_event_file(gen_kprobe_test);
	}


	/* Now unregister and free the event */
	WARN_ON(kprobe_event_delete("gen_kprobe_test"));

	if (trace_event_file_is_valid(gen_kretprobe_test)) {
		/* Disable the event or you can't remove it */
		WARN_ON(trace_array_set_clr_event(gen_kretprobe_test->tr,
						  "kprobes",
						  "gen_kretprobe_test", false));

		/* Now give the file and instance back */
		trace_put_event_file(gen_kretprobe_test);
	}


	/* Now unregister and free the event */
	WARN_ON(kprobe_event_delete("gen_kretprobe_test"));
}

module_init(kprobe_event_gen_test_init)
module_exit(kprobe_event_gen_test_exit)

MODULE_AUTHOR("Tom Zanussi");
MODULE_DESCRIPTION("kprobe event generation test");
MODULE_LICENSE("GPL v2");