From: Masami Hiramatsu (Google) <mhiramat@kernel.org>

Add test cases for checking recursion level must be less than 2.
This ensures that fprobe can prevent infinite recursive call by
itself.

Signed-off-by: Masami Hiramatsu (Google) <mhiramat@kernel.org>
---
 lib/tests/test_fprobe.c |   92 ++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 90 insertions(+), 2 deletions(-)

diff --git a/lib/tests/test_fprobe.c b/lib/tests/test_fprobe.c
index cf92111b5c79..0022da242a83 100644
--- a/lib/tests/test_fprobe.c
+++ b/lib/tests/test_fprobe.c
@@ -17,8 +17,10 @@ static u32 rand1, entry_val, exit_val;
 /* Use indirect calls to avoid inlining the target functions */
 static u32 (*target)(u32 value);
 static u32 (*target2)(u32 value);
+static u32 (*target3)(u32 value);
 static unsigned long target_ip;
 static unsigned long target2_ip;
+static unsigned long target3_ip;
 static int entry_return_value;
 
 static noinline u32 fprobe_selftest_target(u32 value)
@@ -31,13 +33,18 @@ static noinline u32 fprobe_selftest_target2(u32 value)
 	return (value / div_factor) + 1;
 }
 
+static noinline u32 fprobe_selftest_target3(u32 value)
+{
+	return (value / div_factor) + 2;
+}
+
 static notrace int fp_entry_handler(struct fprobe *fp, unsigned long ip,
 				    unsigned long ret_ip,
 				    struct ftrace_regs *fregs, void *data)
 {
 	KUNIT_EXPECT_FALSE(current_test, preemptible());
 	/* This can be called on the fprobe_selftest_target and the fprobe_selftest_target2 */
-	if (ip != target_ip)
+	if (ip != target_ip && ip != target3_ip)
 		KUNIT_EXPECT_EQ(current_test, ip, target2_ip);
 	entry_val = (rand1 / div_factor);
 	if (fp->entry_data_size) {
@@ -190,6 +197,84 @@ static void test_fprobe_skip(struct kunit *test)
 	KUNIT_EXPECT_EQ(test, 0, unregister_fprobe(&fp));
 }
 
+static int recur_count;
+static int recur_high;
+
+static notrace int fp_recur_entry_handler(struct fprobe *fp, unsigned long ip,
+					  unsigned long ret_ip,
+					  struct ftrace_regs *fregs, void *data)
+{
+	/* fgraph/fprobe allows one-level nesting. */
+	if (recur_count > 1)
+		return 0;
+
+	recur_count++;
+	if (recur_count > recur_high)
+		recur_high = recur_count;
+	target3(rand1);
+	recur_count--;
+
+	return 0;
+}
+
+static notrace void fp_recur_exit_handler(struct fprobe *fp, unsigned long ip,
+					  unsigned long ret_ip,
+					  struct ftrace_regs *fregs, void *data)
+{
+	/* fgraph/fprobe allows one-level nesting. */
+	if (recur_count > 1)
+		return;
+
+	recur_count++;
+	if (recur_count > recur_high)
+		recur_high = recur_count;
+	target3(rand1);
+	recur_count--;
+}
+
+static void test_fprobe_recursion_entry(struct kunit *test)
+{
+	struct fprobe fp = {
+		.entry_handler = fp_recur_entry_handler,
+	};
+	struct fprobe fp3 = {
+		.entry_handler = fp_recur_entry_handler,
+	};
+
+	current_test = test;
+	recur_count = 0;
+	recur_high = 0;
+	KUNIT_EXPECT_EQ(test, 0, register_fprobe(&fp, "fprobe_selftest_target", NULL));
+	KUNIT_EXPECT_EQ(test, 0, register_fprobe(&fp3, "fprobe_selftest_target3", NULL));
+
+	target(rand1);
+
+	KUNIT_EXPECT_LE(test, recur_high, 2);
+	KUNIT_EXPECT_EQ(test, 0, unregister_fprobe(&fp3));
+	KUNIT_EXPECT_EQ(test, 0, unregister_fprobe(&fp));
+}
+
+static void test_fprobe_recursion_exit(struct kunit *test)
+{
+	struct fprobe fp = {
+		.exit_handler = fp_recur_exit_handler,
+	};
+	struct fprobe fp3 = {
+		.exit_handler = fp_recur_exit_handler,
+	};
+
+	current_test = test;
+	recur_count = 0;
+	KUNIT_EXPECT_EQ(test, 0, register_fprobe(&fp, "fprobe_selftest_target", NULL));
+	KUNIT_EXPECT_EQ(test, 0, register_fprobe(&fp3, "fprobe_selftest_target3", NULL));
+
+	target(rand1);
+
+	KUNIT_EXPECT_LE(test, recur_high, 2);
+	KUNIT_EXPECT_EQ(test, 0, unregister_fprobe(&fp3));
+	KUNIT_EXPECT_EQ(test, 0, unregister_fprobe(&fp));
+}
+
 static unsigned long get_ftrace_location(void *func)
 {
 	unsigned long size, addr = (unsigned long)func;
@@ -205,8 +290,10 @@ static int fprobe_test_init(struct kunit *test)
 	rand1 = get_random_u32_above(div_factor);
 	target = fprobe_selftest_target;
 	target2 = fprobe_selftest_target2;
+	target3 = fprobe_selftest_target3;
 	target_ip = get_ftrace_location(target);
 	target2_ip = get_ftrace_location(target2);
+	target3_ip = get_ftrace_location(target3);
 
 	return 0;
 }
@@ -217,6 +304,8 @@ static struct kunit_case fprobe_testcases[] = {
 	KUNIT_CASE(test_fprobe_syms),
 	KUNIT_CASE(test_fprobe_data),
 	KUNIT_CASE(test_fprobe_skip),
+	KUNIT_CASE(test_fprobe_recursion_entry),
+	KUNIT_CASE(test_fprobe_recursion_exit),
 	{}
 };
 
@@ -227,4 +316,3 @@ static struct kunit_suite fprobe_test_suite = {
 };
 
 kunit_test_suites(&fprobe_test_suite);
-