The shadow gmap returned by gmap_create_shadow() could get dropped
before taking the gmap->children_lock. This meant that the shadow gmap
was sometimes being used while its reference count was 0.

Fix this by taking the additional reference inside gmap_create_shadow()
while still holding gmap->children_lock, instead of afterwards.

Fixes: e38c884df921 ("KVM: s390: Switch to new gmap")
Signed-off-by: Claudio Imbrenda <imbrenda@linux.ibm.com>
---
 arch/s390/kvm/gmap.c | 15 ++++++++++++---
 arch/s390/kvm/vsie.c |  6 +++++-
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/arch/s390/kvm/gmap.c b/arch/s390/kvm/gmap.c
index da222962ef6d..26cd2b208b6f 100644
--- a/arch/s390/kvm/gmap.c
+++ b/arch/s390/kvm/gmap.c
@@ -1179,6 +1179,8 @@ static int gmap_protect_asce_top_level(struct kvm_s390_mmu_cache *mc, struct gma
  * The shadow table will be removed automatically on any change to the
  * PTE mapping for the source table.
  *
+ * The returned shadow gmap will be returned with one extra reference.
+ *
  * Return: A guest address space structure, ERR_PTR(-ENOMEM) if out of memory,
  * ERR_PTR(-EAGAIN) if the caller has to retry and ERR_PTR(-EFAULT) if the
  * parent gmap table could not be protected.
@@ -1189,10 +1191,13 @@ struct gmap *gmap_create_shadow(struct kvm_s390_mmu_cache *mc, struct gmap *pare
 	struct gmap *sg, *new;
 	int rc;
 
-	scoped_guard(spinlock, &parent->children_lock)
+	scoped_guard(spinlock, &parent->children_lock) {
 		sg = gmap_find_shadow(parent, asce, edat_level);
-	if (sg)
-		return sg;
+		if (sg) {
+			gmap_get(sg);
+			return sg;
+		}
+	}
 	/* Create a new shadow gmap. */
 	new = gmap_new(parent->kvm, asce.r ? 1UL << (64 - PAGE_SHIFT) : asce_end(asce));
 	if (!new)
@@ -1206,6 +1211,7 @@ struct gmap *gmap_create_shadow(struct kvm_s390_mmu_cache *mc, struct gmap *pare
 		sg = gmap_find_shadow(parent, asce, edat_level);
 		if (sg) {
 			gmap_put(new);
+			gmap_get(sg);
 			return sg;
 		}
 		if (asce.r) {
@@ -1219,16 +1225,19 @@ struct gmap *gmap_create_shadow(struct kvm_s390_mmu_cache *mc, struct gmap *pare
 			}
 			gmap_add_child(parent, new);
 			/* Nothing to protect, return right away. */
+			gmap_get(new);
 			return new;
 		}
 	}
 
+	gmap_get(new);
 	new->parent = parent;
 	/* Protect while inserting, protects against invalidation races. */
 	rc = gmap_protect_asce_top_level(mc, new);
 	if (rc) {
 		new->parent = NULL;
 		gmap_put(new);
+		gmap_put(new);
 		return ERR_PTR(rc);
 	}
 	return new;
diff --git a/arch/s390/kvm/vsie.c b/arch/s390/kvm/vsie.c
index faf8b01fa672..d0296491b2f7 100644
--- a/arch/s390/kvm/vsie.c
+++ b/arch/s390/kvm/vsie.c
@@ -1256,6 +1256,7 @@ static struct gmap *acquire_gmap_shadow(struct kvm_vcpu *vcpu, struct vsie_page
 			release_gmap_shadow(vsie_page);
 		}
 	}
+again:
 	gmap = gmap_create_shadow(vcpu->arch.mc, vcpu->kvm->arch.gmap, asce, edat);
 	if (IS_ERR(gmap))
 		return gmap;
@@ -1263,11 +1264,14 @@ static struct gmap *acquire_gmap_shadow(struct kvm_vcpu *vcpu, struct vsie_page
 		/* unlikely race condition, remove the previous shadow */
 		if (vsie_page->gmap_cache.gmap)
 			release_gmap_shadow(vsie_page);
+		if (!gmap->parent) {
+			gmap_put(gmap);
+			goto again;
+		}
 		vcpu->kvm->stat.gmap_shadow_create++;
 		list_add(&vsie_page->gmap_cache.list, &gmap->scb_users);
 		vsie_page->gmap_cache.gmap = gmap;
 		prefix_unmapped(vsie_page);
-		gmap_get(gmap);
 	}
 	return gmap;
 }
-- 
2.52.0