refactor: respond to John's review

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

8 files changed, 143 insertions, 327 deletions

diff --git a/include/libvfio-user.h b/include/libvfio-user.h
index c1c844a..af16e60 100644
--- a/include/libvfio-user.h
+++ b/include/libvfio-user.h
@@ -623,9 +623,8 @@ typedef struct {
 } vfu_migration_callbacks_t;
 
 int
-vfu_setup_device_migration_callbacks(vfu_ctx_t *vfu_ctx, uint64_t flags,
-                                     const vfu_migration_callbacks_t
-                                     *callbacks);
+vfu_setup_device_migration_callbacks(vfu_ctx_t *vfu_ctx,
+    const vfu_migration_callbacks_t *callbacks);
 
 /**
  * Triggers an interrupt.
diff --git a/include/vfio-user.h b/include/vfio-user.h
index c2d2503..a994718 100644
--- a/include/vfio-user.h
+++ b/include/vfio-user.h
@@ -66,8 +66,7 @@ enum vfio_user_command {
     VFIO_USER_DMA_READ                  = 11,
     VFIO_USER_DMA_WRITE                 = 12,
     VFIO_USER_DEVICE_RESET              = 13,
-    VFIO_USER_DIRTY_PAGES               = 14, // TODO remove
-    VFIO_USER_REGION_WRITE_MULTI        = 15, // TODO implement
+    VFIO_USER_REGION_WRITE_MULTI        = 15,
     VFIO_USER_DEVICE_FEATURE            = 16,
     VFIO_USER_MIG_DATA_READ             = 17,
     VFIO_USER_MIG_DATA_WRITE            = 18,
@@ -212,9 +211,9 @@ struct vfio_user_device_feature_dma_logging_range {
 } __attribute__((packed));
 
 struct vfio_user_device_feature_dma_logging_control {
-    uint64_t                                          page_size;
-    uint32_t                                          num_ranges;
-    uint32_t                                          reserved;
+    uint64_t page_size;
+    uint32_t num_ranges;
+    uint32_t reserved;
     struct vfio_user_device_feature_dma_logging_range ranges[];
 } __attribute__((packed));
 
@@ -235,7 +234,7 @@ struct vfio_user_device_feature_dma_logging_report {
 struct vfio_user_device_feature {
 	uint32_t	argsz;
 	uint32_t	flags;
-#ifndef VFIO_REGION_TYPE_MIGRATION_DEPRECATED
+#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[] */
@@ -282,10 +281,9 @@ enum vfio_user_device_mig_state {
 	VFIO_USER_DEVICE_STATE_RUNNING_P2P = 5,
 	VFIO_USER_DEVICE_STATE_PRE_COPY = 6,
 	VFIO_USER_DEVICE_STATE_PRE_COPY_P2P = 7,
+    VFIO_USER_DEVICE_NUM_STATES = 8,
 };
 
-#define VFIO_USER_DEVICE_NUM_STATES 8
-
 struct vfio_user_mig_data {
     uint32_t    argsz;
     uint32_t    size;
diff --git a/lib/dma.c b/lib/dma.c
index c4d8a58..92d87f1 100644
--- a/lib/dma.c
+++ b/lib/dma.c
@@ -258,7 +258,7 @@ dma_map_region(dma_controller_t *dma, dma_memory_region_t *region)
     return 0;
 }
 
-static ssize_t
+ssize_t
 get_bitmap_size(size_t region_size, size_t pgsize)
 {
     if (pgsize == 0) {
@@ -288,202 +288,6 @@ dirty_page_logging_start_on_region(dma_memory_region_t *region, size_t pgsize)
     return 0;
 }
 
-static dma_memory_region_t *
-find_region(dma_controller_t *dma, uint64_t iova, uint64_t length) {
-    for (size_t i = 0; i < (size_t)dma->nregions; i++) {
-        if ((uint64_t)dma->regions[i].info.iova.iov_base == iova &&
-            dma->regions[i].info.iova.iov_len == length) {
-            return &dma->regions[i];
-        }
-    }
-
-    return NULL;
-}
-
-bool
-is_dma_feature(uint32_t feature) {
-    switch (feature) {
-        case VFIO_DEVICE_FEATURE_DMA_LOGGING_START:
-        case VFIO_DEVICE_FEATURE_DMA_LOGGING_STOP:
-        case VFIO_DEVICE_FEATURE_DMA_LOGGING_REPORT:
-            return true;
-    }
-
-    return false;
-}
-
-ssize_t
-dma_get_request_bitmap_size(size_t length, void *buf) {
-    if (length != sizeof(struct vfio_user_device_feature_dma_logging_report)) {
-        return ERROR_INT(EINVAL);
-    }
-
-    struct vfio_user_device_feature_dma_logging_report *req = buf;
-
-    return get_bitmap_size(req->length, req->page_size);
-}
-
-ssize_t
-dma_feature_get(vfu_ctx_t *vfu_ctx, uint32_t feature, void *buf)
-{
-    assert(vfu_ctx != NULL);
-    assert(feature == VFIO_DEVICE_FEATURE_DMA_LOGGING_REPORT);
-
-    struct dma_controller *dma = vfu_ctx->dma;
-
-    assert(dma != NULL);
-
-    struct vfio_user_device_feature_dma_logging_report *req = buf;
-
-    ssize_t bitmap_size = get_bitmap_size(req->length, req->page_size);
-
-    char* bitmap = (char*) buf
-        + sizeof(struct vfio_user_device_feature_dma_logging_report);
-
-    return dma_controller_dirty_page_get(dma,
-                                         (vfu_dma_addr_t) req->iova,
-                                         req->length,
-                                         req->page_size,
-                                         bitmap_size,
-                                         bitmap);
-}
-
-/*
- * Currently we only support IOVA ranges that correspond exactly to a region.
- * Also, once DMA logging has been started on a certain subset of the regions,
- * it must be stopped on all of those regions at the same time before any other
- * regions can start logging.
- */
-ssize_t
-dma_feature_set(vfu_ctx_t *vfu_ctx, uint32_t feature, void *buf)
-{
-    assert(vfu_ctx != NULL);
-    assert(feature == VFIO_DEVICE_FEATURE_DMA_LOGGING_START ||
-           feature == VFIO_DEVICE_FEATURE_DMA_LOGGING_STOP);
-
-    struct dma_controller *dma = vfu_ctx->dma;
-
-    assert(dma != NULL);
-    
-    struct vfio_user_device_feature_dma_logging_control *req = buf;
-
-    size_t num_regions = req->num_ranges;
-    bool is_all_regions = false;
-
-    // If num_ranges is zero, we consider every region to be selected.
-    if (num_regions == 0) {
-        num_regions = dma->nregions;
-        is_all_regions = true;
-    }
-
-    dma_memory_region_t *region;
-
-    if (feature == VFIO_DEVICE_FEATURE_DMA_LOGGING_START) {
-        if (req->page_size == 0) {
-            return ERROR_INT(EINVAL);
-        }
-
-        if (dma->dirty_pgsize > 0) {
-            if (dma->dirty_pgsize != req->page_size) {
-                return ERROR_INT(EINVAL);
-            }
-            return 0;
-        }
-
-        for (size_t i = 0; i < num_regions; i++) {
-            if (is_all_regions) {
-                region = &dma->regions[i];
-            } else {
-                region = find_region(dma, req->ranges[i].iova,
-                                     req->ranges[i].length);
-            }
-
-            if (region == NULL) {
-                return ERROR_INT(EINVAL);
-            }
-
-            if (region->fd == -1) {
-                if (is_all_regions) {
-                    /*
-                    * DMA regions not mapped by the server are not dirty tracked
-                    * (the client must track changes via handling
-                    * VFIO_USER_DMA_WRITE). Therefore, we don't enable dirty page
-                    * tracking for such regions. See #627.
-                    */
-                    continue;
-                } else {
-                    /*
-                     * Explicitly trying to enable dirty page tracking on a
-                     * non-mapped region is invalid.
-                     */
-                    return ERROR_INT(EINVAL);
-                }
-            }
-            
-            if (dirty_page_logging_start_on_region(region,
-                                                   req->page_size) < 0) {
-                int _errno = errno;
-                size_t j;
-
-                for (j = 0; j < i; j++) {
-                    if (is_all_regions) {
-                        region = &dma->regions[j];
-                    } else {
-                        region = find_region(dma, req->ranges[j].iova,
-                                             req->ranges[j].length);
-                    }
-                    free(region->dirty_bitmap);
-                    region->dirty_bitmap = NULL;
-                }
-
-                return ERROR_INT(_errno);
-            }
-        }
-
-        dma->dirty_pgsize = req->page_size;
-
-        vfu_log(dma->vfu_ctx, LOG_DEBUG, "dirty pages: started logging");
-    } else if (feature == VFIO_DEVICE_FEATURE_DMA_LOGGING_STOP) {
-        if (dma->dirty_pgsize == 0) {
-            return 0;
-        }
-
-        for (size_t i = 0; i < num_regions; i++) {
-            if (is_all_regions) {
-                region = &dma->regions[i];
-            } else {
-                region = find_region(dma, req->ranges[i].iova,
-                                     req->ranges[i].length);
-            }
-
-            if (region == NULL) {
-                return ERROR_INT(EINVAL);
-            }
-
-            if (region->fd == -1) {
-                if (is_all_regions) {
-                    continue;
-                } else {
-                    return ERROR_INT(EINVAL);
-                }
-            }
-
-            if (region->dirty_bitmap == NULL) {
-                return ERROR_INT(EINVAL);
-            }
-
-            free(region->dirty_bitmap);
-            region->dirty_bitmap = NULL;
-        }
-
-        dma->dirty_pgsize = 0;
-
-        vfu_log(dma->vfu_ctx, LOG_DEBUG, "dirty pages: stopped logging");
-    }
-
-    return 0;
-}
-
 int
 MOCK_DEFINE(dma_controller_add_region)(dma_controller_t *dma,
                                        vfu_dma_addr_t dma_addr, uint64_t size,
@@ -804,11 +608,6 @@ dma_controller_dirty_page_get(dma_controller_t *dma, vfu_dma_addr_t addr,
         return ERROR_INT(EINVAL);
     }
 
-    if (region->dirty_bitmap == NULL) {
-        vfu_log(dma->vfu_ctx, LOG_ERR, "region %d is not logged", sg.region);
-        return ERROR_INT(EINVAL);
-    }
-
     for (i = 0; i < (size_t)bitmap_size; i++) {
         uint8_t val = region->dirty_bitmap[i];
         uint8_t *outp = (uint8_t *)&bitmap[i];
diff --git a/lib/dma.h b/lib/dma.h
index 55ae72c..a00ff4e 100644
--- a/lib/dma.h
+++ b/lib/dma.h
@@ -135,17 +135,8 @@ MOCK_DECLARE(int, dma_controller_remove_region, dma_controller_t *dma,
 MOCK_DECLARE(void, dma_controller_unmap_region, dma_controller_t *dma,
              dma_memory_region_t *region);
 
-bool
-is_dma_feature(uint32_t feature);
-
-ssize_t
-dma_get_request_bitmap_size(size_t length, void *buf);
-
 ssize_t
-dma_feature_get(vfu_ctx_t *vfu_ctx, uint32_t feature, void *buf);
-
-ssize_t
-dma_feature_set(vfu_ctx_t *vfu_ctx, uint32_t feature, void *buf);
+get_bitmap_size(size_t region_size, size_t pgsize);
 
 // Helper for dma_addr_to_sgl() slow path.
 int
@@ -373,7 +364,7 @@ dma_sgl_put(dma_controller_t *dma, dma_sg_t *sgl, size_t cnt)
         region = &dma->regions[sg->region];
 
         if (sg->writeable) {
-            if (dma->dirty_pgsize > 0 && region->dirty_bitmap != NULL) {
+            if (dma->dirty_pgsize > 0) {
                 _dma_mark_dirty(dma, region, sg);
             }
         }
diff --git a/lib/libvfio-user.c b/lib/libvfio-user.c
index ff57b72..678c08c 100644
--- a/lib/libvfio-user.c
+++ b/lib/libvfio-user.c
@@ -908,7 +908,7 @@ device_reset(vfu_ctx_t *vfu_ctx, vfu_reset_type_t reason)
 }
 
 static uint32_t
-device_feature_flags(uint32_t feature) {
+device_feature_flags_supported(uint32_t feature) {
     switch (feature) {
         case VFIO_DEVICE_FEATURE_MIGRATION:
         case VFIO_DEVICE_FEATURE_DMA_LOGGING_REPORT:
@@ -925,6 +925,29 @@ device_feature_flags(uint32_t feature) {
     };
 }
 
+static bool
+is_migration_feature(uint32_t feature) {
+    switch (feature) {
+        case VFIO_DEVICE_FEATURE_MIGRATION:
+        case VFIO_DEVICE_FEATURE_MIG_DEVICE_STATE:
+            return true;
+    }
+
+    return false;
+}
+
+static bool
+is_dma_feature(uint32_t feature) {
+    switch (feature) {
+        case VFIO_DEVICE_FEATURE_DMA_LOGGING_START:
+        case VFIO_DEVICE_FEATURE_DMA_LOGGING_STOP:
+        case VFIO_DEVICE_FEATURE_DMA_LOGGING_REPORT:
+            return true;
+    }
+
+    return false;
+}
+
 static int
 handle_device_feature(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg)
 {
@@ -939,7 +962,7 @@ handle_device_feature(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg)
     struct vfio_user_device_feature *req = msg->in.iov.iov_base;
 
     uint32_t supported_flags =
-        device_feature_flags(req->flags & VFIO_DEVICE_FEATURE_MASK);
+        device_feature_flags_supported(req->flags & VFIO_DEVICE_FEATURE_MASK);
 
     if ((req->flags & supported_flags) !=
         (req->flags & ~VFIO_DEVICE_FEATURE_MASK) || supported_flags == 0) {
@@ -977,23 +1000,38 @@ handle_device_feature(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg)
             memcpy(msg->out.iov.iov_base, msg->in.iov.iov_base,
                    sizeof(struct vfio_user_device_feature));
 
-            ret = migration_feature_get(vfu_ctx, feature,
-                                        msg->out.iov.iov_base +
-                                        sizeof(struct vfio_user_device_feature));
-
             struct vfio_user_device_feature *res = msg->out.iov.iov_base;
 
-            if (ret < 0) {
-                msg->out.iov.iov_len = 0;
-            } else {
-                res->argsz = sizeof(struct vfio_user_device_feature)
+            res->argsz = sizeof(struct vfio_user_device_feature)
                     + sizeof(struct vfio_user_device_feature_migration);
+
+            if (feature == VFIO_DEVICE_FEATURE_MIGRATION) {
+                struct vfio_user_device_feature_migration *mig =
+                    (void*)res->data;
+
+                // 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.
+                mig->flags = VFIO_MIGRATION_STOP_COPY | VFIO_MIGRATION_PRE_COPY;
+
+                ret = 0;
+            } else {
+                assert(feature == VFIO_DEVICE_FEATURE_MIG_DEVICE_STATE);
+
+                struct vfio_user_device_feature_mig_state *state =
+                    (void*)res->data;
+                
+                state->device_state = migration_get_state(vfu_ctx);
+
+                ret = 0;
             }
         } else if (is_dma_feature(feature)) {
-            ssize_t bitmap_size = dma_get_request_bitmap_size(
-                req->argsz - sizeof(struct vfio_user_device_feature),
-                req->data
-            );
+            struct vfio_user_device_feature_dma_logging_report *rep =
+                (void*)req->data;
+
+            ssize_t bitmap_size = get_bitmap_size(rep->length, rep->page_size);
 
             msg->out.iov.iov_len = sizeof(struct vfio_user_device_feature)
                 + sizeof(struct vfio_user_device_feature_dma_logging_report)
@@ -1012,7 +1050,20 @@ handle_device_feature(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg)
 
             res->argsz = msg->out.iov.iov_len;
 
-            ret = dma_feature_get(vfu_ctx, feature, res->data);
+            dma_controller_t *dma = vfu_ctx->dma;
+
+            assert(dma != NULL);
+
+            char* bitmap = (char*) msg->out.iov.iov_base
+                + sizeof(struct vfio_user_device_feature)
+                + sizeof(struct vfio_user_device_feature_dma_logging_report);
+
+            ret = dma_controller_dirty_page_get(dma,
+                                                (vfu_dma_addr_t) rep->iova,
+                                                rep->length,
+                                                rep->page_size,
+                                                bitmap_size,
+                                                bitmap);
 
             if (ret < 0) {
                 msg->out.iov.iov_len = 0;
@@ -1034,9 +1085,27 @@ handle_device_feature(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg)
         struct vfio_user_device_feature *res = msg->out.iov.iov_base;
 
         if (is_migration_feature(feature)) {
-            ret = migration_feature_set(vfu_ctx, feature, res->data);
+            assert(feature == VFIO_DEVICE_FEATURE_MIG_DEVICE_STATE);
+            
+            struct vfio_user_device_feature_mig_state *state = (void*)res->data;
+
+            ret = migration_set_state(vfu_ctx, state->device_state);
         } else if (is_dma_feature(feature)) {
-            ret = dma_feature_set(vfu_ctx, feature, res->data);
+            dma_controller_t *dma = vfu_ctx->dma;
+
+            assert(dma != NULL);
+
+            if (feature == VFIO_DEVICE_FEATURE_DMA_LOGGING_START) {
+                struct vfio_user_device_feature_dma_logging_control *ctl =
+                    (void*)res->data;
+                ret = dma_controller_dirty_page_logging_start(dma,
+                                                              ctl->page_size);
+            } else {
+                assert(feature == VFIO_DEVICE_FEATURE_DMA_LOGGING_STOP);
+
+                dma_controller_dirty_page_logging_stop(dma);
+                ret = 0;
+            }
         } else {
             return ERROR_INT(EINVAL);
         }
@@ -1207,10 +1276,6 @@ handle_request(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg)
         ret = device_reset(vfu_ctx, VFU_RESET_DEVICE);
         break;
 
-    case VFIO_USER_DIRTY_PAGES:
-        vfu_log(vfu_ctx, LOG_ERR, "VFIO_USER_DIRTY_PAGES deprecated");
-        break;
-
     case VFIO_USER_DEVICE_FEATURE:
         ret = handle_device_feature(vfu_ctx, msg);
         break;
@@ -1326,7 +1391,6 @@ MOCK_DEFINE(cmd_allowed_when_stopped_and_copying)(uint16_t cmd)
 {
     return cmd == VFIO_USER_REGION_READ ||
            cmd == VFIO_USER_REGION_WRITE ||
-           cmd == VFIO_USER_DIRTY_PAGES ||
            cmd == VFIO_USER_DEVICE_FEATURE ||
            cmd == VFIO_USER_MIG_DATA_READ;
 }
@@ -1996,7 +2060,7 @@ vfu_setup_irq_state_callback(vfu_ctx_t *vfu_ctx, enum vfu_dev_irq_type type,
 }
 
 EXPORT int
-vfu_setup_device_migration_callbacks(vfu_ctx_t *vfu_ctx, uint64_t flags,
+vfu_setup_device_migration_callbacks(vfu_ctx_t *vfu_ctx,
                                      const vfu_migration_callbacks_t *callbacks)
 {
     int ret = 0;
@@ -2010,7 +2074,7 @@ vfu_setup_device_migration_callbacks(vfu_ctx_t *vfu_ctx, uint64_t flags,
         return ERROR_INT(EINVAL);
     }
 
-    vfu_ctx->migration = init_migration(callbacks, flags, &ret);
+    vfu_ctx->migration = init_migration(callbacks, &ret);
     if (vfu_ctx->migration == NULL) {
         vfu_log(vfu_ctx, LOG_ERR, "failed to initialize device migration");
         return ERROR_INT(ret);
diff --git a/lib/migration.c b/lib/migration.c
index e450576..5617a53 100644
--- a/lib/migration.c
+++ b/lib/migration.c
@@ -52,8 +52,7 @@ MOCK_DEFINE(vfio_migr_state_transition_is_valid)(uint32_t from, uint32_t to)
  * in vfu_ctx_t.
  */
 struct migration *
-init_migration(const vfu_migration_callbacks_t *callbacks,
-               uint64_t flags, int *err)
+init_migration(const vfu_migration_callbacks_t *callbacks, int *err)
 {
     struct migration *migr;
 
@@ -63,8 +62,6 @@ init_migration(const vfu_migration_callbacks_t *callbacks,
         return NULL;
     }
 
-    migr->flags = flags;
-
     /*
      * FIXME: incorrect, if the client doesn't give a pgsize value, it means "no
      * migration support", handle this
@@ -92,7 +89,6 @@ MOCK_DEFINE(migr_state_transition)(struct migration *migr,
                                    enum vfio_user_device_mig_state state)
 {
     assert(migr != NULL);
-    /* FIXME validate that state transition */
     migr->state = state;
 }
 
@@ -149,7 +145,7 @@ MOCK_DEFINE(migr_trans_to_valid_state)(vfu_ctx_t *vfu_ctx, struct migration *mig
             return ret;
         }
     }
-    migr_state_transition(migr, device_state); // TODO confused
+    migr_state_transition(migr, device_state);
     return 0;
 }
 
@@ -170,75 +166,34 @@ 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
-is_migration_feature(uint32_t feature) {
-    switch (feature) {
-        case VFIO_DEVICE_FEATURE_MIGRATION:
-        case VFIO_DEVICE_FEATURE_MIG_DEVICE_STATE:
-            return true;
-    }
-
-    return false;
-}
-
-ssize_t
-migration_feature_get(vfu_ctx_t *vfu_ctx, uint32_t feature, void *buf)
+size_t
+migration_get_state(vfu_ctx_t *vfu_ctx)
 {
-    assert(vfu_ctx != NULL);
-
-    struct vfio_user_device_feature_migration *res;
-    struct vfio_user_device_feature_mig_state *state;
-
-    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 ERROR_INT(EINVAL);
-    };
+    return vfu_ctx->migration->state;
 }
 
 ssize_t
-migration_feature_set(vfu_ctx_t *vfu_ctx, uint32_t feature, void *buf)
+migration_set_state(vfu_ctx_t *vfu_ctx, uint32_t device_state)
 {
-    assert(vfu_ctx != NULL);
+    struct migration *migr = vfu_ctx->migration;
+    uint32_t state;
+    ssize_t ret = 0;
+    
+    if (device_state > VFIO_USER_DEVICE_NUM_STATES) {
+        return ERROR_INT(EINVAL);
+    }
+    
+    while (migr->state != device_state && ret == 0) {
+        state = next_state[migr->state][device_state];
 
-    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 = 0;
-        
-        if (res->device_state > VFIO_USER_DEVICE_NUM_STATES) {
+        if (state == VFIO_USER_DEVICE_STATE_ERROR) {
             return ERROR_INT(EINVAL);
         }
-        
-        while (migr->state != res->device_state && ret == 0) {
-            state = next_state[migr->state][res->device_state];
-
-            if (state == VFIO_USER_DEVICE_STATE_ERROR) {
-                return ERROR_INT(EINVAL);
-            }
-
-            ret = handle_device_state(vfu_ctx, migr, state, true);
-        };
-        
-        return ret;
-    }
 
-    return ERROR_INT(EINVAL);
+        ret = handle_device_state(vfu_ctx, migr, state, true);
+    };
+    
+    return ret;
 }
 
 ssize_t
@@ -247,6 +202,10 @@ handle_mig_data_read(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_mig_data)) {
+        return ERROR_INT(EINVAL);
+    }
+
     struct migration *migr = vfu_ctx->migration;
     struct vfio_user_mig_data *req = msg->in.iov.iov_base;
 
@@ -294,6 +253,10 @@ handle_mig_data_write(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_mig_data)) {
+        return ERROR_INT(EINVAL);
+    }
+
     struct migration *migr = vfu_ctx->migration;
     struct vfio_user_mig_data *req = msg->in.iov.iov_base;
 
@@ -307,6 +270,12 @@ handle_mig_data_write(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg)
         return ERROR_INT(EINVAL);
     }
 
+    if (req->size > vfu_ctx->client_max_data_xfer_size) {
+        vfu_log(vfu_ctx, LOG_ERR, "transfer size exceeds limit (%d > %ld)",
+                req->size, vfu_ctx->client_max_data_xfer_size);
+        return ERROR_INT(EINVAL);
+    }
+
     return migr->callbacks.write_data(vfu_ctx, &req->data, req->size);
 }
 
diff --git a/lib/migration.h b/lib/migration.h
index d164a3a..0d0dd76 100644
--- a/lib/migration.h
+++ b/lib/migration.h
@@ -45,17 +45,13 @@
 #include "private.h"
 
 struct migration *
-init_migration(const vfu_migration_callbacks_t *callbacks,
-               uint64_t flags, int *err);
+init_migration(const vfu_migration_callbacks_t *callbacks, int *err);
 
-bool
-is_migration_feature(uint32_t feature);
-
-ssize_t
-migration_feature_get(vfu_ctx_t *vfu_ctx, uint32_t feature, void *buf);
+size_t
+migration_get_state(vfu_ctx_t *vfu_ctx);
 
 ssize_t
-migration_feature_set(vfu_ctx_t *vfu_ctx, uint32_t feature, void *buf);
+migration_set_state(vfu_ctx_t *vfu_ctx, uint32_t device_state);
 
 ssize_t
 handle_mig_data_read(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg);
diff --git a/lib/migration_priv.h b/lib/migration_priv.h
index 639925e..bb4e5f0 100644
--- a/lib/migration_priv.h
+++ b/lib/migration_priv.h
@@ -34,17 +34,17 @@
 #include <linux/vfio.h>
 
 struct migration {
-    uint64_t flags;
     enum vfio_user_device_mig_state state;
     size_t pgsize;
     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
+ * This defines valid migration state transitions. Each element in the array
+ * corresponds to a FROM state and each bit of the element to 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        -> {}