[PATCH RFC bpf-next 20/20] trait: register traits in benchmarks and tests

[arthur@xxxxxxxxxxxxxxx]

From: Arthur Fabre <afabre@xxxxxxxxxxxxxx>

Otherwise the verifier rejects the programs now.

Signed-off-by: Arthur Fabre <afabre@xxxxxxxxxxxxxx>
---
 tools/testing/selftests/bpf/bench.c                |  3 ++
 tools/testing/selftests/bpf/bench.h                |  1 +
 .../selftests/bpf/benchs/bench_xdp_traits.c        | 33 +++++++++++++++++++++-
 .../testing/selftests/bpf/prog_tests/xdp_traits.c  | 18 +++++++++++-
 4 files changed, 53 insertions(+), 2 deletions(-)

diff --git a/tools/testing/selftests/bpf/bench.c b/tools/testing/selftests/bpf/bench.c
index 4678b928fc6ad2f0a870a25d9b10c75a1f6d77ba..0961cb71ddf1d682ca61e512e9f4b2df3606747c 100644
--- a/tools/testing/selftests/bpf/bench.c
+++ b/tools/testing/selftests/bpf/bench.c
@@ -752,5 +752,8 @@ int main(int argc, char **argv)
 		bench->report_final(state.results + env.warmup_sec,
 				    state.res_cnt - env.warmup_sec);
 
+	if (bench->cleanup)
+		bench->cleanup();
+
 	return 0;
 }
diff --git a/tools/testing/selftests/bpf/bench.h b/tools/testing/selftests/bpf/bench.h
index 005c401b3e2275030d1d489cd77423cb1fb652ad..6a94af31df17e7cc35bb1432c4b205eb96b631f2 100644
--- a/tools/testing/selftests/bpf/bench.h
+++ b/tools/testing/selftests/bpf/bench.h
@@ -58,6 +58,7 @@ struct bench {
 	void (*measure)(struct bench_res* res);
 	void (*report_progress)(int iter, struct bench_res* res, long delta_ns);
 	void (*report_final)(struct bench_res res[], int res_cnt);
+	void (*cleanup)(void);
 };
 
 struct counter {
diff --git a/tools/testing/selftests/bpf/benchs/bench_xdp_traits.c b/tools/testing/selftests/bpf/benchs/bench_xdp_traits.c
index 0fbcd49edd825f53e6957319d3f05efc218dfb02..b6fa2d8f2504dd9c35f8fb9dc1f1099b55a55ac6 100644
--- a/tools/testing/selftests/bpf/benchs/bench_xdp_traits.c
+++ b/tools/testing/selftests/bpf/benchs/bench_xdp_traits.c
@@ -57,7 +57,18 @@ static void trait_validate(void)
 
 static void trait_setup(void)
 {
-	int err;
+	int err, i, key;
+	union bpf_attr attr;
+
+	/* Register all keys so we can use them all. */
+	bzero(&attr, sizeof(attr));
+	for (i = 0; i < 64; i++) {
+		key = syscall(__NR_bpf, BPF_REGISTER_TRAIT, &attr, sizeof(attr));
+		if (key < 0) {
+			fprintf(stderr, "couldn't register trait: %d\n", key);
+			exit(1);
+		}
+	}
 
 	setup_libbpf();
 
@@ -77,6 +88,23 @@ static void trait_setup(void)
 	}
 }
 
+static void trait_cleanup(void)
+{
+	int err, i;
+	union bpf_attr attr;
+
+	/* Unregister all keys so we can run again. */
+	bzero(&attr, sizeof(attr));
+	for (i = 0; i < 64; i++) {
+		attr.unregister_trait.trait = i;
+		err = syscall(__NR_bpf, BPF_UNREGISTER_TRAIT, &attr, sizeof(attr));
+		if (err < 0) {
+			fprintf(stderr, "couldn't unregister trait %d: %d\n", i, err);
+			exit(1);
+		}
+	}
+}
+
 static void trait_get_setup(void)
 {
 	trait_setup();
@@ -135,6 +163,7 @@ const struct bench bench_xdp_trait_get = {
 	.measure = trait_measure,
 	.report_progress = ops_report_progress,
 	.report_final = ops_report_final,
+	.cleanup = trait_cleanup,
 };
 
 const struct bench bench_xdp_trait_set = {
@@ -146,6 +175,7 @@ const struct bench bench_xdp_trait_set = {
 	.measure = trait_measure,
 	.report_progress = ops_report_progress,
 	.report_final = ops_report_final,
+	.cleanup = trait_cleanup,
 };
 
 const struct bench bench_xdp_trait_move = {
@@ -157,4 +187,5 @@ const struct bench bench_xdp_trait_move = {
 	.measure = trait_measure,
 	.report_progress = ops_report_progress,
 	.report_final = ops_report_final,
+	.cleanup = trait_cleanup,
 };
diff --git a/tools/testing/selftests/bpf/prog_tests/xdp_traits.c b/tools/testing/selftests/bpf/prog_tests/xdp_traits.c
index 4175b28d45e91e82435e646e5edd783980d5fe70..1c1eff235a6159d377a5e8b9e0a4d956c4540e8e 100644
--- a/tools/testing/selftests/bpf/prog_tests/xdp_traits.c
+++ b/tools/testing/selftests/bpf/prog_tests/xdp_traits.c
@@ -6,8 +6,16 @@ static void _test_xdp_traits(void)
 {
 	const char *file = "./test_xdp_traits.bpf.o";
 	struct bpf_object *obj;
-	int err, prog_fd;
+	int err, prog_fd, i, key;
 	char buf[128];
+	union bpf_attr attr;
+
+	/* Register all keys so we can use them all. */
+	bzero(&attr, sizeof(attr));
+	for (i = 0; i < 64; i++) {
+		key = syscall(__NR_bpf, BPF_REGISTER_TRAIT, &attr, sizeof(attr));
+		ASSERT_OK_FD(key, "test_xdp_traits");
+	}
 
 	LIBBPF_OPTS(bpf_test_run_opts, topts,
 		.data_in = &pkt_v4,
@@ -26,6 +34,14 @@ static void _test_xdp_traits(void)
 	ASSERT_EQ(topts.retval, XDP_PASS, "retval");
 
 	bpf_object__close(obj);
+
+	/* Unregister all keys so we can run again. */
+	bzero(&attr, sizeof(attr));
+	for (i = 0; i < 64; i++) {
+		attr.unregister_trait.trait = i;
+		err = syscall(__NR_bpf, BPF_UNREGISTER_TRAIT, &attr, sizeof(attr));
+		ASSERT_OK(err, "test_xdp_traits");
+	}
 }
 
 void test_xdp_traits(void)

-- 
2.43.0