[RFC PATCH 4/4] selftests/bpf: Test changing KV store value layout

[Amery Hung <ameryhung@gmail.com>]

While it is not the most ideal way I imagine how the KV store to be
used. The test tries to show upgrading a bpf program with a change in
the definition of a structure. If using a structure for the value, while
adding a new member is easy, it is less clear how removing or changing
member layout would work.

Signed-off-by: Amery Hung <ameryhung@gmail.com>
---
 .../bpf/prog_tests/test_uptr_kv_store.c       | 77 +++++++++++++++++++
 .../selftests/bpf/progs/test_uptr_kv_store.c  |  9 +++
 .../bpf/progs/test_uptr_kv_store_v1.c         | 46 +++++++++++
 .../selftests/bpf/test_uptr_kv_store_common.h | 13 ++++
 4 files changed, 145 insertions(+)
 create mode 100644 tools/testing/selftests/bpf/progs/test_uptr_kv_store_v1.c

diff --git a/tools/testing/selftests/bpf/prog_tests/test_uptr_kv_store.c b/tools/testing/selftests/bpf/prog_tests/test_uptr_kv_store.c
index 2075b8e47972..2e86bb08b26e 100644
--- a/tools/testing/selftests/bpf/prog_tests/test_uptr_kv_store.c
+++ b/tools/testing/selftests/bpf/prog_tests/test_uptr_kv_store.c
@@ -3,6 +3,7 @@
 #include "uptr_kv_store.h"
 #include "test_uptr_kv_store_common.h"
 #include "test_uptr_kv_store.skel.h"
+#include "test_uptr_kv_store_v1.skel.h"
 
 static void test_uptr_kv_store_basic(void)
 {
@@ -70,8 +71,84 @@ static void test_uptr_kv_store_basic(void)
 	kv_store_close(kvs);
 }
 
+static void test_uptr_kv_store_change_value(void)
+{
+	int err, pid;
+	struct test_uptr_kv_store_v1 *skel_v1;
+	struct test_uptr_kv_store *skel;
+	struct test_struct_v1 val_v1;
+	struct test_struct val, *val_p;
+	struct kv_store *kvs = NULL;
+
+	skel = test_uptr_kv_store__open_and_load();
+	if (!ASSERT_OK_PTR(skel, "skel_open_and_load"))
+		return;
+
+	skel->bss->target_pid = -1;
+	err = test_uptr_kv_store__attach(skel);
+	if (!ASSERT_OK(err, "skel_attach"))
+		return;
+
+	pid = sys_gettid();
+	kvs = kv_store_init(pid, skel->maps.data_map, "/sys/fs/bpf/kv_store_data_map");
+
+	/* update key 0 to test_struct in user space */
+	val.a = 1;
+	val.b = 2;
+	err = kv_store_put(kvs, 0, &val, sizeof(val));
+	ASSERT_OK(err, "kv_store_put struct val");
+	val_p = kv_store_get(kvs, 0);
+	ASSERT_OK_PTR(val_p, "kv_store_get struct val");
+	ASSERT_EQ(val_p->a, val.a, "user space: check get val.a == put val.a");
+	ASSERT_EQ(val_p->b, val.b, "user space: check get val.b == put val.b");
+
+	/* lookup test_struct at key 0 in test_uptr_kv_store */
+	skel->bss->test_key = 0;
+	skel->bss->test_op = KVS_STRUCT_GET;
+	skel->bss->target_pid = pid;
+	sys_gettid();
+	skel->bss->target_pid = -1;
+	ASSERT_EQ(skel->bss->test_struct_val.a, val.a, "bpf: check get val.a == put val.a");
+	ASSERT_EQ(skel->bss->test_struct_val.b, val.b, "bpf: check get val.b == put val.b");
+
+	/* add a new field to test_struct */
+	err = kv_store_update_value_size(kvs, 0, sizeof(val_v1));
+	ASSERT_OK(err, "kv_store_update_value_size");
+
+	/* rollout a new version */
+	skel_v1 = test_uptr_kv_store_v1__open();
+	if (!ASSERT_OK_PTR(skel, "skel_open v1"))
+		goto out;
+
+	kv_store_data_map_set_reuse(kvs, skel_v1->maps.data_map);
+
+	err = test_uptr_kv_store_v1__load(skel_v1);
+	if (!ASSERT_OK(err, "skel_load v1"))
+		goto out;
+
+	skel_v1->bss->target_pid = -1;
+	err = test_uptr_kv_store_v1__attach(skel_v1);
+	if (!ASSERT_OK(err, "skel_attach v1"))
+		goto out;
+
+	/* lookup struct_key_0 in test_uptr_kv_store */
+	skel_v1->bss->test_key = 0;
+	skel_v1->bss->test_op = KVS_STRUCT_GET;
+	skel_v1->bss->target_pid = pid;
+	sys_gettid();
+	skel_v1->bss->target_pid = -1;
+
+	ASSERT_EQ(skel_v1->bss->test_struct_val.a, val.a, "bpf: check get val_v1.a == put val.a");
+	ASSERT_EQ(skel_v1->bss->test_struct_val.b, val.b, "bpf: check get val_v1.b == put val.b");
+
+out:
+	kv_store_close(kvs);
+}
+
 void test_uptr_kv_store(void)
 {
 	if (test__start_subtest("uptr_kv_store_basic"))
 		test_uptr_kv_store_basic();
+	if (test__start_subtest("uptr_kv_store_change_value"))
+		test_uptr_kv_store_change_value();
 }
