fix: implementation fixes from comments

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

5 files changed, 107 insertions, 65 deletions

diff --git a/include/vfio-user.h b/include/vfio-user.h
index 42ffa29..77b3c43 100644
--- a/include/vfio-user.h
+++ b/include/vfio-user.h
@@ -229,30 +229,42 @@ struct vfio_user_bitmap_range {
 struct vfio_user_device_feature {
 	uint32_t	argsz;
 	uint32_t	flags;
+#ifndef VFIO_DEVICE_FEATURE_MASK
 #define VFIO_DEVICE_FEATURE_MASK	(0xffff)  /* 16-bit feature index */
 #define VFIO_DEVICE_FEATURE_GET		(1 << 16) /* Get feature into data[] */
 #define VFIO_DEVICE_FEATURE_SET		(1 << 17) /* Set feature from data[] */
 #define VFIO_DEVICE_FEATURE_PROBE	(1 << 18) /* Probe feature support */
+#endif
 	uint8_t  	data[];
 } __attribute__((packed));
 
 /* Analogous to vfio_device_feature_migration */
 struct vfio_user_device_feature_migration {
     uint64_t flags;
+#ifndef VFIO_MIGRATION_STOP_COPY
 #define VFIO_MIGRATION_STOP_COPY    (1 << 0)
 #define VFIO_MIGRATION_P2P          (1 << 1)
 #define VFIO_MIGRATION_PRE_COPY     (1 << 2)
+#endif
 } __attribute__((packed));
+#ifndef VFIO_DEVICE_FEATURE_MIGRATION
 #define VFIO_DEVICE_FEATURE_MIGRATION 1
+#endif
+_Static_assert(sizeof(struct vfio_user_device_feature_migration) == 8,
+               "bad vfio_user_device_feature_migration size");
 
 /* Analogous to vfio_device_feature_mig_state */
 struct vfio_user_device_feature_mig_state {
     uint32_t    device_state;
     uint32_t    data_fd;
 } __attribute__((packed));
+#ifndef VFIO_DEVICE_FEATURE_MIG_DEVICE_STATE
 #define VFIO_DEVICE_FEATURE_MIG_DEVICE_STATE 2
+#endif
+_Static_assert(sizeof(struct vfio_user_device_feature_migration) == 8,
+               "bad vfio_user_device_feature_mig_state size");
 
-enum vfio_device_mig_state {
+enum vfio_user_device_mig_state {
 	VFIO_DEVICE_STATE_ERROR = 0,
 	VFIO_DEVICE_STATE_STOP = 1,
 	VFIO_DEVICE_STATE_RUNNING = 2,
@@ -263,16 +275,7 @@ enum vfio_device_mig_state {
 	VFIO_DEVICE_STATE_PRE_COPY_P2P = 7,
 };
 
-// FIXME awful names below
-
-// used for read request and write response
-struct vfio_user_mig_data_without_data {
-    uint32_t    argsz;
-    uint64_t    size;
-} __attribute__((packed));
-
-// used for write request and read response
-struct vfio_user_mig_data_with_data {
+struct vfio_user_mig_data {
     uint32_t    argsz;
     uint64_t    size;
     uint8_t     data[];
diff --git a/lib/libvfio-user.c b/lib/libvfio-user.c
index 7920494..145146d 100644
--- a/lib/libvfio-user.c
+++ b/lib/libvfio-user.c
@@ -1028,15 +1028,23 @@ handle_dirty_pages(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg)
 static int
 handle_device_feature(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg)
 {
+    assert(vfu_ctx != NULL);
+    assert(msg != NULL);
+
+    if (msg->in.iov.iov_len < sizeof(struct vfio_user_device_feature)) {
+        return -EINVAL;
+    }
+
     struct vfio_user_device_feature *req = msg->in.iov.iov_base;
 
     if (vfu_ctx->migration == NULL) {
         return -EINVAL;
     }
 
-    if (!migration_feature_supported(req->flags)) {
-        // FIXME what error code to return? we really want "not supported"
-        //   instead of "not permitted"?
+    uint32_t supported_flags =
+        migration_feature_flags(req->flags & VFIO_DEVICE_FEATURE_MASK);
+
+    if ((req->flags & supported_flags) != req->flags || supported_flags == 0) {
         return -EINVAL;
     }
 
@@ -1045,23 +1053,42 @@ handle_device_feature(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg)
     if (req->flags & VFIO_DEVICE_FEATURE_PROBE) {
         msg->out.iov.iov_base = malloc(msg->in.iov.iov_len);
         msg->out.iov.iov_len = msg->in.iov.iov_len;
+
+        if (msg->out.iov.iov_base == NULL) {
+            return -1;
+        }
+
         memcpy(msg->out.iov.iov_base, msg->in.iov.iov_base,
                msg->out.iov.iov_len);
 
         ret = 0;
     } else if (req->flags & VFIO_DEVICE_FEATURE_GET) {
-        msg->out.iov.iov_base = calloc(8, 1);
-        msg->out.iov.iov_len = 8;
+        // all supported outgoing data is currently the same size as 
+        //   vfio_user_device_feature_migration
+        msg->out.iov.iov_len = sizeof(struct vfio_user_device_feature_migration);
+        msg->out.iov.iov_base = calloc(1, msg->out.iov.iov_len);
 
-        ret = migration_feature_get(vfu_ctx, req->flags,
+        if (msg->out.iov.iov_base == NULL) {
+            return -1;
+        }
+
+        ret = migration_feature_get(vfu_ctx,
+                                    req->flags & VFIO_DEVICE_FEATURE_MASK,
                                     msg->out.iov.iov_base);
     } else if (req->flags & VFIO_DEVICE_FEATURE_SET) {
         msg->out.iov.iov_base = malloc(msg->in.iov.iov_len);
         msg->out.iov.iov_len = msg->in.iov.iov_len;
+
+        if (msg->out.iov.iov_base == NULL) {
+            return -1;
+        }
+
         memcpy(msg->out.iov.iov_base, msg->in.iov.iov_base,
                msg->out.iov.iov_len);
 
-        ret = migration_feature_set(vfu_ctx, req->flags, req->data);
+        ret = migration_feature_set(vfu_ctx,
+                                    req->flags & VFIO_DEVICE_FEATURE_MASK,
+                                    req->data);
     }
 
     return ret;
@@ -2051,7 +2078,7 @@ vfu_setup_device_migration_callbacks(vfu_ctx_t *vfu_ctx, uint64_t flags,
         return ERROR_INT(ret);
     }
 
-    return ret;
+    return 0;
 }
 
 #ifdef DEBUG
diff --git a/lib/migration.c b/lib/migration.c
index 4ab9181..69f3f35 100644
--- a/lib/migration.c
+++ b/lib/migration.c
@@ -42,7 +42,7 @@
 bool
 MOCK_DEFINE(vfio_migr_state_transition_is_valid)(uint32_t from, uint32_t to)
 {
-    return transitions[from][to];
+    return (transitions[from] & (1 << to)) != 0;
 }
 
 /*
@@ -91,7 +91,7 @@ init_migration(const vfu_migration_callbacks_t *callbacks,
 
 void
 MOCK_DEFINE(migr_state_transition)(struct migration *migr,
-                                   enum vfio_device_mig_state state)
+                                   enum vfio_user_device_mig_state state)
 {
     assert(migr != NULL);
     /* FIXME validate that state transition */
@@ -99,7 +99,7 @@ MOCK_DEFINE(migr_state_transition)(struct migration *migr,
 }
 
 vfu_migr_state_t
-MOCK_DEFINE(migr_state_vfio_to_vfu)(enum vfio_device_mig_state state)
+MOCK_DEFINE(migr_state_vfio_to_vfu)(enum vfio_user_device_mig_state state)
 {
     switch (state) {
         case VFIO_DEVICE_STATE_STOP:
@@ -163,6 +163,7 @@ MOCK_DEFINE(handle_device_state)(vfu_ctx_t *vfu_ctx, struct migration *migr,
                                  uint32_t device_state, bool notify)
 {
 
+    assert(vfu_ctx != NULL);
     assert(migr != NULL);
 
     if (!vfio_migr_state_transition_is_valid(migr->state, device_state)) {
@@ -171,48 +172,55 @@ MOCK_DEFINE(handle_device_state)(vfu_ctx_t *vfu_ctx, struct migration *migr,
     return migr_trans_to_valid_state(vfu_ctx, migr, device_state, notify);
 }
 
-bool
-migration_feature_supported(uint32_t flags) {
-    switch (flags & VFIO_DEVICE_FEATURE_MASK) {
+uint32_t
+migration_feature_flags(uint32_t feature) {
+    switch (feature) {
         case VFIO_DEVICE_FEATURE_MIGRATION:
-            return !(flags & VFIO_DEVICE_FEATURE_SET); // not supported for set
+            return VFIO_DEVICE_FEATURE_GET | VFIO_DEVICE_FEATURE_PROBE;
         case VFIO_DEVICE_FEATURE_MIG_DEVICE_STATE:
-            return true;
+            return VFIO_DEVICE_FEATURE_GET
+                | VFIO_DEVICE_FEATURE_SET
+                | VFIO_DEVICE_FEATURE_PROBE;
         default:
-            return false;
+            return 0;
     };
 }
 
 ssize_t
-migration_feature_get(vfu_ctx_t *vfu_ctx, uint32_t flags, void *buf)
+migration_feature_get(vfu_ctx_t *vfu_ctx, uint32_t feature, void *buf)
 {
+    assert(vfu_ctx != NULL);
+
     struct vfio_user_device_feature_migration *res;
     struct vfio_user_device_feature_mig_state *state;
 
-    switch (flags & VFIO_DEVICE_FEATURE_MASK) {
+    switch (feature) {
         case VFIO_DEVICE_FEATURE_MIGRATION:
             res = buf;
             // FIXME are these always supported? Can we consider to be
             //   "supported" if said support is just an empty callback?
+            //
+            // We don't need to return RUNNING or ERROR since they are always
+            //   supported.
             res->flags = VFIO_MIGRATION_STOP_COPY | VFIO_MIGRATION_PRE_COPY;
 
             return 0;
-
         case VFIO_DEVICE_FEATURE_MIG_DEVICE_STATE:
             state = buf;
             state->device_state = vfu_ctx->migration->state;
 
             return 0;
-
         default:
             return -EINVAL;
     };
 }
 
 ssize_t
-migration_feature_set(vfu_ctx_t *vfu_ctx, uint32_t flags, void *buf)
+migration_feature_set(vfu_ctx_t *vfu_ctx, uint32_t feature, void *buf)
 {
-    if (flags & VFIO_DEVICE_FEATURE_MIG_DEVICE_STATE) {
+    assert(vfu_ctx != NULL);
+
+    if (feature == VFIO_DEVICE_FEATURE_MIG_DEVICE_STATE) {
         struct vfio_user_device_feature_mig_state *res = buf;
         struct migration *migr = vfu_ctx->migration;
         
@@ -225,8 +233,11 @@ migration_feature_set(vfu_ctx_t *vfu_ctx, uint32_t flags, void *buf)
 ssize_t
 handle_mig_data_read(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg)
 {
+    assert(vfu_ctx != NULL);
+    assert(msg != NULL);
+
     struct migration *migr = vfu_ctx->migration;
-    struct vfio_user_mig_data_without_data *req = msg->in.iov.iov_base;
+    struct vfio_user_mig_data *req = msg->in.iov.iov_base;
 
     if (vfu_ctx->migration == NULL) {
         return -EINVAL;
@@ -234,47 +245,50 @@ handle_mig_data_read(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg)
 
     if (migr->state != VFIO_DEVICE_STATE_PRE_COPY
         && migr->state != VFIO_DEVICE_STATE_STOP_COPY) {
+        vfu_log(vfu_ctx, LOG_ERR, "bad state to read data: %d", migr->state);
+        return -EINVAL;
+    }
+
+    if (req->size > vfu_ctx->client_max_data_xfer_size) {
         return -EINVAL;
     }
 
     msg->out.iov.iov_len = msg->in.iov.iov_len + req->size;
     msg->out.iov.iov_base = calloc(1, msg->out.iov.iov_len);
 
-    struct vfio_user_mig_data_with_data *res = msg->out.iov.iov_base;
+    if (msg->out.iov.iov_base == NULL) {
+        return -EINVAL;
+    }
+
+    struct vfio_user_mig_data *res = msg->out.iov.iov_base;
 
     ssize_t ret = migr->callbacks.read_data(vfu_ctx, &res->data, req->size);
 
     res->size = ret;
     res->argsz = ret + msg->in.iov.iov_len;
 
-    if (ret < 0) {
-        return -1;
-    }
-
     return ret;
 }
 
 ssize_t
 handle_mig_data_write(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg)
 {
+    assert(vfu_ctx != NULL);
+    assert(msg != NULL);
+
     struct migration *migr = vfu_ctx->migration;
-    struct vfio_user_mig_data_with_data *req = msg->in.iov.iov_base;
+    struct vfio_user_mig_data *req = msg->in.iov.iov_base;
 
     if (vfu_ctx->migration == NULL) {
         return -EINVAL;
     }
 
     if (migr->state != VFIO_DEVICE_STATE_RESUMING) {
+        vfu_log(vfu_ctx, LOG_ERR, "bad state to write data: %d", migr->state);
         return -EINVAL;
     }
 
-    ssize_t ret = migr->callbacks.write_data(vfu_ctx, &req->data, req->size);
-
-    if (ret < 0) {
-        return -1;
-    }
-
-    return ret;
+    return migr->callbacks.write_data(vfu_ctx, &req->data, req->size);
 }
 
 bool
diff --git a/lib/migration.h b/lib/migration.h
index f817f66..61d7a1e 100644
--- a/lib/migration.h
+++ b/lib/migration.h
@@ -48,14 +48,14 @@ struct migration *
 init_migration(const vfu_migration_callbacks_t *callbacks,
                uint64_t flags, int *err);
 
-bool
-migration_feature_supported(uint32_t flags);
+uint32_t
+migration_feature_flags(uint32_t feature);
 
 ssize_t
-migration_feature_get(vfu_ctx_t *vfu_ctx, uint32_t flags, void *buf);
+migration_feature_get(vfu_ctx_t *vfu_ctx, uint32_t feature, void *buf);
 
 ssize_t
-migration_feature_set(vfu_ctx_t *vfu_ctx, uint32_t flags, void *buf);
+migration_feature_set(vfu_ctx_t *vfu_ctx, uint32_t feature, void *buf);
 
 ssize_t
 handle_mig_data_read(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg);
@@ -80,7 +80,7 @@ uint64_t
 migration_get_flags(struct migration *migr);
 
 MOCK_DECLARE(void, migr_state_transition, struct migration *migr,
-             enum vfio_device_mig_state state);
+             enum vfio_user_device_mig_state state);
 
 MOCK_DECLARE(bool, vfio_migr_state_transition_is_valid, uint32_t from,
              uint32_t to);
diff --git a/lib/migration_priv.h b/lib/migration_priv.h
index ac58b1f..c20982c 100644
--- a/lib/migration_priv.h
+++ b/lib/migration_priv.h
@@ -35,22 +35,20 @@
 
 struct migration {
     uint64_t flags;
-    enum vfio_device_mig_state state;
+    enum vfio_user_device_mig_state state;
     size_t pgsize;
     vfu_migration_callbacks_t callbacks;
 };
 
-/* valid migration state transitions 
-   indexed by vfio_device_mig_state enum */
-static const bool transitions[8][8] = {
-    {0, 0, 0, 0, 0, 0, 0, 0}, // ERROR
-    {0, 0, 1, 1, 1, 0, 0, 0}, // STOP
-    {0, 1, 0, 0, 0, 0, 1, 0}, // RUNNING
-    {0, 1, 0, 0, 0, 0, 0, 0}, // STOP_COPY
-    {0, 1, 0, 0, 0, 0, 0, 0}, // RESUMING
-    {0, 0, 0, 0, 0, 0, 0, 0}, // RUNNING_P2P
-    {0, 0, 1, 1, 0, 0, 0, 0}, // PRE_COPY
-    {0, 0, 0, 0, 0, 0, 0, 0}  // PRE_COPY_P2P
+/* 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] = {
+    // ERROR    STOP        RUNNING     STOP_COPY
+    0b00000000, 0b00111000, 0b01000010, 0b01000000,
+    // RESUMING RUNNING_P2P PRE_COPY    PRE_COPY_P2P 
+    0b01000000, 0b00000000, 0b00110000, 0b00000000
 };
 
 MOCK_DECLARE(vfu_migr_state_t, migr_state_vfio_to_vfu, uint32_t device_state);