dsp: pdr: Add pldm_entity_association_tree_copy_root_check()

Allocations cannot be treated as infallible. Ensure that copying
the entity association tree signals failure to the caller along with
cleaning up after itself to avoid leaking memory.

Change-Id: Icfd255b45530e42a6a3a0a3205e665eed53708d1
Signed-off-by: Andrew Jeffery <andrew@codeconstruct.com.au>
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7dd9e98..51912fb 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -24,6 +24,7 @@
 3. oem: meta: Add encode_oem_meta_file_io_read_resp()
 4. pdr: Add pldm_entity_association_pdr_remove_contained_entity()
 5. pdr: Add pldm_pdr_remove_fru_record_set_by_rsi()
+6. pldm_entity_association_tree_copy_root_check()
 
 ### Changed
 
@@ -55,6 +56,12 @@
    - Change parameters from individual pointers to a struct.
    - Check the length provided in the message won't exceed the buffer.
 
+2. pldm_entity_association_tree_copy_root()
+
+   The implementation allocates, but gives no indication to the caller if an
+   allocation (and hence the copy) has failed. Users should migrate to
+   pldm_entity_association_tree_copy_root_check().
+
 ### Removed
 
 1. Deprecated functions with the `_check` suffix
diff --git a/include/libpldm/pdr.h b/include/libpldm/pdr.h
index a028e7b..f67cdad 100644
--- a/include/libpldm/pdr.h
+++ b/include/libpldm/pdr.h
@@ -594,6 +594,19 @@
 	pldm_entity_association_tree *org_tree,
 	pldm_entity_association_tree *new_tree);
 
+/** @brief Create a copy of an existing entity association tree
+ *
+ *  @param[in] org_tree - pointer to source tree
+ *  @param[in/out] new_tree - pointer to destination tree
+ *
+ *  @return 0 if the entity association tree was copied, -EINVAL if the argument
+ *          values are invalid, or -ENOMEM if memory required for the copy
+ *          cannot be allocated.
+ */
+int pldm_entity_association_tree_copy_root_check(
+	pldm_entity_association_tree *org_tree,
+	pldm_entity_association_tree *new_tree);
+
 /** @brief Destroy all the nodes of the entity association tree
  *
  *  @param[in] tree - pointer to entity association tree
diff --git a/src/dsp/pdr.c b/src/dsp/pdr.c
index 0f73f02..670e1ed 100644
--- a/src/dsp/pdr.c
+++ b/src/dsp/pdr.c
@@ -1213,26 +1213,49 @@
 	return node;
 }
 
-static void entity_association_tree_copy(pldm_entity_node *org_node,
-					 pldm_entity_node **new_node)
+static int entity_association_tree_copy(pldm_entity_node *org_node,
+					pldm_entity_node **new_node)
 {
+	int rc;
+
 	if (org_node == NULL) {
-		return;
+		return 0;
 	}
+
 	*new_node = malloc(sizeof(pldm_entity_node));
+	if (!*new_node) {
+		return -ENOMEM;
+	}
+
 	(*new_node)->parent = org_node->parent;
 	(*new_node)->entity = org_node->entity;
 	(*new_node)->association_type = org_node->association_type;
 	(*new_node)->remote_container_id = org_node->remote_container_id;
 	(*new_node)->first_child = NULL;
 	(*new_node)->next_sibling = NULL;
-	entity_association_tree_copy(org_node->first_child,
-				     &((*new_node)->first_child));
-	entity_association_tree_copy(org_node->next_sibling,
-				     &((*new_node)->next_sibling));
+
+	rc = entity_association_tree_copy(org_node->first_child,
+					  &((*new_node)->first_child));
+	if (rc) {
+		goto cleanup;
+	}
+
+	rc = entity_association_tree_copy(org_node->next_sibling,
+					  &((*new_node)->next_sibling));
+	if (rc) {
+		entity_association_tree_destroy((*new_node)->first_child);
+		goto cleanup;
+	}
+
+	return 0;
+
+cleanup:
+	free(*new_node);
+	*new_node = NULL;
+	return rc;
 }
 
-LIBPLDM_ABI_STABLE
+LIBPLDM_ABI_DEPRECATED
 void pldm_entity_association_tree_copy_root(
 	pldm_entity_association_tree *org_tree,
 	pldm_entity_association_tree *new_tree)
@@ -1244,6 +1267,19 @@
 	entity_association_tree_copy(org_tree->root, &(new_tree->root));
 }
 
+LIBPLDM_ABI_TESTING
+int pldm_entity_association_tree_copy_root_check(
+	pldm_entity_association_tree *org_tree,
+	pldm_entity_association_tree *new_tree)
+{
+	if (!org_tree || !new_tree) {
+		return -EINVAL;
+	}
+
+	new_tree->last_used_container_id = org_tree->last_used_container_id;
+	return entity_association_tree_copy(org_tree->root, &(new_tree->root));
+}
+
 LIBPLDM_ABI_STABLE
 void pldm_entity_association_tree_destroy_root(
 	pldm_entity_association_tree *tree)
diff --git a/tests/dsp/pdr.cpp b/tests/dsp/pdr.cpp
index 726bdfb..b8e224e 100644
--- a/tests/dsp/pdr.cpp
+++ b/tests/dsp/pdr.cpp
@@ -1464,9 +1464,12 @@
     pldm_entity_association_tree_destroy(tree);
 }
 
+#ifdef LIBPLDM_API_TESTING
 TEST(EntityAssociationPDR, testCopyTree)
 {
     pldm_entity entities[4]{};
+    int rc;
+
     entities[0].entity_type = 1;
     entities[1].entity_type = 2;
     entities[2].entity_type = 2;
@@ -1492,7 +1495,8 @@
     pldm_entity_association_tree_visit(orgTree, &orgOut, &orgNum);
     EXPECT_EQ(orgNum, 4u);
 
-    pldm_entity_association_tree_copy_root(orgTree, newTree);
+    rc = pldm_entity_association_tree_copy_root_check(orgTree, newTree);
+    ASSERT_EQ(rc, 0);
     size_t newNum{};
     pldm_entity* newOut = nullptr;
     pldm_entity_association_tree_visit(newTree, &newOut, &newNum);
@@ -1505,6 +1509,7 @@
     pldm_entity_association_tree_destroy(orgTree);
     pldm_entity_association_tree_destroy(newTree);
 }
+#endif
 
 TEST(EntityAssociationPDR, testExtract)
 {