expose migration region (#305)

This patch exposes the fact that live migration is implemented as a
special device region. Hiding this from the user doesn't offer much
benefit since it only takes just a little bit of extra code for the user
to handle it as a region. We do keep the migration callback
functionality since this feature substantially simplifies supporting
live migration from the device implementation's perspective.

Signed-off-by: Thanos Makatos <thanos.makatos@nutanix.com>
Co-authored-by: John Levon <john.levon@nutanix.com>

10 files changed, 463 insertions, 88 deletions

diff --git a/include/libvfio-user.h b/include/libvfio-user.h
index 75b5943..1cc62d9 100644
--- a/include/libvfio-user.h
+++ b/include/libvfio-user.h
@@ -227,7 +227,12 @@ typedef ssize_t (vfu_region_access_cb_t)(vfu_ctx_t *vfu_ctx, char *buf,
  * given callback.  However, the callback can still be invoked, even on a
  * mappable area, if the client chooses to call VFIO_USER_REGION_READ/WRITE.
  *
- * A VFU_PCI_DEV_CFG_REGION_IDX region, corresponding to PCI config space, has
+ * The following regions are special and are explained below:
+ *  - VFU_PCI_DEV_CFG_REGION_IDX,
+ *  - VFU_PCI_DEV_MIGR_REGION_IDX, and
+ *  - VFU_GENERIC_DEV_MIGR_REG_IDX.
+ *
+ * Region VFU_PCI_DEV_CFG_REGION_IDX, corresponding to PCI config space, has
  * special handling:
  *
  *  - the @size argument is ignored: the region size is always the size defined
@@ -240,6 +245,24 @@ typedef ssize_t (vfu_region_access_cb_t)(vfu_ctx_t *vfu_ctx, char *buf,
  *    and writes are an error
  *  - otherwise, the callback is expected to handle the access
  *
+ * Regions VFU_PCI_DEV_MIGR_REGION_IDX and VFU_GENERIC_DEV_MIGR_REG_IDX,
+ * corresponding to the migration region, enable live migration support for
+ * the device. The migration region must contain at the beginning the migration
+ * registers (struct vfio_device_migration_info defined in <linux/vfio.h>) and
+ * the remaining part of the region can be arbitrarily used by the device
+ * implementation. The region provided must have at least
+ * vfu_get_migr_register_area_size() bytes available at the start of the region
+ * (this size is guaranteed to be page-aligned). If mmap_areas is given, it
+ * must _not_ include this part of the region.
+ *
+ * libvfio-user offers two ways for the migration region to be used:
+ *  1. natively: the device implementation must handle accesses to the
+ *      migration registers and migration data via the region callbacks. The
+ *      semantics of these registers are explained in <linux/vfio.h>.
+ *  2. via the vfu_migration_t callbacks: the device implementation registers
+ *      a set of callbacks by calling vfu_setup_device_migration. The region's
+ *      read/write callbacks are never called.
+ *
  * @vfu_ctx: the libvfio-user context
  * @region_idx: region index
  * @size: size of the region
@@ -262,6 +285,13 @@ vfu_setup_region(vfu_ctx_t *vfu_ctx, int region_idx, size_t size,
                  int fd);
 
 /*
+ * Returns the size of the area needed to hold the migration registers at the
+ * beginning of the migration region; guaranteed to be page aligned.
+ */
+size_t
+vfu_get_migr_register_area_size(void);
+
+/*
  * Callback function that is called when the guest resets the device.
  */
 typedef int (vfu_reset_cb_t)(vfu_ctx_t *vfu_ctx);
@@ -348,6 +378,8 @@ typedef enum {
 } vfu_migr_state_t;
 
 
+#define VFU_MIGR_CALLBACKS_VERS 1
+
 /*
  * Callbacks during the pre-copy and stop-and-copy phases.
  *
@@ -363,6 +395,11 @@ typedef enum {
  */
 typedef struct {
 
+    /*
+     * Set it to VFU_MIGR_CALLBACKS_VERS.
+     */
+    int version;
+
     /* migration state transition callback */
     /* TODO rename to vfu_migration_state_transition_callback */
     /* FIXME maybe we should create a single callback and pass the state? */
@@ -422,21 +459,23 @@ typedef struct {
 
 } vfu_migration_callbacks_t;
 
-typedef struct {
-    size_t                      size;
-    vfu_migration_callbacks_t   callbacks;
-    struct iovec                *mmap_areas;
-    uint32_t                    nr_mmap_areas;
-} vfu_migration_t;
-
-//TODO: Re-visit once migration support is done.
 /**
- * Enable support for device migration.
+ * vfu_setup_device_migration provides an abstraction over the migration
+ * protocol: the user specifies a set of callbacks which are called in response
+ * to client accesses of the migration region; the migration region read/write
+ * callbacks are not called after this function call. Offsets in callbacks are
+ * relative to @data_offset.
+ * 
  * @vfu_ctx: the libvfio-user context
- * @migration: information required to migrate device
+ * @callbacks: migration callbacks
+ * @data_offset: offset in the migration region where data begins.
+ *
+ * @returns 0 on success, -1 on error, sets errno.
  */
 int
-vfu_setup_device_migration(vfu_ctx_t *vfu_ctx, vfu_migration_t *migration);
+vfu_setup_device_migration_callbacks(vfu_ctx_t *vfu_ctx,
+                                     const vfu_migration_callbacks_t *callbacks,
+                                     uint64_t data_offset);
 
 /**
  * Triggers an interrupt.
@@ -550,6 +589,18 @@ vfu_dma_write(vfu_ctx_t *vfu_ctx, dma_sg_t *sg, void *data);
 
 /*
  * Supported PCI regions.
+ *
+ * Note: in VFIO, each region starts at a terabyte offset
+ * (VFIO_PCI_INDEX_TO_OFFSET) and because Linux supports up to 128 TB of user
+ * space virtual memory, there can be up to 128 device regions. PCI regions are
+ * fixed and in retrospect this choice has proven to be problematic because
+ * devices might contain potentially unused regions. New regions can now be
+ * positioned anywhere by using the VFIO_REGION_INFO_CAP_TYPE capability.  In
+ * vfio-user we don't have this problem because the region index is just an
+ * identifier: the VMM memory maps a file descriptor that is passed to it and
+ * the mapping offset is derived from the mmap_areas offset value, rather than a
+ * static mapping from region index to offset. Thus, additional regions can
+ * have static indexes in vfio-user.
  */
 enum {
     VFU_PCI_DEV_BAR0_REGION_IDX,
@@ -561,6 +612,7 @@ enum {
     VFU_PCI_DEV_ROM_REGION_IDX,
     VFU_PCI_DEV_CFG_REGION_IDX,
     VFU_PCI_DEV_VGA_REGION_IDX,
+    VFU_PCI_DEV_MIGR_REGION_IDX,
     VFU_PCI_DEV_NUM_REGIONS,
 };
 
@@ -571,6 +623,11 @@ typedef enum {
     VFU_PCI_TYPE_EXPRESS
 } vfu_pci_type_t;
 
+enum {
+    VFU_GENERIC_DEV_MIGR_REGION_IDX,
+    VFU_GENERIC_DEV_NUM_REGIONS
+};
+
 /**
  * Initialize the context for a PCI device. This function must be called only
  * once per libvfio-user context.
diff --git a/lib/irq.c b/lib/irq.c
index c3428e3..5c7a1e4 100644
--- a/lib/irq.c
+++ b/lib/irq.c
@@ -333,6 +333,7 @@ handle_device_get_irq_info(vfu_ctx_t *vfu_ctx, uint32_t size,
     assert(irq_info_out != NULL);
 
     if (size != sizeof *irq_info_in || size != irq_info_in->argsz) {
+        vfu_log(vfu_ctx, LOG_WARNING, "IRQ info size %d", size);
         return -EINVAL;
     }
 
diff --git a/lib/libvfio-user.c b/lib/libvfio-user.c
index 868af07..00379a6 100644
--- a/lib/libvfio-user.c
+++ b/lib/libvfio-user.c
@@ -215,7 +215,7 @@ region_access(vfu_ctx_t *vfu_ctx, size_t region_index, char *buf,
 
     if (region_index == VFU_PCI_DEV_CFG_REGION_IDX) {
         ret = pci_config_space_access(vfu_ctx, buf, count, offset, is_write);
-    } else if (is_migr_reg(vfu_ctx, region_index)) {
+    } else if (is_migr_reg(vfu_ctx, region_index) && vfu_ctx->migration != NULL) {
         ret = migration_region_access(vfu_ctx, buf, count, offset, is_write);
     } else {
         vfu_region_access_cb_t *cb = vfu_ctx->reg_info[region_index].cb;
@@ -1128,6 +1128,7 @@ vfu_create_ctx(vfu_trans_t trans, const char *path, int flags, void *pvt,
 {
     vfu_ctx_t *vfu_ctx = NULL;
     int err = 0;
+    size_t i;
 
     //FIXME: Validate arguments.
 
@@ -1166,7 +1167,7 @@ vfu_create_ctx(vfu_trans_t trans, const char *path, int flags, void *pvt,
      * to seperate migration region from standard regions in vfu_ctx.reg_info
      * and move it into vfu_ctx.migration.
      */
-    vfu_ctx->nr_regions = VFU_PCI_DEV_NUM_REGIONS + 1;
+    vfu_ctx->nr_regions = VFU_PCI_DEV_NUM_REGIONS;
     vfu_ctx->reg_info = calloc(vfu_ctx->nr_regions, sizeof *vfu_ctx->reg_info);
     if (vfu_ctx->reg_info == NULL) {
         err = -ENOMEM;
@@ -1190,6 +1191,10 @@ vfu_create_ctx(vfu_trans_t trans, const char *path, int flags, void *pvt,
         vfu_ctx->fd = err;
     }
 
+    for (i = 0; i< vfu_ctx->nr_regions; i++) {
+        vfu_ctx->reg_info[i].fd = -1;
+    }
+
     return vfu_ctx;
 
 err_out:
@@ -1235,6 +1240,38 @@ copyin_mmap_areas(vfu_reg_info_t *reg_info,
     return 0;
 }
 
+static bool
+ranges_intersect(size_t off1, size_t size1, size_t off2, size_t size2)
+{
+    /*
+     * For two ranges to intersect, the start of each range must be before the
+     * end of the other range.
+     * TODO already defined in lib/pci_caps.c, maybe introduce a file for misc
+     * utility functions?
+     */
+    return (off1 < (off2 + size2) && off2 < (off1 + size1));
+}
+
+static bool
+maps_over_migr_regs(struct iovec *iov)
+{
+    return ranges_intersect(0, vfu_get_migr_register_area_size(),
+                            (size_t)iov->iov_base, iov->iov_len);
+}
+
+static bool
+validate_sparse_mmaps_for_migr_reg(vfu_reg_info_t *reg)
+{
+    int i;
+
+    for (i = 0; i < reg->nr_mmap_areas; i++) {
+        if (maps_over_migr_regs(&reg->mmap_areas[i])) {
+            return false;
+        }
+    }
+    return true;
+}
+
 int
 vfu_setup_region(vfu_ctx_t *vfu_ctx, int region_idx, size_t size,
                  vfu_region_access_cb_t *cb, int flags,
@@ -1243,7 +1280,7 @@ vfu_setup_region(vfu_ctx_t *vfu_ctx, int region_idx, size_t size,
     struct iovec whole_region = { .iov_base = 0, .iov_len = size };
     vfu_reg_info_t *reg;
     size_t i;
-    int ret;
+    int ret = 0;
 
     assert(vfu_ctx != NULL);
 
@@ -1254,7 +1291,7 @@ vfu_setup_region(vfu_ctx_t *vfu_ctx, int region_idx, size_t size,
     }
 
     if (region_idx < VFU_PCI_DEV_BAR0_REGION_IDX ||
-        region_idx > VFU_PCI_DEV_VGA_REGION_IDX) {
+        region_idx >= VFU_PCI_DEV_NUM_REGIONS) {
         vfu_log(vfu_ctx, LOG_ERR, "invalid region index %d", region_idx);
         return ERROR(EINVAL);
     }
@@ -1267,6 +1304,12 @@ vfu_setup_region(vfu_ctx_t *vfu_ctx, int region_idx, size_t size,
         return ERROR(EINVAL);
     }
 
+    if (region_idx == VFU_PCI_DEV_MIGR_REGION_IDX &&
+        size < vfu_get_migr_register_area_size()) {
+        vfu_log(vfu_ctx, LOG_ERR, "invalid migration region size %d", size);
+        return ERROR(EINVAL);
+    }
+
     for (i = 0; i < nr_mmap_areas; i++) {
         struct iovec *iov = &mmap_areas[i];
         if ((size_t)iov->iov_base + iov->iov_len > size) {
@@ -1289,11 +1332,30 @@ vfu_setup_region(vfu_ctx_t *vfu_ctx, int region_idx, size_t size,
     if (nr_mmap_areas > 0) {
         ret = copyin_mmap_areas(reg, mmap_areas, nr_mmap_areas);
         if (ret < 0) {
-            memset(reg, 0, sizeof (*reg));
-            return ERROR(-ret);
+            goto out;
         }
     }
 
+    if (region_idx == VFU_PCI_DEV_MIGR_REGION_IDX) {
+        if (!validate_sparse_mmaps_for_migr_reg(reg)) {
+            vfu_log(vfu_ctx, LOG_ERR,
+                    "migration registers cannot be memory mapped");
+            ret = -EINVAL;
+            goto out;
+        }
+
+        /*
+         * FIXME keeping for now until we're sure we're OK with fixing the
+         * migration region index.
+         */
+        vfu_ctx->migr_reg = reg;
+    }
+out:
+    if (ret < 0) {
+        free(reg->mmap_areas);
+        memset(reg, 0, sizeof (*reg));
+        return ERROR(-ret);
+    }
     return 0;
 }
 
@@ -1346,40 +1408,31 @@ vfu_setup_device_nr_irqs(vfu_ctx_t *vfu_ctx, enum vfu_dev_irq_type type,
 }
 
 int
-vfu_setup_device_migration(vfu_ctx_t *vfu_ctx, vfu_migration_t *migration)
+vfu_setup_device_migration_callbacks(vfu_ctx_t *vfu_ctx,
+                                     const vfu_migration_callbacks_t *callbacks,
+                                     uint64_t data_offset)
 {
-    vfu_reg_info_t *migr_reg;
     int ret = 0;
 
     assert(vfu_ctx != NULL);
+    assert(callbacks != NULL);
 
-    //FIXME: Validate args.
-
-    if (vfu_ctx->migr_reg != NULL) {
-        vfu_log(vfu_ctx, LOG_ERR, "device migration is already setup");
-        return ERROR(EEXIST);
+    if (vfu_ctx->migr_reg == NULL) {
+        vfu_log(vfu_ctx, LOG_ERR, "no device migration region");
+        return ERROR(EINVAL);
     }
 
-    /* FIXME hacky, find a more robust way to allocate a region index */
-    migr_reg = &vfu_ctx->reg_info[(vfu_ctx->nr_regions - 1)];
-
-    /* FIXME: Are there sparse areas need to be setup flags accordingly */
-    ret = copyin_mmap_areas(migr_reg, migration->mmap_areas,
-                            migration->nr_mmap_areas);
-    if (ret < 0) {
-        return ERROR(-ret);
+    if (callbacks->version != VFU_MIGR_CALLBACKS_VERS) {
+        vfu_log(vfu_ctx, LOG_ERR, "unsupported migration callbacks version %d",
+                callbacks->version);
+        return ERROR(EINVAL);
     }
 
-    migr_reg->flags = VFU_REGION_FLAG_RW;
-    migr_reg->size = sizeof(struct vfio_device_migration_info) + migration->size;
-
-    vfu_ctx->migration = init_migration(migration, &ret);
+    vfu_ctx->migration = init_migration(callbacks, data_offset, &ret);
     if (vfu_ctx->migration == NULL) {
         vfu_log(vfu_ctx, LOG_ERR, "failed to initialize device migration");
-        free(migr_reg->mmap_areas);
         return ERROR(ret);
     }
-    vfu_ctx->migr_reg = migr_reg;
 
     return 0;
 }
diff --git a/lib/migration.c b/lib/migration.c
index 2f532e7..0e1b2a6 100644
--- a/lib/migration.c
+++ b/lib/migration.c
@@ -51,9 +51,14 @@ enum migr_iter_state {
 };
 
 struct migration {
+    /*
+     * TODO if the user provides an FD then should mmap it and use the migration
+     * registers in the file
+     */
     struct vfio_device_migration_info info;
     size_t pgsize;
     vfu_migration_callbacks_t callbacks;
+    uint64_t data_offset;
 
     /*
      * This is only for the saving state. The resuming state is simpler so we
@@ -120,13 +125,24 @@ vfio_migr_state_transition_is_valid(__u32 from, __u32 to)
     return migr_states[from].state & (1 << to);
 }
 
+size_t
+vfu_get_migr_register_area_size(void)
+{
+    return ROUND_UP(sizeof(struct vfio_device_migration_info),
+                    sysconf(_SC_PAGE_SIZE));
+}
+
+/*
+ * TODO no need to dynamically allocate memory, we can keep struct migration
+ * in vfu_ctx_t.
+ */
 struct migration *
-init_migration(const vfu_migration_t * const vfu_migr, int *err)
+init_migration(const vfu_migration_callbacks_t * callbacks,
+               uint64_t data_offset, int *err)
 {
     struct migration *migr;
 
-    *err = 0;
-    if (vfu_migr->size < sizeof(struct vfio_device_migration_info)) {
+    if (data_offset < vfu_get_migr_register_area_size()) {
         *err = EINVAL;
         return NULL;
     }
@@ -140,14 +156,15 @@ init_migration(const vfu_migration_t * const vfu_migr, int *err)
     /*
      * FIXME: incorrect, if the client doesn't give a pgsize value, it means "no
      * migration support", handle this
+     * FIXME must be available even if migration callbacks aren't used
      */
     migr->pgsize = sysconf(_SC_PAGESIZE);
 
-
-    /* FIXME this should be done in vfu_ctx_run or poll */
+    /* FIXME this should be done in vfu_ctx_realize */
     migr->info.device_state = VFIO_DEVICE_STATE_RUNNING;
+    migr->data_offset = data_offset;
 
-    migr->callbacks = vfu_migr->callbacks;
+    migr->callbacks = *callbacks;
     if (migr->callbacks.transition == NULL ||
         migr->callbacks.get_pending_bytes == NULL ||
         migr->callbacks.prepare_data == NULL ||
@@ -345,7 +362,7 @@ handle_data_offset(vfu_ctx_t *vfu_ctx, struct migration *migr,
     case VFIO_DEVICE_STATE_RUNNING | VFIO_DEVICE_STATE_SAVING:
         ret = handle_data_offset_when_saving(vfu_ctx, migr, is_write);
         if (ret == 0 && !is_write) {
-            *offset = migr->iter.offset + sizeof(struct vfio_device_migration_info);
+            *offset = migr->iter.offset + migr->data_offset;
         }
         return ret;
     case VFIO_DEVICE_STATE_RESUMING:
@@ -358,7 +375,7 @@ handle_data_offset(vfu_ctx_t *vfu_ctx, struct migration *migr,
         if (ret < 0) {
             return ret;
         }
-        *offset += sizeof(struct vfio_device_migration_info);
+        *offset += migr->data_offset;
         return 0;
     }
     /* TODO improve error message */
@@ -496,7 +513,19 @@ migration_region_access(vfu_ctx_t *vfu_ctx, char *buf, size_t count,
         ret = migration_region_access_registers(vfu_ctx, buf, count,
                                                 pos, is_write);
     } else {
-        pos -= sizeof(struct vfio_device_migration_info);
+
+        if (pos < (loff_t)migr->data_offset) {
+            /*
+             * TODO we can simply ignore the access to that part and handle
+             * any access to the data region properly.
+             */
+            vfu_log(vfu_ctx, LOG_WARNING,
+                    "bad access to dead space %#lx-%#lx in migration region",
+                    pos, pos + count - 1);
+            return -EINVAL;
+        }
+
+        pos -= migr->data_offset;
         if (is_write) {
             ret = migr->callbacks.write_data(vfu_ctx, buf, count, pos);
         } else {
diff --git a/lib/migration.h b/lib/migration.h
index 0265d9b..bc73940 100644
--- a/lib/migration.h
+++ b/lib/migration.h
@@ -44,7 +44,8 @@
 #include "libvfio-user.h"
 
 struct migration *
-init_migration(const vfu_migration_t * const vfu_migr, int *err);
+init_migration(const vfu_migration_callbacks_t *callbacks,
+               uint64_t data_offset, int *err);
 
 ssize_t
 migration_region_access(vfu_ctx_t *vfu_ctx, char *buf, size_t count,
diff --git a/samples/client.c b/samples/client.c
index 92d86e0..d9b6661 100644
--- a/samples/client.c
+++ b/samples/client.c
@@ -209,7 +209,7 @@ send_device_reset(int sock)
 /* returns whether a VFIO migration capability is found */
 static bool
 get_region_vfio_caps(struct vfio_info_cap_header *header,
-                     struct vfio_region_info_cap_sparse_mmap *sparse)
+                     struct vfio_region_info_cap_sparse_mmap **sparse)
 {
     struct vfio_region_info_cap_type *type;
     unsigned int i;
@@ -218,12 +218,13 @@ get_region_vfio_caps(struct vfio_info_cap_header *header,
     while (true) {
         switch (header->id) {
             case VFIO_REGION_INFO_CAP_SPARSE_MMAP:
-                sparse = (struct vfio_region_info_cap_sparse_mmap*)header;
+                *sparse = (struct vfio_region_info_cap_sparse_mmap *)header;
                 printf("%s: Sparse cap nr_mmap_areas %d\n", __func__,
-                       sparse->nr_areas);
-                for (i = 0; i < sparse->nr_areas; i++) {
+                       (*sparse)->nr_areas);
+                for (i = 0; i < (*sparse)->nr_areas; i++) {
                     printf("%s: area %d offset %#llx size %llu\n", __func__,
-                           i, sparse->areas[i].offset, sparse->areas[i].size);
+                           i, (*sparse)->areas[i].offset,
+                           (*sparse)->areas[i].size);
                 }
                 break;
             case VFIO_REGION_INFO_CAP_TYPE:
@@ -261,14 +262,10 @@ do_get_device_region_info(int sock, struct vfio_region_info *region_info,
 }
 
 static void
-mmap_sparse_areas(int *fds, size_t nr_fds,
-                  struct vfio_region_info_cap_sparse_mmap *sparse)
+mmap_sparse_areas(int *fds, struct vfio_region_info_cap_sparse_mmap *sparse)
 {
     size_t i;
 
-    assert(nr_fds == 2);
-    assert(sparse->nr_areas == 2);
-
     for (i = 0; i < sparse->nr_areas; i++) {
 
         ssize_t ret;
@@ -318,9 +315,6 @@ get_device_region_info(int sock, uint32_t index)
         region_info->index = index;
         do_get_device_region_info(sock, region_info, fds, &nr_fds);
         assert(region_info->argsz == size);
-        assert(nr_fds == 2);
-        assert(fds[0] >= 0);
-        assert(fds[1] >= 0);
     } else {
         nr_fds = 0;
     }
@@ -332,9 +326,12 @@ get_device_region_info(int sock, uint32_t index)
     if (cap_sz) {
         struct vfio_region_info_cap_sparse_mmap *sparse = NULL;
         if (get_region_vfio_caps((struct vfio_info_cap_header*)(region_info + 1),
-                                 sparse)) {
+                                 &sparse)) {
             if (sparse != NULL) {
-                mmap_sparse_areas(fds, nr_fds, sparse);
+                assert((index == VFU_PCI_DEV_BAR1_REGION_IDX && nr_fds == 2) ||
+                       (index == VFU_PCI_DEV_MIGR_REGION_IDX && nr_fds == 1));
+                assert(nr_fds == sparse->nr_areas);
+                mmap_sparse_areas(fds, sparse);
             }
             return true;
         }
diff --git a/samples/server.c b/samples/server.c
index 6170753..9c2ee37 100644
--- a/samples/server.c
+++ b/samples/server.c
@@ -398,20 +398,41 @@ migration_data_written(UNUSED vfu_ctx_t *vfu_ctx, UNUSED __u64 count)
     return 0;
 }
 
+size_t
+nr_pages(size_t size)
+{
+    return (size / sysconf(_SC_PAGE_SIZE) + (size % sysconf(_SC_PAGE_SIZE) > 1));
+}
+
+size_t
+page_align(size_t size) {
+    return  nr_pages(size) * sysconf(_SC_PAGE_SIZE);
+}
+
 int main(int argc, char *argv[])
 {
     int ret;
     bool verbose = false;
     char opt;
     struct sigaction act = {.sa_handler = _sa_handler};
-    size_t bar1_size = 0x3000;
+    const size_t bar1_size = 0x3000;
+    size_t migr_regs_size, migr_data_size, migr_size;
     struct server_data server_data = {
         .migration = {
             .state = VFU_MIGR_STATE_RUNNING
         }
     };
     vfu_ctx_t *vfu_ctx;
-    FILE *fp;
+    FILE *bar1_fp, *migr_fp;
+    const vfu_migration_callbacks_t migr_callbacks = {
+        .version = VFU_MIGR_CALLBACKS_VERS,
+        .transition = &migration_device_state_transition,
+        .get_pending_bytes = &migration_get_pending_bytes,
+        .prepare_data = &migration_prepare_data,
+        .read_data = &migration_read_data,
+        .data_written = &migration_data_written,
+        .write_data = &migration_write_data
+    };
 
     while ((opt = getopt(argc, argv, "v")) != -1) {
         switch (opt) {
@@ -463,25 +484,26 @@ int main(int argc, char *argv[])
      * this under Linux. If we really want to prohibit it we have to use
      * separate files for the same region.
      */
-    if ((fp = tmpfile()) == NULL) {
+    if ((bar1_fp = tmpfile()) == NULL) {
         err(EXIT_FAILURE, "failed to create BAR1 file");
     }
     server_data.bar1_size = bar1_size;
-    if (ftruncate(fileno(fp), server_data.bar1_size) == -1) {
+    if (ftruncate(fileno(bar1_fp), server_data.bar1_size) == -1) {
         err(EXIT_FAILURE, "failed to truncate BAR1 file");
     }
     server_data.bar1 = mmap(NULL, server_data.bar1_size, PROT_READ | PROT_WRITE,
-                            MAP_SHARED, fileno(fp), 0);
+                            MAP_SHARED, fileno(bar1_fp), 0);
     if (server_data.bar1 == MAP_FAILED) {
         err(EXIT_FAILURE, "failed to mmap BAR1");
     }
-    struct iovec mmap_areas[] = {
+    struct iovec bar1_mmap_areas[] = {
         { .iov_base  = (void*)0, .iov_len = 0x1000 },
         { .iov_base  = (void*)0x2000, .iov_len = 0x1000 }
     };
     ret = vfu_setup_region(vfu_ctx, VFU_PCI_DEV_BAR1_REGION_IDX,
                            server_data.bar1_size, &bar1_access,
-                           VFU_REGION_FLAG_RW, mmap_areas, 2, fileno(fp));
+                           VFU_REGION_FLAG_RW, bar1_mmap_areas, 2,
+                           fileno(bar1_fp));
     if (ret < 0) {
         err(EXIT_FAILURE, "failed to setup BAR1 region");
     }
@@ -501,21 +523,41 @@ int main(int argc, char *argv[])
         err(EXIT_FAILURE, "failed to setup irq counts");
     }
 
-    vfu_migration_t migration = {
-        .size = bar1_size + sizeof(time_t),
-        .mmap_areas = mmap_areas,
-        .nr_mmap_areas = 2,
-        .callbacks = {
-            .transition = &migration_device_state_transition,
-            .get_pending_bytes = &migration_get_pending_bytes,
-            .prepare_data = &migration_prepare_data,
-            .read_data = &migration_read_data,
-            .data_written = &migration_data_written,
-            .write_data = &migration_write_data
-        }
-    };
+    /* setup migration */
 
-    ret = vfu_setup_device_migration(vfu_ctx, &migration);
+    /*
+     * The migration registers aren't memory mappable, so in order to make the
+     * rest of the migration region memory mappable we must effectively reserve
+     * an entire page.
+     */
+    migr_regs_size = vfu_get_migr_register_area_size();
+    migr_data_size = page_align(bar1_size + sizeof(time_t));
+    migr_size = migr_regs_size + migr_data_size;
+    if ((migr_fp = tmpfile()) == NULL) {
+        err(EXIT_FAILURE, "failed to create migration file");
+    }
+    if (ftruncate(fileno(migr_fp), migr_size) == -1) {
+        err(EXIT_FAILURE, "failed to truncate migration file");
+    }
+    server_data.bar1 = mmap(NULL, server_data.bar1_size, PROT_READ | PROT_WRITE,
+                            MAP_SHARED, fileno(bar1_fp), 0);
+    if (server_data.bar1 == MAP_FAILED) {
+        err(EXIT_FAILURE, "failed to mmap migration file");
+    }
+    struct iovec migr_mmap_areas[] = {
+        [0] = {
+            .iov_base  = (void *)migr_regs_size,
+            .iov_len = migr_data_size
+        },
+    };
+    ret = vfu_setup_region(vfu_ctx, VFU_PCI_DEV_MIGR_REGION_IDX, migr_size,
+                           NULL, VFU_REGION_FLAG_RW, migr_mmap_areas,
+                           ARRAY_SIZE(migr_mmap_areas), fileno(migr_fp));
+    if (ret < 0) {
+        err(EXIT_FAILURE, "failed to setup migration region");
+    }
+    ret = vfu_setup_device_migration_callbacks(vfu_ctx, &migr_callbacks,
+                                               migr_regs_size);
     if (ret < 0) {
         err(EXIT_FAILURE, "failed to setup device migration");
     }
diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt
index 2a339f0..a4f49ad 100644
--- a/test/CMakeLists.txt
+++ b/test/CMakeLists.txt
@@ -53,6 +53,8 @@ target_link_libraries(unit-tests PUBLIC "-Wl,--wrap=close")
 target_link_libraries(unit-tests PUBLIC "-Wl,--wrap=tran_sock_send_iovec")
 target_link_libraries(unit-tests PUBLIC "-Wl,--wrap=free")
 target_link_libraries(unit-tests PUBLIC "-Wl,--wrap=process_request")
+target_link_libraries(unit-tests PUBLIC "-Wl,--wrap=bind")
+target_link_libraries(unit-tests PUBLIC "-Wl,--wrap=listen")
 
 enable_testing()
 add_test(NAME unit-tests COMMAND unit-tests)
diff --git a/test/mocks.c b/test/mocks.c
index c38ec7b..303a6c2 100644
--- a/test/mocks.c
+++ b/test/mocks.c
@@ -33,6 +33,8 @@
 #include <setjmp.h>
 #include <cmocka.h>
 #include <stdio.h>
+#include <sys/types.h>
+#include <sys/socket.h>
 
 #include "mocks.h"
 #include "dma.h"
@@ -124,6 +126,10 @@ __wrap_exec_command(vfu_ctx_t *vfu_ctx, struct vfio_user_header *hdr,
 int
 __wrap_close(int fd)
 {
+    if (!is_patched(close)) {
+        return __real_close(fd);
+    }
+
     check_expected(fd);
     return mock();
 }
@@ -168,6 +174,19 @@ __wrap_process_request(vfu_ctx_t *vfu_ctx)
     return mock();
 }
 
+int __wrap_bind(int sockfd __attribute__((unused)),
+    const struct sockaddr *addr __attribute__((unused)),
+    socklen_t addrlen __attribute__((unused)))
+{
+    return 0;
+}
+
+int __wrap_listen(int sockfd __attribute__((unused)),
+    int backlog __attribute__((unused)))
+{
+    return 0;
+}
+
 /* FIXME should be something faster than unsorted array, look at tsearch(3). */
 static struct function funcs[] = {
     {.addr = &__wrap_dma_controller_add_region},
@@ -179,7 +198,9 @@ static struct function funcs[] = {
     {.addr = &__wrap_close},
     {.addr = &__wrap_tran_sock_send_iovec},
     {.addr = &__wrap_free},
-    {.addr = &__wrap_process_request}
+    {.addr = &__wrap_process_request},
+    {.addr = &__wrap_bind},
+    {.addr = &__wrap_listen}
 };
 
 static struct function*
diff --git a/test/unit-tests.c b/test/unit-tests.c
index c989d91..e3ad86e 100644
--- a/test/unit-tests.c
+++ b/test/unit-tests.c
@@ -39,6 +39,7 @@
 #include <alloca.h>
 #include <string.h>
 #include <linux/pci_regs.h>
+#include <sys/param.h>
 
 #include "dma.h"
 #include "libvfio-user.h"
@@ -1202,11 +1203,158 @@ test_migration_state_transitions(void **state __attribute__ ((unused)))
  * and provide a function to execute before and after each unit test.
  */
 static int
-setup(void **state __attribute__((unused))) {
+setup(void **state __attribute__((unused)))
+{
     unpatch_all();
     return 0;
 }
 
+static struct test_setup_migr_reg_dat {
+    vfu_ctx_t *v;
+    size_t rs; /* migration registers size */
+    size_t ds; /* migration data size */
+    size_t s; /* migration region size*/
+    const vfu_migration_callbacks_t c;
+} migr_reg_data = {
+    .c = {
+        .version = VFU_MIGR_CALLBACKS_VERS,
+        .transition = (void *)0x1,
+        .get_pending_bytes = (void *)0x2,
+        .prepare_data = (void *)0x3,
+        .read_data = (void *)0x4,
+        .write_data = (void *)0x5,
+        .data_written = (void *)0x6
+    }
+};
+
+static int
+setup_test_setup_migration_region(void **state)
+{
+    struct test_setup_migr_reg_dat *p = &migr_reg_data;
+    p->v = vfu_create_ctx(VFU_TRANS_SOCK, "test", 0, NULL,
+        VFU_DEV_TYPE_PCI);
+    if (p->v == NULL) {
+        return -1;
+    }
+    p->rs = ROUND_UP(sizeof(struct vfio_device_migration_info), sysconf(_SC_PAGE_SIZE));
+    p->ds = sysconf(_SC_PAGE_SIZE);
+    p->s = p->rs + p->ds;
+    *state = p;
+    return setup(state);
+}
+
+static vfu_ctx_t *
+get_vfu_ctx(void **state)
+{
+    return (*((struct test_setup_migr_reg_dat **)(state)))->v;
+}
+
+static int
+teardown_test_setup_migration_region(void **state)
+{
+    struct test_setup_migr_reg_dat *p = *state;
+    vfu_destroy_ctx(p->v);
+    return 0;
+}
+
+static void
+test_setup_migration_region_too_small(void **state)
+{
+    vfu_ctx_t *v = get_vfu_ctx(state);
+    int r = vfu_setup_region(v, VFU_PCI_DEV_MIGR_REGION_IDX,
+        vfu_get_migr_register_area_size() - 1, NULL,
+        VFU_REGION_FLAG_READ | VFU_REGION_FLAG_WRITE, NULL, 0, -1);
+    assert_int_equal(-1, r);
+    assert_int_equal(EINVAL, errno);
+}
+
+static void
+test_setup_migration_region_size_ok(void **state)
+{
+    vfu_ctx_t *v = get_vfu_ctx(state);
+    int r = vfu_setup_region(v, VFU_PCI_DEV_MIGR_REGION_IDX,
+        vfu_get_migr_register_area_size(), NULL,
+        VFU_REGION_FLAG_READ | VFU_REGION_FLAG_WRITE, NULL, 0, -1);
+    assert_int_equal(0, r);
+}
+
+static void
+test_setup_migration_region_fully_mappable(void **state)
+{
+    struct test_setup_migr_reg_dat *p = *state;
+    int r = vfu_setup_region(p->v, VFU_PCI_DEV_MIGR_REGION_IDX, p->s,
+        NULL, VFU_REGION_FLAG_READ | VFU_REGION_FLAG_WRITE, NULL, 0,
+        0xdeadbeef);
+    assert_int_equal(-1, r);
+    assert_int_equal(EINVAL, errno);
+}
+
+static void
+test_setup_migration_region_sparsely_mappable_over_migration_registers(void **state)
+{
+    struct test_setup_migr_reg_dat *p = *state;
+    struct iovec mmap_areas[] = {
+        [0] = {
+            .iov_base = 0,
+            .iov_len = p->rs
+        }
+    };
+    int r = vfu_setup_region(p->v, VFU_PCI_DEV_MIGR_REGION_IDX, p->s, NULL,
+        VFU_REGION_FLAG_READ | VFU_REGION_FLAG_WRITE, mmap_areas, 1, 0xdeadbeef);
+    assert_int_equal(-1, r);
+    assert_int_equal(EINVAL, errno);
+}
+
+static void
+test_setup_migration_region_sparsely_mappable_valid(void **state)
+{
+    struct test_setup_migr_reg_dat *p = *state;
+    struct iovec mmap_areas[] = {
+        [0] = {
+            .iov_base = (void *)p->rs,
+            .iov_len = p->ds
+        }
+    };
+    int r = vfu_setup_region(p->v, VFU_PCI_DEV_MIGR_REGION_IDX, p->s, NULL,
+        VFU_REGION_FLAG_READ | VFU_REGION_FLAG_WRITE, mmap_areas, 1,
+        0xdeadbeef);
+    assert_int_equal(0, r);
+}
+
+static void
+test_setup_migration_callbacks_without_migration_region(void **state)
+{
+    struct test_setup_migr_reg_dat *p = *state;
+    assert_int_equal(-1, vfu_setup_device_migration_callbacks(p->v, &p->c, 0));
+    assert_int_equal(EINVAL, errno);
+}
+
+static void
+test_setup_migration_callbacks_bad_data_offset(void **state)
+{
+    struct test_setup_migr_reg_dat *p = *state;
+    int r = vfu_setup_region(p->v, VFU_PCI_DEV_MIGR_REGION_IDX, p->s, NULL,
+        VFU_REGION_FLAG_READ | VFU_REGION_FLAG_WRITE, NULL, 0, -1);
+    assert_int_equal(0, r);
+    r = vfu_setup_device_migration_callbacks(p->v, &p->c,
+        vfu_get_migr_register_area_size() - 1);
+    assert_int_equal(-1, r);
+}
+
+static void
+test_setup_migration_callbacks(void **state)
+{
+    struct test_setup_migr_reg_dat *p = *state;
+    int r = vfu_setup_region(p->v, VFU_PCI_DEV_MIGR_REGION_IDX, p->s, NULL,
+        VFU_REGION_FLAG_READ | VFU_REGION_FLAG_WRITE, NULL, 0, -1);
+    assert_int_equal(0, r);
+    r = vfu_setup_device_migration_callbacks(p->v, &p->c,
+        vfu_get_migr_register_area_size());
+    assert_int_equal(0, r);
+    assert_non_null(p->v->migration);
+    /* FIXME can't validate p->v->migration because it's a private strcut, need to move it out of lib/migration.c */
+}
+
 int main(void)
 {
    const struct CMUnitTest tests[] = {
@@ -1232,7 +1380,31 @@ int main(void)
         cmocka_unit_test_setup(test_dma_map_sg, setup),
         cmocka_unit_test_setup(test_dma_addr_to_sg, setup),
         cmocka_unit_test_setup(test_vfu_setup_device_dma_cb, setup),
-        cmocka_unit_test_setup(test_migration_state_transitions, setup)
+        cmocka_unit_test_setup(test_migration_state_transitions, setup),
+        cmocka_unit_test_setup_teardown(test_setup_migration_region_too_small,
+            setup_test_setup_migration_region,
+            teardown_test_setup_migration_region),
+        cmocka_unit_test_setup_teardown(test_setup_migration_region_size_ok,
+            setup_test_setup_migration_region,
+            teardown_test_setup_migration_region),
+        cmocka_unit_test_setup_teardown(test_setup_migration_region_fully_mappable,
+            setup_test_setup_migration_region,
+            teardown_test_setup_migration_region),
+        cmocka_unit_test_setup_teardown(test_setup_migration_region_sparsely_mappable_over_migration_registers,
+            setup_test_setup_migration_region,
+            teardown_test_setup_migration_region),
+        cmocka_unit_test_setup_teardown(test_setup_migration_region_sparsely_mappable_valid,
+            setup_test_setup_migration_region,
+            teardown_test_setup_migration_region),
+        cmocka_unit_test_setup_teardown(test_setup_migration_callbacks_without_migration_region,
+            setup_test_setup_migration_region,
+            teardown_test_setup_migration_region),
+        cmocka_unit_test_setup_teardown(test_setup_migration_callbacks_bad_data_offset,
+            setup_test_setup_migration_region,
+            teardown_test_setup_migration_region),
+        cmocka_unit_test_setup_teardown(test_setup_migration_callbacks,
+            setup_test_setup_migration_region,
+            teardown_test_setup_migration_region),
     };
 
     return cmocka_run_group_tests(tests, NULL, NULL);