[RFC,11/11] kvm: powerpc: book3s: Fix module ownership

diff --git a/arch/powerpc/include/asm/kvm_ppc.h b/arch/powerpc/include/asm/kvm_ppc.h
index a4a5893..2022720 100644
--- a/arch/powerpc/include/asm/kvm_ppc.h
+++ b/arch/powerpc/include/asm/kvm_ppc.h
@@ -171,6 +171,7 @@  extern int kvmppc_xics_int_on(struct kvm *kvm, u32 irq);
 extern int kvmppc_xics_int_off(struct kvm *kvm, u32 irq);
 
 struct kvmppc_ops {
+	struct module *owner;
 	bool is_hv_enabled;
 	int (*get_sregs)(struct kvm_vcpu *vcpu, struct kvm_sregs *sregs);
 	int (*set_sregs)(struct kvm_vcpu *vcpu, struct kvm_sregs *sregs);
diff --git a/arch/powerpc/kvm/book3s.c b/arch/powerpc/kvm/book3s.c
index 34e189c..363df6a 100644
--- a/arch/powerpc/kvm/book3s.c
+++ b/arch/powerpc/kvm/book3s.c
@@ -831,6 +831,10 @@  int kvm_arch_check_processor_compat(void *opaque)
 {
 	int r,cpu;
 	struct kvmppc_ops *kvm_ops = (struct kvmppc_ops *)opaque;
+
+	if (!kvm_ops)
+		return 0;
+
 	for_each_online_cpu(cpu) {
 		smp_call_function_single(cpu,
 					 kvm_ops->check_processor_compat,
@@ -840,6 +844,7 @@  int kvm_arch_check_processor_compat(void *opaque)
 	}
 	return r;
 }
+EXPORT_SYMBOL_GPL(kvm_arch_check_processor_compat);
 
 EXPORT_SYMBOL_GPL(kvm_get_dirty_log);
 EXPORT_SYMBOL_GPL(kvmppc_core_pending_dec);
@@ -851,3 +856,19 @@  EXPORT_SYMBOL_GPL(kvmppc_core_prepare_to_enter);
 EXPORT_SYMBOL_GPL(kvmppc_core_queue_dec);
 EXPORT_SYMBOL_GPL(kvmppc_free_lpid);
 
+static int kvmppc_book3s_init(void)
+{
+	int r;
+
+	r = kvm_init(NULL, sizeof(struct kvm_vcpu), 0, THIS_MODULE);
+	return r;
+}
+
+static void kvmppc_book3s_exit(void)
+{
+	kvm_exit();
+}
+
+module_init(kvmppc_book3s_init);
+module_exit(kvmppc_book3s_exit);
+
diff --git a/arch/powerpc/kvm/book3s_hv.c b/arch/powerpc/kvm/book3s_hv.c
index 0a684a7..7bdc780 100644
--- a/arch/powerpc/kvm/book3s_hv.c
+++ b/arch/powerpc/kvm/book3s_hv.c
@@ -2092,23 +2092,20 @@  static int kvmppc_book3s_init_hv(void)
 {
 	int r;
 
-	if (!kvmppc_pr_ops) {
-		r = kvm_init(&kvm_ops_hv, sizeof(struct kvm_vcpu),
-			     0, THIS_MODULE);
-		if (r)
-			return r;
-	}
+	r = kvm_arch_check_processor_compat(&kvm_ops_hv);
+	if (r < 0)
+		return r;
+
+	kvm_ops_hv.owner = THIS_MODULE;
 	kvmppc_hv_ops = &kvm_ops_hv;
-	r = kvmppc_mmu_hv_init();
 
+	r = kvmppc_mmu_hv_init();
 	return r;
 }
 
 static void kvmppc_book3s_exit_hv(void)
 {
 	kvmppc_hv_ops = NULL;
-	if (!kvmppc_pr_ops)
-		kvm_exit();
 }
 
 module_init(kvmppc_book3s_init_hv);
diff --git a/arch/powerpc/kvm/book3s_pr.c b/arch/powerpc/kvm/book3s_pr.c
index e49e4b0..c79fada 100644
--- a/arch/powerpc/kvm/book3s_pr.c
+++ b/arch/powerpc/kvm/book3s_pr.c
@@ -1551,17 +1551,14 @@  static int kvmppc_book3s_init_pr(void)
 {
 	int r;
 
-	if (!kvmppc_hv_ops) {
-		r = kvm_init(&kvm_ops_pr, sizeof(struct kvm_vcpu),
-			     0, THIS_MODULE);
-		if (r)
-			return r;
-	}
-	/* Assign the global value */
+	r = kvm_arch_check_processor_compat(&kvm_ops_pr);
+	if (r < 0)
+		return r;
+
+	kvm_ops_pr.owner = THIS_MODULE;
 	kvmppc_pr_ops = &kvm_ops_pr;
 
 	r = kvmppc_mmu_hpte_sysinit();
-
 	return r;
 }
 
@@ -1569,8 +1566,6 @@  static void kvmppc_book3s_exit_pr(void)
 {
 	kvmppc_pr_ops = NULL;
 	kvmppc_mmu_hpte_sysexit();
-	if (!kvmppc_hv_ops)
-		kvm_exit();
 }
 
 module_init(kvmppc_book3s_init_pr);
diff --git a/arch/powerpc/kvm/powerpc.c b/arch/powerpc/kvm/powerpc.c
index 1209229..677fa7e 100644
--- a/arch/powerpc/kvm/powerpc.c
+++ b/arch/powerpc/kvm/powerpc.c
@@ -26,6 +26,7 @@ 
 #include <linux/fs.h>
 #include <linux/slab.h>
 #include <linux/file.h>
+#include <linux/module.h>
 #include <asm/cputable.h>
 #include <asm/uaccess.h>
 #include <asm/kvm_ppc.h>
@@ -270,25 +271,32 @@  void kvm_arch_hardware_unsetup(void)
 
 int kvm_arch_init_vm(struct kvm *kvm, unsigned long type)
 {
+	struct kvmppc_ops *kvm_ops = NULL;
 	/*
 	 * if we have both HV and PR enabled, default is HV
 	 */
 	if (type == 0) {
 		if (kvmppc_hv_ops)
-			kvm->arch.kvm_ops = kvmppc_hv_ops;
+			kvm_ops = kvmppc_hv_ops;
 		else
-			kvm->arch.kvm_ops = kvmppc_pr_ops;
+			kvm_ops = kvmppc_pr_ops;
+		if (!kvm_ops)
+			goto err_out;
 	} else	if (type == KVM_VM_PPC_HV) {
 		if (!kvmppc_hv_ops)
 			goto err_out;
-		kvm->arch.kvm_ops = kvmppc_hv_ops;
+		kvm_ops = kvmppc_hv_ops;
 	} else if (type == KVM_VM_PPC_PR) {
 		if (!kvmppc_pr_ops)
 			goto err_out;
-		kvm->arch.kvm_ops = kvmppc_pr_ops;
+		kvm_ops = kvmppc_pr_ops;
 	} else
 		goto err_out;
 
+	if (kvm_ops->owner && !try_module_get(kvm_ops->owner))
+		return -ENOENT;
+
+	kvm->arch.kvm_ops = kvm_ops;
 	return kvmppc_core_init_vm(kvm);
 err_out:
 	return -EINVAL;
@@ -311,6 +319,9 @@  void kvm_arch_destroy_vm(struct kvm *kvm)
 	kvmppc_core_destroy_vm(kvm);
 
 	mutex_unlock(&kvm->lock);
+
+	/* drop the module reference */
+	module_put(kvm->arch.kvm_ops->owner);
 }
 
 void kvm_arch_sync_events(struct kvm *kvm)