From: Paolo Bonzini <pbonzini@redhat.com>
Date: Fri, 12 Nov 2021 04:02:24 -0500
Subject: KVM: SEV: unify cgroup cleanup code for svm_vm_migrate_from
Git-commit: 501cfe06790651af4470d8ca77f162515867cd4d
Patch-mainline: v5.16-rc1
References: jsc#SLE-19924

Use the same cleanup code independent of whether the cgroup to be
uncharged and unref'd is the source or the destination cgroup.  Use a
bool to track whether the destination cgroup has been charged, which also
fixes a bug in the error case: the destination cgroup must be uncharged
only if it does not match the source.

Fixes: b56639318bb2 ("KVM: SEV: Add support for SEV intra host migration")
Signed-off-by: Paolo Bonzini <pbonzini@redhat.com>
Acked-by: Joerg Roedel <jroedel@suse.de>
---
 arch/x86/kvm/svm/sev.c |   25 +++++++++++--------------
 1 file changed, 11 insertions(+), 14 deletions(-)

--- a/arch/x86/kvm/svm/sev.c
+++ b/arch/x86/kvm/svm/sev.c
@@ -1615,12 +1615,6 @@ static void sev_migrate_from(struct kvm_
 	src->handle = 0;
 	src->pages_locked = 0;
 
-	if (dst->misc_cg != src->misc_cg)
-		sev_misc_cg_uncharge(src);
-
-	put_misc_cg(src->misc_cg);
-	src->misc_cg = NULL;
-
 	INIT_LIST_HEAD(&dst->regions_list);
 	list_replace_init(&src->regions_list, &dst->regions_list);
 }
@@ -1668,9 +1662,10 @@ static int sev_es_migrate_from(struct kv
 int svm_vm_migrate_from(struct kvm *kvm, unsigned int source_fd)
 {
 	struct kvm_sev_info *dst_sev = &to_kvm_svm(kvm)->sev_info;
-	struct kvm_sev_info *src_sev;
+	struct kvm_sev_info *src_sev, *cg_cleanup_sev;
 	struct file *source_kvm_file;
 	struct kvm *source_kvm;
+	bool charged = false;
 	int ret;
 
 	ret = sev_lock_for_migration(kvm);
@@ -1700,10 +1695,12 @@ int svm_vm_migrate_from(struct kvm *kvm,
 
 	src_sev = &to_kvm_svm(source_kvm)->sev_info;
 	dst_sev->misc_cg = get_current_misc_cg();
+	cg_cleanup_sev = dst_sev;
 	if (dst_sev->misc_cg != src_sev->misc_cg) {
 		ret = sev_misc_cg_try_charge(dst_sev);
 		if (ret)
-			goto out_dst_put_cgroup;
+			goto out_dst_cgroup;
+		charged = true;
 	}
 
 	ret = sev_lock_vcpus_for_migration(kvm);
@@ -1720,6 +1717,7 @@ int svm_vm_migrate_from(struct kvm *kvm,
 	}
 	sev_migrate_from(dst_sev, src_sev);
 	kvm_vm_dead(source_kvm);
+	cg_cleanup_sev = src_sev;
 	ret = 0;
 
 out_source_vcpu:
@@ -1727,12 +1725,11 @@ out_source_vcpu:
 out_dst_vcpu:
 	sev_unlock_vcpus_for_migration(kvm);
 out_dst_cgroup:
-	if (ret < 0) {
-		sev_misc_cg_uncharge(dst_sev);
-out_dst_put_cgroup:
-		put_misc_cg(dst_sev->misc_cg);
-		dst_sev->misc_cg = NULL;
-	}
+	/* Operates on the source on success, on the destination on failure.  */
+	if (charged)
+		sev_misc_cg_uncharge(cg_cleanup_sev);
+	put_misc_cg(cg_cleanup_sev->misc_cg);
+	cg_cleanup_sev->misc_cg = NULL;
 out_source:
 	sev_unlock_after_migration(source_kvm);
 out_fput: