[RFC PATCH 4/4] kho: Add test cases for versioned subtrees

[Logan Odell <loganodell@google.com>]

Add new test cases to test_kho.c to verify saving and restoring
multiple versions of a subtree. The test preserves versions 1 and 2
of "kho_test_ver" and retrieves them, verifying their contents.
It also ensures no memory leaks occur during the restore process.

Signed-off-by: Logan Odell <loganodell@google.com>
---
 lib/test_kho.c | 266 +++++++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 245 insertions(+), 21 deletions(-)

diff --git a/lib/test_kho.c b/lib/test_kho.c
index 7ef9e4061869..db531da63576 100644
--- a/lib/test_kho.c
+++ b/lib/test_kho.c
@@ -310,52 +310,276 @@ static int kho_test_restore(phys_addr_t fdt_phys)
 	return 0;
 }
 
+static void kho_test_cleanup(void)
+{
+	/* unpreserve and free the data stored in folios */
+	kho_test_unpreserve_data(&kho_test_state);
+	for (int i = 0; i < kho_test_state.nr_folios; i++)
+		folio_put(kho_test_state.folios[i]);
+
+	kvfree(kho_test_state.folios);
+
+	/* Unpreserve and release the FDT folio */
+	kho_unpreserve_folio(kho_test_state.fdt);
+	folio_put(kho_test_state.fdt);
+}
+
+#define KHO_TEST_VER_NAME	"kho_test_ver"
+#define KHO_TEST_VER_1		1
+#define KHO_TEST_VER_2		2
+#define KHO_TEST_VER_MAGIC_1	0x11111111
+#define KHO_TEST_VER_MAGIC_2	0x22222222
+
+struct kho_test_versioned_state {
+	struct folio *fdt1;
+	struct folio *fdt2;
+};
+
+static struct kho_test_versioned_state kho_test_ver_state;
+
+static int kho_test_prepare_ver_fdt(struct folio *folio, int version, unsigned int magic)
+{
+	void *fdt = folio_address(folio);
+	int err;
+
+	err = fdt_create(fdt, PAGE_SIZE);
+	err |= fdt_finish_reservemap(fdt);
+	err |= fdt_begin_node(fdt, "");
+	err |= fdt_property(fdt, "compatible", KHO_TEST_COMPAT, sizeof(KHO_TEST_COMPAT));
+	err |= fdt_property(fdt, "magic", &magic, sizeof(magic));
+	err |= fdt_property(fdt, "version", &version, sizeof(version));
+	err |= fdt_end_node(fdt);
+	err |= fdt_finish(fdt);
+
+	return err;
+}
+
+static int kho_test_ver_save(void)
+{
+	struct kho_test_versioned_state *state = &kho_test_ver_state;
+	int err;
+
+	state->fdt1 = folio_alloc(GFP_KERNEL, 0);
+	if (!state->fdt1)
+		return -ENOMEM;
+
+	state->fdt2 = folio_alloc(GFP_KERNEL, 0);
+	if (!state->fdt2) {
+		err = -ENOMEM;
+		goto err_free_fdt1;
+	}
+
+	err = kho_preserve_folio(state->fdt1);
+	if (err)
+		goto err_free_fdt2;
+
+	err = kho_preserve_folio(state->fdt2);
+	if (err)
+		goto err_unpreserve_fdt1;
+
+	err = kho_test_prepare_ver_fdt(state->fdt1, KHO_TEST_VER_1, KHO_TEST_VER_MAGIC_1);
+	if (err)
+		goto err_unpreserve_fdt2;
+
+	err = kho_test_prepare_ver_fdt(state->fdt2, KHO_TEST_VER_2, KHO_TEST_VER_MAGIC_2);
+	if (err)
+		goto err_unpreserve_fdt2;
+
+	err = kho_add_subtree_version(KHO_TEST_VER_NAME, KHO_TEST_VER_1,
+				      folio_address(state->fdt1));
+	if (err)
+		goto err_unpreserve_fdt2;
+
+	err = kho_add_subtree_version(KHO_TEST_VER_NAME, KHO_TEST_VER_2,
+				      folio_address(state->fdt2));
+	if (err)
+		goto err_remove_ver1;
+
+	pr_info("KHO versioned test save succeeded\n");
+	return 0;
+
+err_remove_ver1:
+	kho_remove_subtree_version(KHO_TEST_VER_NAME, KHO_TEST_VER_1, folio_address(state->fdt1));
+err_unpreserve_fdt2:
+	kho_unpreserve_folio(state->fdt2);
+err_unpreserve_fdt1:
+	kho_unpreserve_folio(state->fdt1);
+err_free_fdt2:
+	folio_put(state->fdt2);
+	state->fdt2 = NULL;
+err_free_fdt1:
+	folio_put(state->fdt1);
+	state->fdt1 = NULL;
+	return err;
+}
+
+static int kho_test_ver_restore_one(phys_addr_t fdt_phys, int expected_version,
+				    unsigned int expected_magic)
+{
+	void *fdt = phys_to_virt(fdt_phys);
+	const unsigned int *magic;
+	const int *version;
+	int node, len;
+
+	node = fdt_path_offset(fdt, "/");
+	if (node < 0)
+		return -EINVAL;
+
+	if (fdt_node_check_compatible(fdt, node, KHO_TEST_COMPAT))
+		return -EINVAL;
+
+	magic = fdt_getprop(fdt, node, "magic", &len);
+	if (!magic || len != sizeof(*magic) || *magic != expected_magic)
+		return -EINVAL;
+
+	version = fdt_getprop(fdt, node, "version", &len);
+	if (!version || len != sizeof(*version) || *version != expected_version)
+		return -EINVAL;
+
+	return 0;
+}
+
+static int kho_test_ver_restore(void)
+{
+	phys_addr_t fdt_phys1, fdt_phys2;
+	struct folio *fdt1 = NULL, *fdt2 = NULL;
+	int err;
+
+	err = kho_retrieve_subtree_version(KHO_TEST_VER_NAME, KHO_TEST_VER_1, &fdt_phys1);
+	if (err) {
+		if (err != -ENOENT)
+			pr_err("Failed to retrieve version 1: %d\n", err);
+		return err;
+	}
+
+	fdt1 = kho_restore_folio(fdt_phys1);
+	if (!fdt1) {
+		pr_err("Failed to restore FDT 1 folio\n");
+		return -EINVAL;
+	}
+
+	err = kho_retrieve_subtree_version(KHO_TEST_VER_NAME, KHO_TEST_VER_2, &fdt_phys2);
+	if (err) {
+		pr_err("Failed to retrieve version 2: %d\n", err);
+		goto out_put_fdt1;
+	}
+
+	fdt2 = kho_restore_folio(fdt_phys2);
+	if (!fdt2) {
+		pr_err("Failed to restore FDT 2 folio\n");
+		err = -EINVAL;
+		goto out_put_fdt1;
+	}
+
+	err = kho_test_ver_restore_one(fdt_phys1, KHO_TEST_VER_1, KHO_TEST_VER_MAGIC_1);
+	if (err) {
+		pr_err("Version 1 verification failed: %d\n", err);
+		goto out_put_fdt2;
+	}
+
+	err = kho_test_ver_restore_one(fdt_phys2, KHO_TEST_VER_2, KHO_TEST_VER_MAGIC_2);
+	if (err) {
+		pr_err("Version 2 verification failed: %d\n", err);
+		goto out_put_fdt2;
+	}
+
+	pr_info("KHO versioned test restore succeeded\n");
+
+out_put_fdt2:
+	folio_put(fdt2);
+out_put_fdt1:
+	folio_put(fdt1);
+	return err;
+}
+
+static void kho_test_ver_cleanup(void)
+{
+	struct kho_test_versioned_state *state = &kho_test_ver_state;
+
+	if (state->fdt1) {
+		kho_unpreserve_folio(state->fdt1);
+		folio_put(state->fdt1);
+		state->fdt1 = NULL;
+	}
+	if (state->fdt2) {
+		kho_unpreserve_folio(state->fdt2);
+		folio_put(state->fdt2);
+		state->fdt2 = NULL;
+	}
+}
+
+static void kho_test_ver_exit(void)
+{
+	struct kho_test_versioned_state *state = &kho_test_ver_state;
+
+	if (state->fdt1)
+		kho_remove_subtree_version(KHO_TEST_VER_NAME, KHO_TEST_VER_1,
+					   folio_address(state->fdt1));
+	if (state->fdt2)
+		kho_remove_subtree_version(KHO_TEST_VER_NAME, KHO_TEST_VER_2,
+					   folio_address(state->fdt2));
+	kho_test_ver_cleanup();
+}
+
 static int __init kho_test_init(void)
 {
 	phys_addr_t fdt_phys;
-	int err;
+	int err, ver_err;
+	bool restored = false;
 
 	if (!kho_is_enabled())
 		return 0;
 
+	// Try to restore non-versioned
 	err = kho_retrieve_subtree(KHO_TEST_FDT, &fdt_phys);
 	if (!err) {
 		err = kho_test_restore(fdt_phys);
-		if (err)
+		if (err) {
 			pr_err("KHO restore failed\n");
-		else
-			pr_info("KHO restore succeeded\n");
-
+			return err;
+		}
+		pr_info("KHO restore succeeded\n");
+		restored = true;
+	} else if (err != -ENOENT) {
+		pr_warn("failed to retrieve %s FDT: %d\n", KHO_TEST_FDT, err);
 		return err;
 	}
 
-	if (err != -ENOENT) {
-		pr_warn("failed to retrieve %s FDT: %d\n", KHO_TEST_FDT, err);
-		return err;
+	// Try to restore versioned
+	ver_err = kho_test_ver_restore();
+	if (!ver_err) {
+		restored = true;
+	} else if (ver_err != -ENOENT) {
+		pr_warn("KHO versioned restore failed: %d\n", ver_err);
+		return ver_err;
 	}
 
-	return kho_test_save();
-}
-module_init(kho_test_init);
+	if (restored)
+		return 0;
 
-static void kho_test_cleanup(void)
-{
-	/* unpreserve and free the data stored in folios */
-	kho_test_unpreserve_data(&kho_test_state);
-	for (int i = 0; i < kho_test_state.nr_folios; i++)
-		folio_put(kho_test_state.folios[i]);
+	// If we got here, it means both returned -ENOENT (first boot).
+	// So we save both.
 
-	kvfree(kho_test_state.folios);
+	err = kho_test_save();
+	if (err)
+		return err;
 
-	/* Unpreserve and release the FDT folio */
-	kho_unpreserve_folio(kho_test_state.fdt);
-	folio_put(kho_test_state.fdt);
+	ver_err = kho_test_ver_save();
+	if (ver_err) {
+		kho_remove_subtree(folio_address(kho_test_state.fdt));
+		kho_test_cleanup();
+		return ver_err;
+	}
+
+	return 0;
 }
+module_init(kho_test_init);
 
 static void __exit kho_test_exit(void)
 {
 	kho_remove_subtree(folio_address(kho_test_state.fdt));
 	kho_test_cleanup();
+	kho_test_ver_exit();
 }
 module_exit(kho_test_exit);
 
-- 
2.55.0.508.g3f0d502094-goog