From: "Maciej S. Szmigiero" <maciej.szmigiero@oracle.com>

It's possible for load_cleanup SaveVMHandler to get called without
load_setup handler being called first.

Since we'll be soon running cleanup operations there that access objects
that need earlier initialization in load_setup let's make sure these
cleanups only run when load_setup handler had indeed been called
earlier.

Signed-off-by: Maciej S. Szmigiero <maciej.szmigiero@oracle.com>
---
 hw/vfio/migration.c           | 21 +++++++++++++++++++--
 include/hw/vfio/vfio-common.h |  1 +
 2 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/hw/vfio/migration.c b/hw/vfio/migration.c
index 01aa11013e42..9e2657073012 100644
--- a/hw/vfio/migration.c
+++ b/hw/vfio/migration.c
@@ -688,16 +688,33 @@ static void vfio_save_state(QEMUFile *f, void *opaque)
 static int vfio_load_setup(QEMUFile *f, void *opaque, Error **errp)
 {
     VFIODevice *vbasedev = opaque;
+    VFIOMigration *migration = vbasedev->migration;
+    int ret;
+
+    assert(!migration->load_setup);
+
+    ret = vfio_migration_set_state(vbasedev, VFIO_DEVICE_STATE_RESUMING,
+                                   migration->device_state, errp);
+    if (ret) {
+        return ret;
+    }
 
-    return vfio_migration_set_state(vbasedev, VFIO_DEVICE_STATE_RESUMING,
-                                    vbasedev->migration->device_state, errp);
+    migration->load_setup = true;
+
+    return 0;
 }
 
 static int vfio_load_cleanup(void *opaque)
 {
     VFIODevice *vbasedev = opaque;
+    VFIOMigration *migration = vbasedev->migration;
+
+    if (!migration->load_setup) {
+        return 0;
+    }
 
     vfio_migration_cleanup(vbasedev);
+    migration->load_setup = false;
     trace_vfio_load_cleanup(vbasedev->name);
 
     return 0;
diff --git a/include/hw/vfio/vfio-common.h b/include/hw/vfio/vfio-common.h
index e0ce6ec3a9b3..246250ed8b75 100644
--- a/include/hw/vfio/vfio-common.h
+++ b/include/hw/vfio/vfio-common.h
@@ -66,6 +66,7 @@ typedef struct VFIOMigration {
     VMChangeStateEntry *vm_state;
     NotifierWithReturn migration_state;
     uint32_t device_state;
+    bool load_setup;
     int data_fd;
     void *data_buffer;
     size_t data_buffer_size;