diff --git a/tools/testing/selftests/bpf/progs/test_uptr_kv_store.c b/tools/testing/selftests/bpf/progs/test_uptr_kv_store.c
index b358cb7fb616..2ed993ab4b01 100644
--- a/tools/testing/selftests/bpf/progs/test_uptr_kv_store.c
+++ b/tools/testing/selftests/bpf/progs/test_uptr_kv_store.c
@@ -8,6 +8,7 @@ pid_t target_pid = 0;
 int test_op;
 int test_key;
 int test_int_val;
+struct test_struct test_struct_val;
 
 SEC("tp_btf/sys_enter")
 int on_enter(__u64 *ctx)
@@ -28,6 +29,14 @@ int on_enter(__u64 *ctx)
 	case KVS_INT_GET:
 		kv_store_get(data, test_key, &test_int_val, 4);
 		break;
+	case KVS_STRUCT_PUT:
+		kv_store_put(data, test_key, &test_struct_val,
+			     sizeof(test_struct_val));
+		break;
+	case KVS_STRUCT_GET:
+		kv_store_get(data, test_key, &test_struct_val,
+			     sizeof(test_struct_val));
+		break;
 	}
 
 	return 0;
diff --git a/tools/testing/selftests/bpf/progs/test_uptr_kv_store_v1.c b/tools/testing/selftests/bpf/progs/test_uptr_kv_store_v1.c
new file mode 100644
index 000000000000..e3dd11e7d11b
--- /dev/null
+++ b/tools/testing/selftests/bpf/progs/test_uptr_kv_store_v1.c
@@ -0,0 +1,46 @@
+#include <vmlinux.h>
+#include <bpf/bpf_helpers.h>
+
+#include "uptr_kv_store.h"
+#include "test_uptr_kv_store_common.h"
+
+pid_t target_pid = 0;
+int test_op;
+int test_key;
+int test_int_val;
+struct test_struct_v1 test_struct_val;
+
+SEC("tp_btf/sys_enter")
+int on_enter(__u64 *ctx)
+{
+	struct kv_store_data_map_value *data;
+	struct task_struct *task;
+
+	task = bpf_get_current_task_btf();
+	if (task->pid != target_pid)
+		return 0;
+
+	data = bpf_task_storage_get(&data_map, task, 0, 0);
+
+	switch (test_op) {
+	case KVS_INT_PUT:
+		kv_store_put(data, test_key, &test_int_val, 4);
+		break;
+	case KVS_INT_GET:
+		kv_store_get(data, test_key, &test_int_val, 4);
+		break;
+	case KVS_STRUCT_PUT:
+		kv_store_put(data, test_key, &test_struct_val,
+			     sizeof(test_struct_val));
+		break;
+	case KVS_STRUCT_GET:
+		kv_store_get(data, test_key, &test_struct_val,
+			     sizeof(test_struct_val));
+		break;
+	}
+
+	return 0;
+}
+
+char _license[] SEC("license") = "GPL";
+
diff --git a/tools/testing/selftests/bpf/test_uptr_kv_store_common.h b/tools/testing/selftests/bpf/test_uptr_kv_store_common.h
index ff7d010ed08f..db91e862789f 100644
--- a/tools/testing/selftests/bpf/test_uptr_kv_store_common.h
+++ b/tools/testing/selftests/bpf/test_uptr_kv_store_common.h
@@ -4,6 +4,19 @@
 enum test_kvs_op {
 	KVS_INT_GET,
 	KVS_INT_PUT,
+	KVS_STRUCT_GET,
+	KVS_STRUCT_PUT,
+};
+
+struct test_struct {
+	int a;
+	int b;
+};
+
+struct test_struct_v1 {
+	int a;
+	int b;
+	int c;
 };
 
 #endif
-- 
2.47.1