unsigned short ibc;
 };
 
-struct kvm_s390_module_hook {
-       int (*hook)(struct kvm_vcpu *vcpu);
-       struct module *owner;
-};
+typedef int (*crypto_hook)(struct kvm_vcpu *vcpu);
 
 struct kvm_s390_crypto {
        struct kvm_s390_crypto_cb *crycb;
-       struct kvm_s390_module_hook *pqap_hook;
+       struct rw_semaphore pqap_hook_rwsem;
+       crypto_hook *pqap_hook;
        __u32 crycbd;
        __u8 aes_kw;
        __u8 dea_kw;
 
 static int handle_pqap(struct kvm_vcpu *vcpu)
 {
        struct ap_queue_status status = {};
+       crypto_hook pqap_hook;
        unsigned long reg0;
        int ret;
        uint8_t fc;
                return kvm_s390_inject_program_int(vcpu, PGM_SPECIFICATION);
 
        /*
-        * Verify that the hook callback is registered, lock the owner
-        * and call the hook.
+        * If the hook callback is registered, there will be a pointer to the
+        * hook function pointer in the kvm_s390_crypto structure. Lock the
+        * owner, retrieve the hook function pointer and call the hook.
         */
+       down_read(&vcpu->kvm->arch.crypto.pqap_hook_rwsem);
        if (vcpu->kvm->arch.crypto.pqap_hook) {
-               if (!try_module_get(vcpu->kvm->arch.crypto.pqap_hook->owner))
-                       return -EOPNOTSUPP;
-               ret = vcpu->kvm->arch.crypto.pqap_hook->hook(vcpu);
-               module_put(vcpu->kvm->arch.crypto.pqap_hook->owner);
+               pqap_hook = *vcpu->kvm->arch.crypto.pqap_hook;
+               ret = pqap_hook(vcpu);
                if (!ret && vcpu->run->s.regs.gprs[1] & 0x00ff0000)
                        kvm_s390_set_psw_cc(vcpu, 3);
+               up_read(&vcpu->kvm->arch.crypto.pqap_hook_rwsem);
                return ret;
        }
+       up_read(&vcpu->kvm->arch.crypto.pqap_hook_rwsem);
        /*
         * A vfio_driver must register a hook.
         * No hook means no driver to enable the SIE CRYCB and no queues.
 
        vfio_ap_matrix_init(&matrix_dev->info, &matrix_mdev->matrix);
        init_waitqueue_head(&matrix_mdev->wait_for_kvm);
        mdev_set_drvdata(mdev, matrix_mdev);
-       matrix_mdev->pqap_hook.hook = handle_pqap;
-       matrix_mdev->pqap_hook.owner = THIS_MODULE;
+       matrix_mdev->pqap_hook = handle_pqap;
        mutex_lock(&matrix_dev->lock);
        list_add(&matrix_mdev->node, &matrix_dev->mdev_list);
        mutex_unlock(&matrix_dev->lock);
                }
 
                kvm_get_kvm(kvm);
+               matrix_mdev->kvm = kvm;
                matrix_mdev->kvm_busy = true;
                mutex_unlock(&matrix_dev->lock);
+
+               down_write(&matrix_mdev->kvm->arch.crypto.pqap_hook_rwsem);
+               kvm->arch.crypto.pqap_hook = &matrix_mdev->pqap_hook;
+               up_write(&matrix_mdev->kvm->arch.crypto.pqap_hook_rwsem);
+
                kvm_arch_crypto_set_masks(kvm,
                                          matrix_mdev->matrix.apm,
                                          matrix_mdev->matrix.aqm,
                                          matrix_mdev->matrix.adm);
+
                mutex_lock(&matrix_dev->lock);
-               kvm->arch.crypto.pqap_hook = &matrix_mdev->pqap_hook;
-               matrix_mdev->kvm = kvm;
                matrix_mdev->kvm_busy = false;
                wake_up_all(&matrix_mdev->wait_for_kvm);
        }
        if (matrix_mdev->kvm) {
                matrix_mdev->kvm_busy = true;
                mutex_unlock(&matrix_dev->lock);
-               kvm_arch_crypto_clear_masks(matrix_mdev->kvm);
+
+               if (matrix_mdev->kvm->arch.crypto.crycbd) {
+                       down_write(&matrix_mdev->kvm->arch.crypto.pqap_hook_rwsem);
+                       matrix_mdev->kvm->arch.crypto.pqap_hook = NULL;
+                       up_write(&matrix_mdev->kvm->arch.crypto.pqap_hook_rwsem);
+
+                       kvm_arch_crypto_clear_masks(matrix_mdev->kvm);
+               }
+
                mutex_lock(&matrix_dev->lock);
                vfio_ap_mdev_reset_queues(matrix_mdev->mdev);
-               matrix_mdev->kvm->arch.crypto.pqap_hook = NULL;
                kvm_put_kvm(matrix_mdev->kvm);
                matrix_mdev->kvm = NULL;
                matrix_mdev->kvm_busy = false;