fix: migration FSM transitions

Signed-off-by: William Henderson <william.henderson@nutanix.com>

4 files changed, 87 insertions, 16 deletions

diff --git a/include/libvfio-user.h b/include/libvfio-user.h
index bb015c2..c1c844a 100644
--- a/include/libvfio-user.h
+++ b/include/libvfio-user.h
@@ -88,12 +88,6 @@ dma_sg_size(void);
  */
 #define LIBVFIO_USER_FLAG_ATTACH_NB  (1 << 0)
 
-/*
- * If set in the flags given to `init_migration`, the server will start in the
- * RESUMING state, ready to receive migration data.
- */
-#define LIBVFIO_USER_MIG_FLAG_START_RESUMING (1 << 0)
-
 typedef enum {
     VFU_TRANS_SOCK,
     // For internal testing only
diff --git a/include/vfio-user.h b/include/vfio-user.h
index 1652a84..c2d2503 100644
--- a/include/vfio-user.h
+++ b/include/vfio-user.h
@@ -284,6 +284,8 @@ enum vfio_user_device_mig_state {
 	VFIO_USER_DEVICE_STATE_PRE_COPY_P2P = 7,
 };
 
+#define VFIO_USER_DEVICE_NUM_STATES 8
+
 struct vfio_user_mig_data {
     uint32_t    argsz;
     uint32_t    size;
diff --git a/lib/migration.c b/lib/migration.c
index 7b55813..879978a 100644
--- a/lib/migration.c
+++ b/lib/migration.c
@@ -71,11 +71,7 @@ init_migration(const vfu_migration_callbacks_t *callbacks,
     migr->pgsize = sysconf(_SC_PAGESIZE);
 
     /* FIXME this should be done in vfu_ctx_realize */
-    if (flags & LIBVFIO_USER_MIG_FLAG_START_RESUMING) {
-        migr->state = VFIO_USER_DEVICE_STATE_RESUMING;
-    } else {
-        migr->state = VFIO_USER_DEVICE_STATE_RUNNING;
-    }
+    migr->state = VFIO_USER_DEVICE_STATE_RUNNING;
 
     migr->callbacks = *callbacks;
     if (migr->callbacks.transition == NULL ||
@@ -220,8 +216,20 @@ migration_feature_set(vfu_ctx_t *vfu_ctx, uint32_t feature, void *buf)
     if (feature == VFIO_DEVICE_FEATURE_MIG_DEVICE_STATE) {
         struct vfio_user_device_feature_mig_state *res = buf;
         struct migration *migr = vfu_ctx->migration;
+        uint32_t state;
+        ssize_t ret;
+        
+        do {
+            state = next_state[migr->state][res->device_state];
+
+            if (state == VFIO_USER_DEVICE_STATE_ERROR) {
+                return -EINVAL;
+            }
+
+            ret = handle_device_state(vfu_ctx, migr, state, true);
+        } while (migr->state != res->device_state && ret == 0);
         
-        return handle_device_state(vfu_ctx, migr, res->device_state, true);
+        return ret;
     }
 
     return -EINVAL;
diff --git a/lib/migration_priv.h b/lib/migration_priv.h
index 2566b39..6272995 100644
--- a/lib/migration_priv.h
+++ b/lib/migration_priv.h
@@ -40,10 +40,12 @@ struct migration {
     vfu_migration_callbacks_t callbacks;
 };
 
-/* valid migration state transitions
-   each number corresponds to a FROM state and each bit a TO state
-   if the bit is set then the transition is allowed
-   the indices of each state are those in the vfio_user_device_mig_state enum */
+/*
+ * valid migration state transitions
+ * each number corresponds to a FROM state and each bit a TO state
+ * if the bit is set then the transition is allowed
+ * the indices of each state are those in the vfio_user_device_mig_state enum
+ */
 static const char transitions[8] = {
     0b00000000, // ERROR        -> {}
     0b00011100, // STOP         -> {RUNNING, STOP_COPY, RESUMING}
@@ -55,6 +57,71 @@ static const char transitions[8] = {
     0b00000000  // PRE_COPY_P2P -> {}
 };
 
+/*
+ * The spec dictates that, if no direct transition is allowed, and the
+ * transition is not one of the explicitly disallowed ones (i.e. anything to
+ * ERROR, anything from ERROR, and STOP_COPY -> PRE_COPY), we should take the
+ * shortest allowed path.
+ * 
+ * This can be indexed as `next_state[current][target] == next`. If next is
+ * ERROR, then the transition is not allowed.
+ */
+static const uint32_t next_state[VFIO_USER_DEVICE_NUM_STATES][VFIO_USER_DEVICE_NUM_STATES] = {
+    [VFIO_USER_DEVICE_STATE_ERROR] = {0, 0, 0, 0, 0, 0, 0, 0},
+    [VFIO_USER_DEVICE_STATE_STOP] = {
+        [VFIO_USER_DEVICE_STATE_ERROR] = VFIO_USER_DEVICE_STATE_ERROR,
+        [VFIO_USER_DEVICE_STATE_STOP] = VFIO_USER_DEVICE_STATE_STOP,
+        [VFIO_USER_DEVICE_STATE_RUNNING] = VFIO_USER_DEVICE_STATE_RUNNING,
+        [VFIO_USER_DEVICE_STATE_STOP_COPY] = VFIO_USER_DEVICE_STATE_STOP_COPY,
+        [VFIO_USER_DEVICE_STATE_RESUMING] = VFIO_USER_DEVICE_STATE_RESUMING,
+        [VFIO_USER_DEVICE_STATE_RUNNING_P2P] = VFIO_USER_DEVICE_STATE_ERROR,
+        [VFIO_USER_DEVICE_STATE_PRE_COPY] = VFIO_USER_DEVICE_STATE_RUNNING,
+        [VFIO_USER_DEVICE_STATE_PRE_COPY_P2P] = VFIO_USER_DEVICE_STATE_ERROR,
+    },
+    [VFIO_USER_DEVICE_STATE_RUNNING] = {
+        [VFIO_USER_DEVICE_STATE_ERROR] = VFIO_USER_DEVICE_STATE_ERROR,
+        [VFIO_USER_DEVICE_STATE_STOP] = VFIO_USER_DEVICE_STATE_STOP,
+        [VFIO_USER_DEVICE_STATE_RUNNING] = VFIO_USER_DEVICE_STATE_RUNNING,
+        [VFIO_USER_DEVICE_STATE_STOP_COPY] = VFIO_USER_DEVICE_STATE_STOP,
+        [VFIO_USER_DEVICE_STATE_RESUMING] = VFIO_USER_DEVICE_STATE_STOP,
+        [VFIO_USER_DEVICE_STATE_RUNNING_P2P] = VFIO_USER_DEVICE_STATE_ERROR,
+        [VFIO_USER_DEVICE_STATE_PRE_COPY] = VFIO_USER_DEVICE_STATE_PRE_COPY,
+        [VFIO_USER_DEVICE_STATE_PRE_COPY_P2P] = VFIO_USER_DEVICE_STATE_ERROR,
+    },
+    [VFIO_USER_DEVICE_STATE_STOP_COPY] = {
+        [VFIO_USER_DEVICE_STATE_ERROR] = VFIO_USER_DEVICE_STATE_ERROR,
+        [VFIO_USER_DEVICE_STATE_STOP] = VFIO_USER_DEVICE_STATE_STOP,
+        [VFIO_USER_DEVICE_STATE_RUNNING] = VFIO_USER_DEVICE_STATE_STOP,
+        [VFIO_USER_DEVICE_STATE_STOP_COPY] = VFIO_USER_DEVICE_STATE_STOP_COPY,
+        [VFIO_USER_DEVICE_STATE_RESUMING] = VFIO_USER_DEVICE_STATE_STOP,
+        [VFIO_USER_DEVICE_STATE_RUNNING_P2P] = VFIO_USER_DEVICE_STATE_ERROR,
+        [VFIO_USER_DEVICE_STATE_PRE_COPY] = VFIO_USER_DEVICE_STATE_STOP,
+        [VFIO_USER_DEVICE_STATE_PRE_COPY_P2P] = VFIO_USER_DEVICE_STATE_ERROR,
+    },
+    [VFIO_USER_DEVICE_STATE_RESUMING] = {
+        [VFIO_USER_DEVICE_STATE_ERROR] = VFIO_USER_DEVICE_STATE_ERROR,
+        [VFIO_USER_DEVICE_STATE_STOP] = VFIO_USER_DEVICE_STATE_STOP,
+        [VFIO_USER_DEVICE_STATE_RUNNING] = VFIO_USER_DEVICE_STATE_STOP,
+        [VFIO_USER_DEVICE_STATE_STOP_COPY] = VFIO_USER_DEVICE_STATE_STOP,
+        [VFIO_USER_DEVICE_STATE_RESUMING] = VFIO_USER_DEVICE_STATE_RESUMING,
+        [VFIO_USER_DEVICE_STATE_RUNNING_P2P] = VFIO_USER_DEVICE_STATE_ERROR,
+        [VFIO_USER_DEVICE_STATE_PRE_COPY] = VFIO_USER_DEVICE_STATE_STOP,
+        [VFIO_USER_DEVICE_STATE_PRE_COPY_P2P] = VFIO_USER_DEVICE_STATE_ERROR,
+    },
+    [VFIO_USER_DEVICE_STATE_RUNNING_P2P] = {0, 0, 0, 0, 0, 0, 0, 0},
+    [VFIO_USER_DEVICE_STATE_PRE_COPY] = {
+        [VFIO_USER_DEVICE_STATE_ERROR] = VFIO_USER_DEVICE_STATE_ERROR,
+        [VFIO_USER_DEVICE_STATE_STOP] = VFIO_USER_DEVICE_STATE_RUNNING,
+        [VFIO_USER_DEVICE_STATE_RUNNING] = VFIO_USER_DEVICE_STATE_RUNNING,
+        [VFIO_USER_DEVICE_STATE_STOP_COPY] = VFIO_USER_DEVICE_STATE_STOP_COPY,
+        [VFIO_USER_DEVICE_STATE_RESUMING] = VFIO_USER_DEVICE_STATE_RUNNING,
+        [VFIO_USER_DEVICE_STATE_RUNNING_P2P] = VFIO_USER_DEVICE_STATE_ERROR,
+        [VFIO_USER_DEVICE_STATE_PRE_COPY] = VFIO_USER_DEVICE_STATE_PRE_COPY,
+        [VFIO_USER_DEVICE_STATE_PRE_COPY_P2P] = VFIO_USER_DEVICE_STATE_ERROR,
+    },
+    [VFIO_USER_DEVICE_STATE_PRE_COPY_P2P] = {0, 0, 0, 0, 0, 0, 0, 0},
+};
+
 MOCK_DECLARE(vfu_migr_state_t, migr_state_vfio_to_vfu, uint32_t device_state);
 
 MOCK_DECLARE(int, state_trans_notify, vfu_ctx_t *vfu_ctx,