[PATCH] maple_tree: Add a mas_destroy() call to mas_expected_entries() failure path

[Liam Howlett <liam.howlett@oracle.com>]

In an exceedingly rare case, there is a possibility that allocating all
of the nodes may fail in a way that the maple state is left with some
allocations completed.  This would happen if the single allocation
succeeds but bulk allocation fails, or if multiple bulk allocations are
required and somewhere along the way one fails.  The partial return is
already cleaned up, but the successful allocations will remain in the
maple state.  When this happens, mas_expected_entries() may leak memory.
Fix this by moving mas_destroy() above mas_expected_entries() and add a
call to mas_destroy() to clear out all allocated memory.

Signed-off-by: Liam R. Howlett <Liam.Howlett@oracle.com>
---
 lib/maple_tree.c | 89 ++++++++++++++++++++++++------------------------
 1 file changed, 45 insertions(+), 44 deletions(-)

diff --git a/lib/maple_tree.c b/lib/maple_tree.c
index 4c383c780162..d00ad50b258e 100644
--- a/lib/maple_tree.c
+++ b/lib/maple_tree.c
@@ -5717,6 +5717,50 @@ int mas_preallocate(struct ma_state *mas, void *entry, gfp_t gfp)
 	return ret;
 }
 
+/*
+ * mas_destroy() - destroy a maple state.
+ * @mas: The maple state
+ *
+ * Upon completion, check the left-most node and rebalance against the node to
+ * the right if necessary.  Frees any allocated nodes associated with this maple
+ * state.
+ */
+void mas_destroy(struct ma_state *mas)
+{
+	struct maple_alloc *node;
+
+	/*
+	 * When using mas_for_each() to insert an expected number of elements,
+	 * it is possible that the number inserted is less than the expected
+	 * number.  To fix an invalid final node, a check is performed here to
+	 * rebalance the previous node with the final node.
+	 */
+	if (mas->mas_flags & MA_STATE_REBALANCE) {
+		unsigned char end;
+
+		if (mas_is_start(mas))
+			mas_start(mas);
+
+		mtree_range_walk(mas);
+		end = mas_data_end(mas) + 1;
+		if (end < mt_min_slot_count(mas->node) - 1)
+			mas_destroy_rebalance(mas, end);
+
+		mas->mas_flags &= ~MA_STATE_REBALANCE;
+	}
+	mas->mas_flags &= ~MA_STATE_BULK;
+
+	while (mas->alloc && !((unsigned long)mas->alloc & 0x1)) {
+		node = mas->alloc;
+		mas->alloc = node->slot[0];
+		if (node->node_count > 0)
+			mt_free_bulk(node->node_count,
+				     (void __rcu **)&node->slot[1]);
+		kmem_cache_free(maple_node_cache, node);
+	}
+	mas->alloc = NULL;
+}
+
 /*
  * mas_expected_entries() - Set the expected number of entries that will be inserted.
  * @mas: The maple state
@@ -5770,54 +5814,11 @@ int mas_expected_entries(struct ma_state *mas, unsigned long nr_entries)
 
 	ret = xa_err(mas->node);
 	mas->node = enode;
+	mas_destroy(mas);
 	return ret;
 
 }
 
-/*
- * mas_destroy() - destroy a maple state.
- * @mas: The maple state
- *
- * Upon completion, check the left-most node and rebalance against the node to
- * the right if necessary.  Frees any allocated nodes associated with this maple
- * state.
- */
-void mas_destroy(struct ma_state *mas)
-{
-	struct maple_alloc *node;
-
-	/*
-	 * When using mas_for_each() to insert an expected number of elements,
-	 * it is possible that the number inserted is less than the expected
-	 * number.  To fix an invalid final node, a check is performed here to
-	 * rebalance the previous node with the final node.
-	 */
-	if (mas->mas_flags & MA_STATE_REBALANCE) {
-		unsigned char end;
-
-		if (mas_is_start(mas))
-			mas_start(mas);
-
-		mtree_range_walk(mas);
-		end = mas_data_end(mas) + 1;
-		if (end < mt_min_slot_count(mas->node) - 1)
-			mas_destroy_rebalance(mas, end);
-
-		mas->mas_flags &= ~MA_STATE_REBALANCE;
-	}
-	mas->mas_flags &= ~MA_STATE_BULK;
-
-	while (mas->alloc && !((unsigned long)mas->alloc & 0x1)) {
-		node = mas->alloc;
-		mas->alloc = node->slot[0];
-		if (node->node_count > 0)
-			mt_free_bulk(node->node_count,
-				     (void __rcu **)&node->slot[1]);
-		kmem_cache_free(maple_node_cache, node);
-	}
-	mas->alloc = NULL;
-}
-
 /**
  * mas_next() - Get the next entry.
  * @mas: The maple state
-- 
2.35.1