The pa_count field in struct ext4_prealloc_space is
reference counter for preallocation descriptors.
Converting it from atomic_t to refcount_t provides
underflow and overflow protection.

This was found using Coccinelle script
atomic_as_refcounter.cocci

Testing was done on a QEMU x86 VM revealing a
pre-existing race condition leading to a
use-after-free.

Signed-off-by: rafad900 <rafad900@gmail.com>
---
 fs/ext4/mballoc.c | 22 +++++++++++-----------
 fs/ext4/mballoc.h |  2 +-
 2 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/fs/ext4/mballoc.c b/fs/ext4/mballoc.c
index ed1bd00e11cd..59b3d4abaab0 100644
--- a/fs/ext4/mballoc.c
+++ b/fs/ext4/mballoc.c
@@ -4823,7 +4823,7 @@ ext4_mb_check_group_pa(ext4_fsblk_t goal_block,
 	ext4_fsblk_t cur_distance, new_distance;
 
 	if (cpa == NULL) {
-		atomic_inc(&pa->pa_count);
+		refcount_inc(&pa->pa_count);
 		return pa;
 	}
 	cur_distance = abs(goal_block - cpa->pa_pstart);
@@ -4833,8 +4833,8 @@ ext4_mb_check_group_pa(ext4_fsblk_t goal_block,
 		return cpa;
 
 	/* drop the previous reference */
-	atomic_dec(&cpa->pa_count);
-	atomic_inc(&pa->pa_count);
+	refcount_dec(&cpa->pa_count);
+	refcount_inc(&pa->pa_count);
 	return pa;
 }
 
@@ -4993,7 +4993,7 @@ ext4_mb_use_preallocated(struct ext4_allocation_context *ac)
 	}
 
 	if (tmp_pa->pa_free && likely(ext4_mb_pa_goal_check(ac, tmp_pa))) {
-		atomic_inc(&tmp_pa->pa_count);
+		refcount_inc(&tmp_pa->pa_count);
 		ext4_mb_use_inode_pa(ac, tmp_pa);
 		spin_unlock(&tmp_pa->pa_lock);
 		read_unlock(&ei->i_prealloc_lock);
@@ -5139,7 +5139,7 @@ static void ext4_mb_mark_pa_deleted(struct super_block *sb,
 static inline void ext4_mb_pa_free(struct ext4_prealloc_space *pa)
 {
 	BUG_ON(!pa);
-	BUG_ON(atomic_read(&pa->pa_count));
+	BUG_ON(refcount_read(&pa->pa_count));
 	BUG_ON(pa->pa_deleted == 0);
 	kmem_cache_free(ext4_pspace_cachep, pa);
 }
@@ -5165,7 +5165,7 @@ static void ext4_mb_put_pa(struct ext4_allocation_context *ac,
 
 	/* in this short window concurrent discard can set pa_deleted */
 	spin_lock(&pa->pa_lock);
-	if (!atomic_dec_and_test(&pa->pa_count) || pa->pa_free != 0) {
+	if (!refcount_dec_and_test(&pa->pa_count) || pa->pa_free != 0) {
 		spin_unlock(&pa->pa_lock);
 		return;
 	}
@@ -5535,7 +5535,7 @@ ext4_mb_discard_group_preallocations(struct super_block *sb,
 	list_for_each_entry_safe(pa, tmp,
 				&grp->bb_prealloc_list, pa_group_list) {
 		spin_lock(&pa->pa_lock);
-		if (atomic_read(&pa->pa_count)) {
+		if (refcount_read(&pa->pa_count)) {
 			spin_unlock(&pa->pa_lock);
 			*busy = 1;
 			continue;
@@ -5637,7 +5637,7 @@ void ext4_discard_preallocations(struct inode *inode)
 		BUG_ON(pa->pa_node_lock.inode_lock != &ei->i_prealloc_lock);
 
 		spin_lock(&pa->pa_lock);
-		if (atomic_read(&pa->pa_count)) {
+		if (refcount_read(&pa->pa_count)) {
 			/* this shouldn't happen often - nobody should
 			 * use preallocation while we're discarding it */
 			spin_unlock(&pa->pa_lock);
@@ -5720,7 +5720,7 @@ static int ext4_mb_pa_alloc(struct ext4_allocation_context *ac)
 	pa = kmem_cache_zalloc(ext4_pspace_cachep, GFP_NOFS);
 	if (!pa)
 		return -ENOMEM;
-	atomic_set(&pa->pa_count, 1);
+	refcount_set(&pa->pa_count, 1);
 	ac->ac_pa = pa;
 	return 0;
 }
@@ -5731,7 +5731,7 @@ static void ext4_mb_pa_put_free(struct ext4_allocation_context *ac)
 
 	BUG_ON(!pa);
 	ac->ac_pa = NULL;
-	WARN_ON(!atomic_dec_and_test(&pa->pa_count));
+	WARN_ON(!refcount_dec_and_test(&pa->pa_count));
 	/*
 	 * current function is only called due to an error or due to
 	 * len of found blocks < len of requested blocks hence the PA has not
@@ -5948,7 +5948,7 @@ ext4_mb_discard_lg_preallocations(struct super_block *sb,
 				pa_node.lg_list,
 				lockdep_is_held(&lg->lg_prealloc_lock)) {
 		spin_lock(&pa->pa_lock);
-		if (atomic_read(&pa->pa_count)) {
+		if (refcount_read(&pa->pa_count)) {
 			/*
 			 * This is the pa that we just used
 			 * for block allocation. So don't
diff --git a/fs/ext4/mballoc.h b/fs/ext4/mballoc.h
index 39333ce72cbd..f708d303228e 100644
--- a/fs/ext4/mballoc.h
+++ b/fs/ext4/mballoc.h
@@ -126,7 +126,7 @@ struct ext4_prealloc_space {
 		struct rcu_head	pa_rcu;
 	} u;
 	spinlock_t		pa_lock;
-	atomic_t		pa_count;
+	refcount_t		pa_count;
 	unsigned		pa_deleted;
 	ext4_fsblk_t		pa_pstart;	/* phys. block */
 	ext4_lblk_t		pa_lstart;	/* log. block */
-- 
2.43.0