refactor message handling path (#376)

Capture message handling inside a new vfu_msg_t private structure and pass that
around to the handlers. This provides no functional change, but greatly
simplifies and cleans up that path, especially around fd and iovec handling.

As part of fixing up the unit tests, start using global variables to reduce the
amount of boiler-plate.

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

8 files changed, 1142 insertions, 1187 deletions

diff --git a/lib/irq.c b/lib/irq.c
index a274191..fde17dc 100644
--- a/lib/irq.c
+++ b/lib/irq.c
@@ -54,43 +54,41 @@ vfio_irq_idx_to_str(int index)
     }
 }
 
-static int
-dev_get_irqinfo(vfu_ctx_t *vfu_ctx, struct vfio_irq_info *irq_info_in,
-                struct vfio_irq_info *irq_info_out)
+int
+handle_device_get_irq_info(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg)
 {
+    struct vfio_irq_info *in_info;
+    struct vfio_irq_info *out_info;
+
     assert(vfu_ctx != NULL);
-    assert(irq_info_in != NULL);
-    assert(irq_info_out != NULL);
-
-    // Ensure provided argsz is sufficiently big and index is within bounds.
-    if ((irq_info_in->argsz < sizeof(struct vfio_irq_info)) ||
-        (irq_info_in->index >= VFU_DEV_NUM_IRQS)) {
-        vfu_log(vfu_ctx, LOG_DEBUG, "bad irq_info (size=%d index=%d)",
-                irq_info_in->argsz, irq_info_in->index);
+    assert(msg != NULL);
+
+    in_info = msg->in_data;
+
+    // FIXME: don't do the ->argsz/in_size check elsewhere
+    if (msg->in_size < sizeof(*in_info) || msg->in_size != in_info->argsz) {
         return ERROR_INT(EINVAL);
     }
 
-    irq_info_out->count = vfu_ctx->irq_count[irq_info_in->index];
-    irq_info_out->flags = VFIO_IRQ_INFO_EVENTFD;
+    if (in_info->index >= VFU_DEV_NUM_IRQS) {
+        vfu_log(vfu_ctx, LOG_DEBUG, "bad irq_info index %d\n", in_info->index);
+        return ERROR_INT(EINVAL);
+    }
 
-    return 0;
-}
+    msg->out_size = sizeof (*out_info);
+    msg->out_data = calloc(1, sizeof(*out_info));
 
-int
-handle_device_get_irq_info(vfu_ctx_t *vfu_ctx, uint32_t size,
-                           struct vfio_irq_info *irq_info_in,
-                           struct vfio_irq_info *irq_info_out)
-{
-    assert(vfu_ctx != NULL);
-    assert(irq_info_in != NULL);
-    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 ERROR_INT(EINVAL);
+    if (msg->out_data == NULL) {
+        return -1;
     }
 
-    return dev_get_irqinfo(vfu_ctx, irq_info_in, irq_info_out);
+    out_info = msg->out_data;
+    out_info->argsz = sizeof (*out_info);
+    out_info->flags = VFIO_IRQ_INFO_EVENTFD;
+    out_info->index = in_info->index;
+    out_info->count = vfu_ctx->irq_count[in_info->index];
+
+    return 0;
 }
 
 static void
@@ -191,6 +189,10 @@ irqs_set_data_bool(vfu_ctx_t *vfu_ctx, struct vfio_irq_set *irq_set, void *data)
     eventfd_t val;
 
     assert(data != NULL);
+
+    // FIXME: we don't check that data is actually within ->argsz, could
+    // dereference junk
+
     for (i = irq_set->start, d8 = data; i < (irq_set->start + irq_set->count);
          i++, d8++) {
         efd = vfu_ctx->irqs->efds[i];
@@ -318,21 +320,21 @@ invalid:
 }
 
 int
-handle_device_set_irqs(vfu_ctx_t *vfu_ctx, uint32_t size,
-                       int *fds, size_t nr_fds, struct vfio_irq_set *irq_set)
+handle_device_set_irqs(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg)
 {
+    struct vfio_irq_set *irq_set = msg->in_data;
     uint32_t data_type;
     int ret;
 
     assert(vfu_ctx != NULL);
-    assert(irq_set != NULL);
+    assert(msg != NULL);
 
-    if (size < sizeof(*irq_set) || size != irq_set->argsz) {
-        vfu_log(vfu_ctx, LOG_ERR, "%s: bad size %u", __func__, size);
+    if (msg->in_size < sizeof(*irq_set) || msg->in_size != irq_set->argsz) {
+        vfu_log(vfu_ctx, LOG_ERR, "bad size %zu", msg->in_size);
         return ERROR_INT(EINVAL);
     }
 
-    ret = device_set_irqs_validate(vfu_ctx, irq_set, nr_fds);
+    ret = device_set_irqs_validate(vfu_ctx, irq_set, msg->nr_in_fds);
     if (ret != 0) {
         return ret;
     }
@@ -350,7 +352,7 @@ handle_device_set_irqs(vfu_ctx_t *vfu_ctx, uint32_t size,
     data_type = irq_set->flags & VFIO_IRQ_SET_DATA_TYPE_MASK;
 
     if ((data_type == VFIO_IRQ_SET_DATA_NONE && irq_set->count == 0) ||
-        (data_type == VFIO_IRQ_SET_DATA_EVENTFD && nr_fds == 0)) {
+        (data_type == VFIO_IRQ_SET_DATA_EVENTFD && msg->nr_in_fds == 0)) {
         irqs_disable(vfu_ctx, irq_set->index, irq_set->start, irq_set->count);
         return 0;
     }
@@ -363,7 +365,7 @@ handle_device_set_irqs(vfu_ctx_t *vfu_ctx, uint32_t size,
     case VFIO_IRQ_SET_DATA_NONE:
         return irqs_set_data_none(vfu_ctx, irq_set);
     case VFIO_IRQ_SET_DATA_EVENTFD:
-        return irqs_set_data_eventfd(vfu_ctx, irq_set, fds);
+        return irqs_set_data_eventfd(vfu_ctx, irq_set, msg->in_fds);
     case VFIO_IRQ_SET_DATA_BOOL:
         return irqs_set_data_bool(vfu_ctx, irq_set, irq_set + 1);
         break;
diff --git a/lib/irq.h b/lib/irq.h
index c230023..e4f430b 100644
--- a/lib/irq.h
+++ b/lib/irq.h
@@ -39,12 +39,10 @@ void
 irqs_reset(vfu_ctx_t *vfu_ctx);
 
 int
-handle_device_get_irq_info(vfu_ctx_t *vfu_ctx, uint32_t size,
-                           struct vfio_irq_info *irq_info_in,
-                           struct vfio_irq_info *irq_info_out);
+handle_device_get_irq_info(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg);
+
 int
-handle_device_set_irqs(vfu_ctx_t *vfu_ctx, uint32_t size,
-                       int *fds, size_t nr_fds, struct vfio_irq_set *irq_set);
+handle_device_set_irqs(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg);
 
 #endif /* LIB_VFIO_USER_IRQ_H */
 
diff --git a/lib/libvfio-user.c b/lib/libvfio-user.c
index 7df7804..8a80dfc 100644
--- a/lib/libvfio-user.c
+++ b/lib/libvfio-user.c
@@ -265,6 +265,7 @@ is_valid_region_access(vfu_ctx_t *vfu_ctx, size_t size, uint16_t cmd,
         return false;
     }
 
+    // FIXME: ->count unbounded (alloc), except for region size
     // FIXME: need to audit later for wraparound
     if (ra->offset + ra->count > vfu_ctx->reg_info[index].size) {
         vfu_log(vfu_ctx, LOG_ERR, "out of bounds region access %#lx-%#lx "
@@ -286,46 +287,51 @@ is_valid_region_access(vfu_ctx_t *vfu_ctx, size_t size, uint16_t cmd,
 }
 
 static int
-handle_region_access(vfu_ctx_t *vfu_ctx, uint32_t size, uint16_t cmd,
-                     void **data, size_t *len,
-                     struct vfio_user_region_access *ra)
+handle_region_access(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg)
 {
+    struct vfio_user_region_access *in_ra = msg->in_data;
+    struct vfio_user_region_access *out_ra;
     ssize_t ret;
     char *buf;
 
     assert(vfu_ctx != NULL);
-    assert(data != NULL);
-    assert(ra != NULL);
+    assert(msg != NULL);
 
-    if (!is_valid_region_access(vfu_ctx, size, cmd, ra)) {
+    if (!is_valid_region_access(vfu_ctx, msg->in_size, msg->hdr.cmd, in_ra)) {
         return ERROR_INT(EINVAL);
     }
 
-    if (ra->count == 0) {
+    if (in_ra->count == 0) {
         return 0;
     }
 
-    *len = sizeof(*ra);
-    if (cmd == VFIO_USER_REGION_READ) {
-        *len += ra->count;
+    msg->out_size = sizeof(*in_ra);
+    if (msg->hdr.cmd == VFIO_USER_REGION_READ) {
+        msg->out_size += in_ra->count;
     }
-    *data = calloc(1, *len);
-    if (*data == NULL) {
+    msg->out_data = calloc(1, msg->out_size);
+    if (msg->out_data == NULL) {
         return -1;
     }
-    if (cmd == VFIO_USER_REGION_READ) {
-        buf = (char *)(((struct vfio_user_region_access*)(*data)) + 1);
+
+    out_ra = msg->out_data;
+    out_ra->region = in_ra->region;
+    out_ra->offset = in_ra->offset;
+    out_ra->count = in_ra->count;
+
+    if (msg->hdr.cmd == VFIO_USER_REGION_READ) {
+        buf = (char *)(&out_ra->data);
     } else {
-        buf = (char *)(ra + 1);
+        buf = (char *)(&in_ra->data);
     }
 
-    ret = region_access(vfu_ctx, ra->region, buf, ra->count, ra->offset,
-                        cmd == VFIO_USER_REGION_WRITE);
+    ret = region_access(vfu_ctx, in_ra->region, buf, in_ra->count,
+                        in_ra->offset, msg->hdr.cmd == VFIO_USER_REGION_WRITE);
 
-    if (ret != ra->count) {
+    if (ret != in_ra->count) {
         vfu_log(vfu_ctx, LOG_ERR, "failed to %s %#x-%#lx: %m",
-                cmd == VFIO_USER_REGION_WRITE ? "write" : "read",
-                ra->count, ra->offset + ra->count - 1);
+                msg->hdr.cmd == VFIO_USER_REGION_WRITE ? "write" : "read",
+                in_ra->count, in_ra->offset + in_ra->count - 1);
         /* FIXME we should return whatever has been accessed, not an error */
         if (ret >= 0) {
             ret = ERROR_INT(EINVAL);
@@ -333,8 +339,42 @@ handle_region_access(vfu_ctx_t *vfu_ctx, uint32_t size, uint16_t cmd,
         return ret;
     }
 
-    ra = *data;
-    ra->count = ret;
+    out_ra->count = ret;
+
+    return 0;
+}
+
+int
+handle_device_get_info(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg)
+{
+    struct vfio_device_info *in_info;
+    struct vfio_device_info *out_info;
+
+    assert(vfu_ctx != NULL);
+    assert(msg != NULL);
+
+    in_info = msg->in_data;
+
+    if (msg->in_size < sizeof(*in_info) || in_info->argsz < sizeof(*in_info)) {
+        return ERROR_INT(EINVAL);
+    }
+
+    msg->out_size = sizeof (*out_info);
+    msg->out_data = calloc(1, sizeof(*out_info));
+
+    if (msg->out_data == NULL) {
+        return -1;
+    }
+
+    out_info = msg->out_data;
+    out_info->argsz = sizeof(*in_info);
+    out_info->flags = VFIO_DEVICE_FLAGS_PCI | VFIO_DEVICE_FLAGS_RESET;
+    out_info->num_regions = vfu_ctx->nr_regions;
+    out_info->num_irqs = VFU_DEV_NUM_IRQS;
+
+    vfu_log(vfu_ctx, LOG_DEBUG, "devinfo flags %#x, num_regions %d, "
+            "num_irqs %d", out_info->flags, out_info->num_regions,
+            out_info->num_irqs);
 
     return 0;
 }
@@ -348,109 +388,75 @@ region_to_offset(uint32_t region)
 }
 
 int
-dev_get_reginfo(vfu_ctx_t *vfu_ctx, uint32_t index, uint32_t argsz,
-                struct vfio_region_info **vfio_reg, int **fds, size_t *nr_fds)
+handle_device_get_region_info(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg)
 {
+    struct vfio_region_info *in_info;
+    struct vfio_region_info *out_info;
     vfu_reg_info_t *vfu_reg;
     size_t caps_size;
 
     assert(vfu_ctx != NULL);
-    assert(vfio_reg != NULL);
+    assert(msg != NULL);
 
-    vfu_reg = &vfu_ctx->reg_info[index];
+    in_info = msg->in_data;
 
-    if (index >= vfu_ctx->nr_regions) {
-        vfu_log(vfu_ctx, LOG_DEBUG, "bad region index %d in get region info",
-                index);
+    if (msg->in_size < sizeof(*in_info) || in_info->argsz < sizeof(*in_info)) {
         return ERROR_INT(EINVAL);
     }
 
-    if (argsz < sizeof(struct vfio_region_info)) {
-        vfu_log(vfu_ctx, LOG_DEBUG, "bad argsz %d", argsz);
+    if (in_info->index >= vfu_ctx->nr_regions) {
+        vfu_log(vfu_ctx, LOG_DEBUG, "bad region index %d in get region info",
+                in_info->index);
         return ERROR_INT(EINVAL);
     }
 
-    /*
-     * TODO We assume that the client expects to receive argsz bytes.
-     */
-    *vfio_reg = calloc(1, argsz);
-    if (!*vfio_reg) {
+    // FIXME: we don't cap client-provided in_info->argsz
+    msg->out_size = in_info->argsz;
+    msg->out_data = calloc(1, msg->out_size);
+
+    if (msg->out_data == NULL) {
         return -1;
     }
-    caps_size = get_vfio_caps_size(index == VFU_PCI_DEV_MIGR_REGION_IDX,
+
+    out_info = msg->out_data;
+
+    vfu_reg = &vfu_ctx->reg_info[in_info->index];
+
+    caps_size = get_vfio_caps_size(in_info->index == VFU_PCI_DEV_MIGR_REGION_IDX,
                                    vfu_reg);
-    (*vfio_reg)->argsz = caps_size + sizeof(struct vfio_region_info);
-    (*vfio_reg)->index = index;
-    (*vfio_reg)->offset = region_to_offset((*vfio_reg)->index);
-    (*vfio_reg)->size = vfu_reg->size;
 
-    (*vfio_reg)->flags = 0;
+    /* This might be more than the buffer we actually return. */
+    out_info->argsz = sizeof(*out_info) + caps_size;
+    out_info->index = in_info->index;
+    out_info->offset = region_to_offset(out_info->index);
+    out_info->size = vfu_reg->size;
+
+    out_info->flags = 0;
 
     if (vfu_reg->flags & VFU_REGION_FLAG_READ) {
-        (*vfio_reg)->flags |= VFIO_REGION_INFO_FLAG_READ;
+        out_info->flags |= VFIO_REGION_INFO_FLAG_READ;
     }
     if (vfu_reg->flags & VFU_REGION_FLAG_WRITE) {
-        (*vfio_reg)->flags |= VFIO_REGION_INFO_FLAG_WRITE;
+        out_info->flags |= VFIO_REGION_INFO_FLAG_WRITE;
     }
 
     if (vfu_reg->fd != -1) {
-        (*vfio_reg)->flags |= VFIO_REGION_INFO_FLAG_MMAP;
+        out_info->flags |= VFIO_REGION_INFO_FLAG_MMAP;
     }
 
-    *nr_fds = 0;
     if (caps_size > 0) {
-        (*vfio_reg)->flags |= VFIO_REGION_INFO_FLAG_CAPS;
-        if (argsz >= (*vfio_reg)->argsz) {
-            dev_get_caps(vfu_ctx, vfu_reg, index == VFU_PCI_DEV_MIGR_REGION_IDX,
-                         *vfio_reg, fds, nr_fds);
+        out_info->flags |= VFIO_REGION_INFO_FLAG_CAPS;
+        /* Only actually provide the caps if they fit. */
+        if (in_info->argsz >= out_info->argsz) {
+            dev_get_caps(vfu_ctx, vfu_reg,
+                         in_info->index == VFU_PCI_DEV_MIGR_REGION_IDX,
+                         out_info, &msg->out_fds, &msg->nr_out_fds);
         }
     }
 
-    vfu_log(vfu_ctx, LOG_DEBUG, "region_info[%d] offset %#llx flags %#x size %llu "
-            "argsz %u",
-            (*vfio_reg)->index, (*vfio_reg)->offset, (*vfio_reg)->flags,
-            (*vfio_reg)->size, (*vfio_reg)->argsz);
-
-    return 0;
-}
-
-/* TODO merge with dev_get_reginfo */
-static int
-handle_device_get_region_info(vfu_ctx_t *vfu_ctx, uint32_t size,
-                              struct vfio_region_info *reg_info_in,
-                              struct vfio_region_info **reg_info_out,
-                              int **fds, size_t *nr_fds)
-{
-    if (size < sizeof(*reg_info_in)) {
-        return ERROR_INT(EINVAL);
-    }
-
-    return dev_get_reginfo(vfu_ctx, reg_info_in->index, reg_info_in->argsz,
-                           reg_info_out, fds, nr_fds);
-}
-
-int
-handle_device_get_info(vfu_ctx_t *vfu_ctx, uint32_t in_size,
-                       struct vfio_device_info *in_dev_info,
-                       struct vfio_device_info *out_dev_info)
-{
-    assert(vfu_ctx != NULL);
-    assert(in_dev_info != NULL);
-    assert(out_dev_info != NULL);
-
-    if (in_size < sizeof(*in_dev_info) ||
-        in_dev_info->argsz < sizeof(*in_dev_info)) {
-        return ERROR_INT(EINVAL);
-    }
-
-    out_dev_info->argsz = sizeof(*in_dev_info);
-    out_dev_info->flags = VFIO_DEVICE_FLAGS_PCI | VFIO_DEVICE_FLAGS_RESET;
-    out_dev_info->num_regions = vfu_ctx->nr_regions;
-    out_dev_info->num_irqs = VFU_DEV_NUM_IRQS;
-
-    vfu_log(vfu_ctx, LOG_DEBUG, "devinfo flags %#x, num_regions %d, "
-            "num_irqs %d", out_dev_info->flags, out_dev_info->num_regions,
-            out_dev_info->num_irqs);
+    vfu_log(vfu_ctx, LOG_DEBUG, "region_info[%d] offset %#llx flags %#x "
+            "size %llu " "argsz %u", out_info->index, out_info->offset,
+            out_info->flags, out_info->size, out_info->argsz);
 
     return 0;
 }
@@ -469,40 +475,28 @@ consume_fd(int *fds, size_t nr_fds, size_t index)
    return fd;
 }
 
-/*
- * Handles a DMA map/unmap request.
- *
- * @vfu_ctx: LM context
- * @size: size, in bytes, of the memory pointed to be @dma_regions
- * @map: whether this is a DMA map operation
- * @fds: array of file descriptors.
- * @nr_fds: size of above array.
- * @dma_regions: memory that contains the DMA regions to be mapped/unmapped
- *
- * @returns 0 on success, -1 and errno on failure.
- */
 int
-handle_dma_map_or_unmap(vfu_ctx_t *vfu_ctx, uint32_t size, bool map,
-                        int *fds, size_t nr_fds,
-                        struct vfio_user_dma_region *dma_regions)
+handle_dma_map_or_unmap(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg)
 {
-    int nr_dma_regions;
-    int ret, i;
+    struct vfio_user_dma_region *dma_regions = msg->in_data;
+    bool is_map = (msg->hdr.cmd == VFIO_USER_DMA_MAP);
+    size_t nr_dma_regions;
     size_t fdi;
+    size_t i;
+    int ret;
 
     assert(vfu_ctx != NULL);
-    assert(nr_fds == 0 || fds != NULL);
 
-    if (vfu_ctx->dma == NULL) {
-        return 0;
+    if (msg->in_size % sizeof(struct vfio_user_dma_region) != 0) {
+        vfu_log(vfu_ctx, LOG_ERR, "bad size of DMA regions %zu", msg->in_size);
+        return ERROR_INT(EINVAL);
     }
 
-    if (size % sizeof(struct vfio_user_dma_region) != 0) {
-        vfu_log(vfu_ctx, LOG_ERR, "bad size of DMA regions %d", size);
-        return ERROR_INT(EINVAL);
+    if (vfu_ctx->dma == NULL) {
+        return 0;
     }
 
-    nr_dma_regions = (int)(size / sizeof(struct vfio_user_dma_region));
+    nr_dma_regions = msg->in_size / sizeof(struct vfio_user_dma_region);
     if (nr_dma_regions == 0) {
         return 0;
     }
@@ -516,12 +510,12 @@ handle_dma_map_or_unmap(vfu_ctx_t *vfu_ctx, uint32_t size, bool map,
                 region->offset, region->prot, region->flags);
 
         vfu_log(vfu_ctx, LOG_DEBUG, "%s DMA region %s",
-                map ? "adding" : "removing", rstr);
+                is_map ? "adding" : "removing", rstr);
 
-        if (map) {
+        if (is_map) {
             int fd = -1;
             if (region->flags == VFIO_USER_F_DMA_REGION_MAPPABLE) {
-                fd = consume_fd(fds, nr_fds, fdi++);
+                fd = consume_fd(msg->in_fds, msg->nr_in_fds, fdi++);
                 if (fd < 0) {
                     vfu_log(vfu_ctx, LOG_ERR, "failed to add DMA region %s: "
                             "mappable but fd not provided", rstr);
@@ -592,19 +586,18 @@ handle_dirty_pages_get(vfu_ctx_t *vfu_ctx,
     if (size % sizeof(struct vfio_iommu_type1_dirty_bitmap_get) != 0) {
         return ERROR_INT(EINVAL);
     }
-    *nr_iovecs = 1 + size / sizeof(struct vfio_iommu_type1_dirty_bitmap_get);
+    *nr_iovecs = size / sizeof(struct vfio_iommu_type1_dirty_bitmap_get);
     *iovecs = malloc(*nr_iovecs * sizeof(struct iovec));
     if (*iovecs == NULL) {
         return -1;
     }
 
-    for (i = 1; i < *nr_iovecs; i++) {
-        struct vfio_iommu_type1_dirty_bitmap_get *r = &ranges[(i - 1)]; /* FIXME ugly indexing */
+    for (i = 0; i < *nr_iovecs; i++) {
+        struct vfio_iommu_type1_dirty_bitmap_get *r = &ranges[i];
         ret = dma_controller_dirty_page_get(vfu_ctx->dma,
-                                            (vfu_dma_addr_t)r->iova,
-                                            r->size, r->bitmap.pgsize,
-                                            r->bitmap.size,
-                                            (char**)&((*iovecs)[i].iov_base));
+                                            (vfu_dma_addr_t)r->iova, r->size,
+                                            r->bitmap.pgsize, r->bitmap.size,
+                                            (char **)&((*iovecs)[i].iov_base));
         if (ret != 0) {
             ret = errno;
             goto out;
@@ -625,82 +618,147 @@ out:
 }
 
 int
-MOCK_DEFINE(handle_dirty_pages)(vfu_ctx_t *vfu_ctx, uint32_t size,
-                                struct iovec **iovecs, size_t *nr_iovecs,
-                                struct vfio_iommu_type1_dirty_bitmap *dirty_bitmap)
+MOCK_DEFINE(handle_dirty_pages)(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg)
 {
+    struct vfio_iommu_type1_dirty_bitmap *dirty_bitmap = msg->in_data;
     int ret;
 
     assert(vfu_ctx != NULL);
-    assert(iovecs != NULL);
-    assert(nr_iovecs != NULL);
-    assert(dirty_bitmap != NULL);
+    assert(msg != NULL);
 
-    if (size < sizeof(*dirty_bitmap) || size != dirty_bitmap->argsz) {
-        vfu_log(vfu_ctx, LOG_ERR, "invalid header size %u", size);
+    if (msg->in_size < sizeof(*dirty_bitmap) ||
+        msg->in_size != dirty_bitmap->argsz) {
+        vfu_log(vfu_ctx, LOG_ERR, "invalid header size %zu", msg->in_size);
         return ERROR_INT(EINVAL);
     }
 
-    if (dirty_bitmap->flags & VFIO_IOMMU_DIRTY_PAGES_FLAG_START) {
+    switch (dirty_bitmap->flags) {
+    case VFIO_IOMMU_DIRTY_PAGES_FLAG_START:
         ret = dma_controller_dirty_page_logging_start(vfu_ctx->dma,
-                                                      migration_get_pgsize(vfu_ctx->migration));
-    } else if (dirty_bitmap->flags & VFIO_IOMMU_DIRTY_PAGES_FLAG_STOP) {
+                  migration_get_pgsize(vfu_ctx->migration));
+        break;
+
+    case VFIO_IOMMU_DIRTY_PAGES_FLAG_STOP:
         dma_controller_dirty_page_logging_stop(vfu_ctx->dma);
         ret = 0;
-    } else if (dirty_bitmap->flags & VFIO_IOMMU_DIRTY_PAGES_FLAG_GET_BITMAP) {
-        ret = handle_dirty_pages_get(vfu_ctx, iovecs, nr_iovecs,
-                                     (struct vfio_iommu_type1_dirty_bitmap_get*)(dirty_bitmap + 1),
-                                     size - sizeof(*dirty_bitmap));
-    } else {
+        break;
+
+    case VFIO_IOMMU_DIRTY_PAGES_FLAG_GET_BITMAP: {
+        struct vfio_iommu_type1_dirty_bitmap_get *get;
+        get = (void *)(dirty_bitmap + 1);
+
+        ret = handle_dirty_pages_get(vfu_ctx, &msg->out_iovecs,
+                                     &msg->nr_out_iovecs, get,
+                                     msg->in_size - sizeof(*dirty_bitmap));
+        break;
+    }
+
+    default:
         vfu_log(vfu_ctx, LOG_ERR, "bad flags %#x", dirty_bitmap->flags);
         ret = ERROR_INT(EINVAL);
+        break;
     }
 
     return ret;
 }
 
-static bool
-is_header_valid(vfu_ctx_t *vfu_ctx, struct vfio_user_header *hdr, size_t size)
+static vfu_msg_t *
+alloc_msg(struct vfio_user_header *hdr, int *fds, size_t nr_fds)
 {
-    assert(hdr != NULL);
+    vfu_msg_t *msg;
+    size_t i;
 
-    if (size < sizeof(hdr)) {
-        vfu_log(vfu_ctx, LOG_ERR, "short header read %ld", size);
-        return false;
+    msg = calloc(1, sizeof(*msg));
+
+    if (msg == NULL) {
+        return NULL;
     }
 
-    if (hdr->flags.type != VFIO_USER_F_TYPE_COMMAND) {
-        vfu_log(vfu_ctx, LOG_ERR, "msg%#hx: not a command req", hdr->msg_id);
-        return false;
+    msg->hdr = *hdr;
+    msg->nr_in_fds = nr_fds;
+
+    if (nr_fds > 0) {
+        msg->in_fds = calloc(msg->nr_in_fds, sizeof(int));
+
+        if (msg->in_fds == NULL) {
+            free(msg);
+            return NULL;
+        }
+
+        for (i = 0; i < msg->nr_in_fds; i++) {
+            msg->in_fds[i] = fds[i];
+        }
     }
 
-    if (hdr->msg_size < sizeof(hdr)) {
-        vfu_log(vfu_ctx, LOG_ERR, "msg%#hx: bad size %d in header",
-                hdr->msg_id, hdr->msg_size);
-        return false;
+    return msg;
+}
+
+static void
+free_msg(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg)
+{
+    int saved_errno = errno;
+    size_t i;
+
+    if (msg == NULL) {
+        return;
     }
 
-    return true;
+    free(msg->in_data);
+
+    for (i = 0; i < msg->nr_in_fds; i++) {
+        if (msg->in_fds[i] != -1) {
+            if (msg->processed_cmd) {
+                vfu_log(vfu_ctx, LOG_DEBUG,
+                        "closing unexpected fd %d (index %zu) from cmd %u",
+                        msg->in_fds[i], i, msg->hdr.cmd);
+            }
+            close(msg->in_fds[i]);
+        }
+    }
+
+    free(msg->in_fds);
+
+    for (i = 0; i < msg->nr_out_fds; i++) {
+        assert(msg->out_fds[i] != -1);
+        close(msg->out_fds[i]);
+    }
+
+    free(msg->out_fds);
+
+    assert(msg->out_data == NULL || msg->out_iovecs == NULL);
+
+    free(msg->out_data);
+
+    /*
+     * Each iov_base refers to data we don't want to free, but we *do* want to
+     * free the allocated array of iovecs if there is one.
+     */
+    free(msg->out_iovecs);
+
+    free(msg);
+
+    errno = saved_errno;
 }
 
 /*
- * Populates @hdr to contain the header for the next command to be processed.
- * Stores any passed FDs into @fds and the number in @nr_fds.
- *
- * Returns 0 if there is no command to process, -1 and errno on error, or the
- * number of bytes read.
+ * Note that we avoid any malloc() before we see data, as this is used for
+ * polling by SPDK.
  */
 int
-MOCK_DEFINE(get_next_command)(vfu_ctx_t *vfu_ctx, struct vfio_user_header *hdr,
-                              int *fds, size_t *nr_fds)
+MOCK_DEFINE(get_request_header)(vfu_ctx_t *vfu_ctx, vfu_msg_t **msgp)
 {
+    int fds[VFIO_USER_CLIENT_MAX_FDS_LIMIT] = { 0 };
+    struct vfio_user_header hdr = { 0, };
+    size_t nr_fds = VFIO_USER_CLIENT_MAX_FDS_LIMIT;
+    size_t i;
     int ret;
 
-    ret = vfu_ctx->tran->get_request(vfu_ctx, hdr, fds, nr_fds);
+    ret = vfu_ctx->tran->get_request_header(vfu_ctx, &hdr, fds, &nr_fds);
+
     if (unlikely(ret < 0)) {
         switch (errno) {
         case EAGAIN:
-            return 0;
+            return -1;
 
         case ENOMSG:
             vfu_reset_ctx(vfu_ctx, "closed");
@@ -716,9 +774,50 @@ MOCK_DEFINE(get_next_command)(vfu_ctx_t *vfu_ctx, struct vfio_user_header *hdr,
         }
     }
 
+    *msgp = alloc_msg(&hdr, fds, nr_fds);
+
+    if (*msgp == NULL) {
+        ret = -1;
+        goto out;
+    }
+
+    return 0;
+
+out:
+    if (ret != 0) {
+        int saved_errno = errno;
+        for (i = 0; i < nr_fds; i++) {
+            close(fds[i]);
+        }
+        errno = saved_errno;
+    }
+
     return ret;
 }
 
+static bool
+is_valid_header(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg)
+{
+    if (msg->hdr.flags.type != VFIO_USER_F_TYPE_COMMAND) {
+        vfu_log(vfu_ctx, LOG_ERR, "msg%#hx: not a command req", msg->hdr.msg_id);
+        return false;
+    }
+
+    if (msg->hdr.msg_size < sizeof(msg->hdr)) {
+        vfu_log(vfu_ctx, LOG_ERR, "msg%#hx: bad size %d in header",
+                msg->hdr.msg_id, msg->hdr.msg_size);
+        return false;
+    }
+
+    if (msg->hdr.msg_size > SERVER_MAX_MSG_SIZE) {
+        vfu_log(vfu_ctx, LOG_ERR, "msg%#hx: size of %u is too large",
+                msg->hdr.msg_id, msg->hdr.msg_size);
+        return false;
+    }
+
+    return true;
+}
+
 bool
 MOCK_DEFINE(cmd_allowed_when_stopped_and_copying)(uint16_t cmd)
 {
@@ -730,8 +829,6 @@ MOCK_DEFINE(cmd_allowed_when_stopped_and_copying)(uint16_t cmd)
 bool
 MOCK_DEFINE(should_exec_command)(vfu_ctx_t *vfu_ctx, uint16_t cmd)
 {
-    assert(vfu_ctx != NULL);
-
     if (device_is_stopped_and_copying(vfu_ctx->migration)) {
         if (!cmd_allowed_when_stopped_and_copying(cmd)) {
             vfu_log(vfu_ctx, LOG_ERR,
@@ -747,138 +844,38 @@ MOCK_DEFINE(should_exec_command)(vfu_ctx_t *vfu_ctx, uint16_t cmd)
     return true;
 }
 
-/*
- * Theoretically, there are still situations in which free() might change errno.
- * See https://sourceware.org/bugzilla/show_bug.cgi?id=17924
- */
-static void
-sfree(void *data)
-{
-    int saved_errno = errno;
-    free(data);
-    errno = saved_errno;
-}
-
 int
-MOCK_DEFINE(exec_command)(vfu_ctx_t *vfu_ctx, struct vfio_user_header *hdr,
-                          size_t size, int *fds, size_t nr_fds, int **fds_out,
-                          size_t *nr_fds_out, struct iovec *_iovecs,
-                          struct iovec **iovecs, size_t *nr_iovecs,
-                          bool *free_iovec_data)
+MOCK_DEFINE(exec_command)(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg)
 {
-    int ret;
-    struct vfio_irq_info *irq_info;
-    struct vfio_device_info *dev_info;
-    struct vfio_region_info *dev_region_info_in, *dev_region_info_out = NULL;
-    void *cmd_data = NULL;
-    size_t cmd_data_size;
-
-    assert(vfu_ctx != NULL);
-    assert(hdr != NULL);
-    assert(fds != NULL);
-    assert(_iovecs != NULL);
-    assert(iovecs != NULL);
-    assert(free_iovec_data != NULL);
-
-    if (!is_header_valid(vfu_ctx, hdr, size)) {
-        return ERROR_INT(EINVAL);
-    }
-
-    if (!should_exec_command(vfu_ctx, hdr->cmd)) {
-        return ERROR_INT(EINVAL);
-    }
-
-    cmd_data_size = hdr->msg_size - sizeof(*hdr);
-
-    if (cmd_data_size > 0) {
-        ret = vfu_ctx->tran->recv_body(vfu_ctx, hdr, &cmd_data);
+    int ret = 0;
 
-        if (ret < 0) {
-            if (errno == ENOMSG) {
-                vfu_reset_ctx(vfu_ctx, "closed");
-                return ERROR_INT(ENOTCONN);
-            } else if (errno == ECONNRESET) {
-                vfu_reset_ctx(vfu_ctx, "reset");
-                return ERROR_INT(ENOTCONN);
-            } else {
-                return -1;
-            }
-        }
-    }
+    msg->processed_cmd = true;
 
-    switch (hdr->cmd) {
+    switch (msg->hdr.cmd) {
     case VFIO_USER_DMA_MAP:
     case VFIO_USER_DMA_UNMAP:
-        ret = handle_dma_map_or_unmap(vfu_ctx, cmd_data_size,
-                                      hdr->cmd == VFIO_USER_DMA_MAP,
-                                      fds, nr_fds, cmd_data);
+        ret = handle_dma_map_or_unmap(vfu_ctx, msg);
         break;
 
     case VFIO_USER_DEVICE_GET_INFO:
-        dev_info = calloc(1, sizeof(*dev_info));
-        if (dev_info == NULL) {
-            ret = ERROR_INT(errno);
-            break;
-        }
-        ret = handle_device_get_info(vfu_ctx, cmd_data_size, cmd_data,
-                                     dev_info);
-        if (ret >= 0) {
-            _iovecs[1].iov_base = dev_info;
-            _iovecs[1].iov_len = dev_info->argsz;
-            *iovecs = _iovecs;
-            *nr_iovecs = 2;
-        } else {
-            sfree(dev_info);
-        }
+        ret = handle_device_get_info(vfu_ctx, msg);
         break;
 
     case VFIO_USER_DEVICE_GET_REGION_INFO:
-        dev_region_info_in = cmd_data;
-        ret = handle_device_get_region_info(vfu_ctx, cmd_data_size,
-                                            dev_region_info_in,
-                                            &dev_region_info_out, fds_out,
-                                            nr_fds_out);
-        if (ret == 0) {
-            _iovecs[1].iov_base = dev_region_info_out;
-            _iovecs[1].iov_len = dev_region_info_in->argsz;
-            *iovecs = _iovecs;
-            *nr_iovecs = 2;
-        }
+        ret = handle_device_get_region_info(vfu_ctx, msg);
         break;
 
     case VFIO_USER_DEVICE_GET_IRQ_INFO:
-        irq_info = calloc(1, sizeof(*irq_info));
-        if (irq_info == NULL) {
-            ret = ERROR_INT(errno);
-            break;
-        }
-        ret = handle_device_get_irq_info(vfu_ctx, cmd_data_size, cmd_data,
-                                         irq_info);
-        if (ret == 0) {
-            _iovecs[1].iov_base = irq_info;
-            _iovecs[1].iov_len = sizeof(*irq_info);
-            *iovecs = _iovecs;
-            *nr_iovecs = 2;
-        } else {
-            sfree(irq_info);
-        }
+        ret = handle_device_get_irq_info(vfu_ctx, msg);
         break;
 
     case VFIO_USER_DEVICE_SET_IRQS:
-        ret = handle_device_set_irqs(vfu_ctx, cmd_data_size, fds, nr_fds,
-                                     cmd_data);
+        ret = handle_device_set_irqs(vfu_ctx, msg);
         break;
 
     case VFIO_USER_REGION_READ:
     case VFIO_USER_REGION_WRITE:
-        ret = handle_region_access(vfu_ctx, cmd_data_size, hdr->cmd,
-                                   &(_iovecs[1].iov_base),
-                                   &(_iovecs[1].iov_len),
-                                   cmd_data);
-        if (ret == 0) {
-            *iovecs = _iovecs;
-            *nr_iovecs = 2;
-        }
+        ret = handle_region_access(vfu_ctx, msg);
         break;
 
     case VFIO_USER_DEVICE_RESET:
@@ -888,94 +885,79 @@ MOCK_DEFINE(exec_command)(vfu_ctx_t *vfu_ctx, struct vfio_user_header *hdr,
     case VFIO_USER_DIRTY_PAGES:
         // FIXME: don't allow migration calls if migration == NULL
         if (vfu_ctx->dma != NULL) {
-            ret = handle_dirty_pages(vfu_ctx, cmd_data_size, iovecs,
-                                     nr_iovecs, cmd_data);
+            ret = handle_dirty_pages(vfu_ctx, msg);
         } else {
             ret = 0;
         }
-        if (ret >= 0) {
-            *free_iovec_data = false;
-        }
         break;
 
     default:
-        vfu_log(vfu_ctx, LOG_ERR, "bad command %d", hdr->cmd);
+        msg->processed_cmd = false;
+        vfu_log(vfu_ctx, LOG_ERR, "bad command %d", msg->hdr.cmd);
         ret = ERROR_INT(EINVAL);
         break;
     }
 
-    sfree(cmd_data);
     return ret;
 }
 
+/*
+ * Handle requests over the vfio-user socket. This can return immediately if we
+ * are non-blocking, and there is no request from the client ready to read from
+ * the socket. Otherwise, we synchronously process the request in place, and
+ * possibly reply.
+ */
 int
 MOCK_DEFINE(process_request)(vfu_ctx_t *vfu_ctx)
 {
-    struct vfio_user_header hdr = { 0, };
-    int *fds = NULL, *fds_out = NULL;
-    size_t nr_fds, i;
-    size_t nr_fds_out = 0;
-    struct iovec _iovecs[2] = { { 0, } };
-    struct iovec *iovecs = NULL;
-    size_t nr_iovecs = 0;
-    bool free_iovec_data = true;
-    int saved_errno;
+    vfu_msg_t *msg = NULL;
     int ret;
 
     assert(vfu_ctx != NULL);
 
-    /*
-     * FIXME if migration device state is VFIO_DEVICE_STATE_STOP then only
-     * migration-related operations should execute. However, some operations
-     * are harmless (e.g. get region info). At the minimum we should fail
-     * accesses to device regions other than the migration region. I'd expect
-     * DMA unmap and get dirty pages to be required even in the stop-and-copy
-     * state.
-     */
+    ret = get_request_header(vfu_ctx, &msg);
 
-    nr_fds = vfu_ctx->client_max_fds;
-    fds = alloca(nr_fds * sizeof(int));
+    if (ret < 0) {
+        if (errno == EAGAIN || errno == EWOULDBLOCK) {
+            ret = 0;
+        }
+        goto out;
+    }
 
-    ret = get_next_command(vfu_ctx, &hdr, fds, &nr_fds);
-    if (ret <= 0) {
-        return ret;
+    if (!is_valid_header(vfu_ctx, msg)) {
+        ret = ERROR_INT(EINVAL);
+        goto out;
     }
 
-    ret = exec_command(vfu_ctx, &hdr, ret, fds, nr_fds, &fds_out, &nr_fds_out,
-                       _iovecs, &iovecs, &nr_iovecs, &free_iovec_data);
+    msg->in_size = msg->hdr.msg_size - sizeof(msg->hdr);
 
-    saved_errno = errno;
+    if (msg->in_size > 0) {
+        ret = vfu_ctx->tran->recv_body(vfu_ctx, msg);
 
-    for (i = 0; i < nr_fds; i++) {
-        if (fds[i] != -1) {
-            vfu_log(vfu_ctx, LOG_DEBUG,
-                    "closing unexpected fd %d (index %zu) from cmd %u",
-                    fds[i], i, hdr.cmd);
-            close(fds[i]);
+        if (ret < 0) {
+            goto out;
         }
     }
 
-    errno = saved_errno;
+    if (!should_exec_command(vfu_ctx, msg->hdr.cmd)) {
+        ret = ERROR_INT(EINVAL);
+        goto out;
+    }
 
-    if (ret < 0) {
-        vfu_log(vfu_ctx, LOG_ERR, "msg%#hx: cmd %d failed: %m", hdr.msg_id,
-                hdr.cmd);
+    ret = exec_command(vfu_ctx, msg);
 
-        if (errno == ENOTCONN) {
-            goto out;
-        }
-    } else {
-        ret = 0;
+    if (ret < 0) {
+        vfu_log(vfu_ctx, LOG_ERR, "msg%#hx: cmd %d failed: %m", msg->hdr.msg_id,
+                msg->hdr.cmd);
     }
 
-    if (hdr.flags.no_reply) {
+    if (msg->hdr.flags.no_reply) {
         /*
          * A failed client request is not a failure of process_request() itself.
          */
         ret = 0;
     } else {
-        ret = vfu_ctx->tran->reply(vfu_ctx, hdr.msg_id, iovecs, nr_iovecs,
-                                   fds_out, nr_fds_out, ret == 0 ? 0 : errno);
+        ret = vfu_ctx->tran->reply(vfu_ctx, msg, ret == 0 ? 0 : errno);
 
         if (ret < 0) {
             vfu_log(vfu_ctx, LOG_ERR, "failed to reply: %m");
@@ -991,22 +973,8 @@ MOCK_DEFINE(process_request)(vfu_ctx_t *vfu_ctx)
     }
 
 out:
-    saved_errno = errno;
-    if (iovecs != NULL) {
-        if (free_iovec_data) {
-            size_t i;
-            for (i = 1; i < nr_iovecs; i++) {
-                free(iovecs[i].iov_base);
-            }
-        }
-        if (iovecs != _iovecs) {
-            free(iovecs);
-        }
-    }
-    free(fds_out);
-    errno = saved_errno;
-
-    return ret == 0 ? 0 : ERROR_INT(errno);
+    free_msg(vfu_ctx, msg);
+    return ret;
 }
 
 int
@@ -1166,7 +1134,6 @@ vfu_destroy_ctx(vfu_ctx_t *vfu_ctx)
     free(vfu_ctx->migration);
     free(vfu_ctx->irqs);
     free(vfu_ctx);
-    // FIXME: Maybe close any open irq efds? Unmap stuff?
 }
 
 void *
diff --git a/lib/private.h b/lib/private.h
index 102fdcd..b397612 100644
--- a/lib/private.h
+++ b/lib/private.h
@@ -38,19 +38,37 @@
 #include "pci_caps.h"
 #include "common.h"
 
-static inline int
-ERROR_INT(int err)
-{
-    errno = err;
-    return -1;
-}
+#define SERVER_MAX_MSG_SIZE 65536
 
-static inline void *
-ERROR_PTR(int err)
-{
-    errno = err;
-    return NULL;
-}
+/*
+ * Structure used to hold an in-flight request+reply.
+ *
+ * Incoming request body and fds are stored in in_*.
+ *
+ * Outgoing requests are either stored in out_data, or out_iovecs. In the latter
+ * case, the iovecs refer to data that should not be freed.
+ */
+typedef struct {
+    /* in/out */
+    struct vfio_user_header hdr;
+
+    bool processed_cmd;
+
+    int *in_fds;
+    size_t nr_in_fds;
+
+    void *in_data;
+    size_t in_size;
+
+    int *out_fds;
+    size_t nr_out_fds;
+
+    void *out_data;
+    size_t out_size;
+
+    struct iovec *out_iovecs;
+    size_t nr_out_iovecs;
+} vfu_msg_t;
 
 struct transport_ops {
     int (*init)(vfu_ctx_t *vfu_ctx);
@@ -59,15 +77,12 @@ struct transport_ops {
 
     int (*attach)(vfu_ctx_t *vfu_ctx);
 
-    int (*get_request)(vfu_ctx_t *vfu_ctx, struct vfio_user_header *hdr,
-                       int *fds, size_t *nr_fds);
+    int (*get_request_header)(vfu_ctx_t *vfu_ctx, struct vfio_user_header *hdr,
+                              int *fds, size_t *nr_fds);
 
-    int (*recv_body)(vfu_ctx_t *vfu_ctx, const struct vfio_user_header *hdr,
-                     void **datap);
+    int (*recv_body)(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg);
 
-    int (*reply)(vfu_ctx_t *vfu_ctx, uint16_t msg_id,
-                 struct iovec *iovecs, size_t nr_iovecs,
-                 int *fds, int count, int err);
+    int (*reply)(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg, int err);
 
     int (*send_msg)(vfu_ctx_t *vfu_ctx, uint16_t msg_id,
                     enum vfio_user_command cmd,
@@ -139,6 +154,20 @@ struct vfu_ctx {
     vfu_dev_type_t          dev_type;
 };
 
+static inline int
+ERROR_INT(int err)
+{
+    errno = err;
+    return -1;
+}
+
+static inline void *
+ERROR_PTR(int err)
+{
+    errno = err;
+    return NULL;
+}
+
 void
 dump_buffer(const char *prefix, const char *buf, uint32_t count);
 
@@ -146,59 +175,26 @@ int
 consume_fd(int *fds, size_t nr_fds, size_t index);
 
 int
-dev_get_reginfo(vfu_ctx_t *vfu_ctx, uint32_t index, uint32_t argsz,
-                struct vfio_region_info **vfio_reg, int **fds, size_t *nr_fds);
+handle_dma_map_or_unmap(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg);
 
 int
-handle_dma_map_or_unmap(vfu_ctx_t *vfu_ctx, uint32_t size, bool map,
-                        int *fds, size_t nr_fds,
-                        struct vfio_user_dma_region *dma_regions);
+handle_device_get_info(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg);
 
 int
-handle_device_get_info(vfu_ctx_t *vfu_ctx, uint32_t size,
-                       struct vfio_device_info *in_dev_info,
-                       struct vfio_device_info *out_dev_info);
+handle_device_get_region_info(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg);
 
-int
-handle_device_set_irqs(vfu_ctx_t *vfu_ctx, uint32_t size,
-                       int *fds, size_t nr_fds, struct vfio_irq_set *irq_set);
+MOCK_DECLARE(int, handle_dirty_pages, vfu_ctx_t *vfu_ctx, vfu_msg_t *msg);
 
-MOCK_DECLARE(bool, should_exec_command, vfu_ctx_t *vfu_ctx, uint16_t cmd);
+MOCK_DECLARE(int, get_request_header, vfu_ctx_t *vfu_ctx, vfu_msg_t **msgp);
 
 MOCK_DECLARE(bool, cmd_allowed_when_stopped_and_copying, uint16_t cmd);
 
-MOCK_DECLARE(int, get_next_command, vfu_ctx_t *vfu_ctx,
-             struct vfio_user_header *hdr, int *fds, size_t *nr_fds);
-
-MOCK_DECLARE(int, exec_command, vfu_ctx_t *vfu_ctx,
-             struct vfio_user_header *hdr, size_t size, int *fds, size_t nr_fds,
-             int **fds_out, size_t *nr_fds_out, struct iovec *_iovecs,
-             struct iovec **iovecs, size_t *nr_iovecs, bool *free_iovec_data);
-
-MOCK_DECLARE(int, process_request, vfu_ctx_t *vfu_ctx);
-
-MOCK_DECLARE(int, handle_dirty_pages, vfu_ctx_t *vfu_ctx, uint32_t size,
-             struct iovec **iovecs, size_t *nr_iovecs,
-             struct vfio_iommu_type1_dirty_bitmap *dirty_bitmap);
-
 MOCK_DECLARE(bool, should_exec_command, vfu_ctx_t *vfu_ctx, uint16_t cmd);
 
-MOCK_DECLARE(bool, cmd_allowed_when_stopped_and_copying, uint16_t cmd);
-
-MOCK_DECLARE(int, get_next_command, vfu_ctx_t *vfu_ctx,
-             struct vfio_user_header *hdr, int *fds, size_t *nr_fds);
-
-MOCK_DECLARE(int, exec_command, vfu_ctx_t *vfu_ctx,
-             struct vfio_user_header *hdr, size_t size, int *fds, size_t nr_fds,
-             int **fds_out, size_t *nr_fds_out, struct iovec *_iovecs,
-             struct iovec **iovecs, size_t *nr_iovecs, bool *free_iovec_data);
+MOCK_DECLARE(int, exec_command, vfu_ctx_t *vfu_ctx, vfu_msg_t *msg);
 
 MOCK_DECLARE(int, process_request, vfu_ctx_t *vfu_ctx);
 
-MOCK_DECLARE(int, handle_dirty_pages, vfu_ctx_t *vfu_ctx, uint32_t size,
-             struct iovec **iovecs, size_t *nr_iovecs,
-             struct vfio_iommu_type1_dirty_bitmap *dirty_bitmap);
-
 #endif /* LIB_VFIO_USER_PRIVATE_H */
 
 /* ex: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab: */
diff --git a/lib/tran_sock.c b/lib/tran_sock.c
index 3785e67..391859c 100644
--- a/lib/tran_sock.c
+++ b/lib/tran_sock.c
@@ -51,8 +51,6 @@
 // FIXME: is this the value we want?
 #define SERVER_MAX_FDS 8
 
-#define SERVER_MAX_MSG_SIZE 65536
-
 typedef struct {
     int listen_fd;
     int conn_fd;
@@ -427,7 +425,7 @@ tran_sock_init(vfu_ctx_t *vfu_ctx)
         goto out;
     }
 
-    /* start listening business */
+    /* start listening for business */
     ret = bind(ts->listen_fd, (struct sockaddr *)&addr, sizeof(addr));
     if (ret < 0) {
         ret = errno;
@@ -754,8 +752,8 @@ tran_sock_attach(vfu_ctx_t *vfu_ctx)
 }
 
 static int
-tran_sock_get_request(vfu_ctx_t *vfu_ctx, struct vfio_user_header *hdr,
-                      int *fds, size_t *nr_fds)
+tran_sock_get_request_header(vfu_ctx_t *vfu_ctx, struct vfio_user_header *hdr,
+                             int *fds, size_t *nr_fds)
 {
     tran_sock_t *ts;
     int sock_flags = 0;
@@ -782,69 +780,83 @@ tran_sock_get_request(vfu_ctx_t *vfu_ctx, struct vfio_user_header *hdr,
 }
 
 static int
-tran_sock_recv_body(vfu_ctx_t *vfu_ctx, const struct vfio_user_header *hdr,
-                    void **datap)
+tran_sock_recv_body(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg)
 {
-    size_t body_size;
     tran_sock_t *ts;
-    void *data;
     int ret;
 
     assert(vfu_ctx != NULL);
     assert(vfu_ctx->tran_data != NULL);
-    assert(hdr != NULL);
-
-    if (hdr->msg_size > SERVER_MAX_MSG_SIZE) {
-        vfu_log(vfu_ctx, LOG_ERR, "msg%#hx: size of %u is too large",
-                hdr->msg_id, hdr->msg_size);
-        return ERROR_INT(EINVAL);
-    }
+    assert(msg != NULL);
 
     ts = vfu_ctx->tran_data;
 
-    body_size = hdr->msg_size - sizeof(*hdr);
-
-    data = malloc(body_size);
+    msg->in_data = malloc(msg->in_size);
 
-    if (data == NULL) {
+    if (msg->in_data == NULL) {
         return -1;
     }
 
-    ret = recv(ts->conn_fd, data, body_size, 0);
+    ret = recv(ts->conn_fd, msg->in_data, msg->in_size, 0);
 
     if (ret < 0) {
         ret = errno;
-        free(data);
+        free(msg->in_data);
+        msg->in_data = NULL;
         return ERROR_INT(ret);
     } else if (ret == 0) {
-        free(data);
+        free(msg->in_data);
+        msg->in_data = NULL;
         return ERROR_INT(ENOMSG);
-    } else if (ret != (int)body_size) {
+    } else if (ret != (int)msg->in_size)  {
         vfu_log(vfu_ctx, LOG_ERR, "msg%#hx: short read: expected=%zu, actual=%d",
-                hdr->msg_id, body_size, ret);
-        free(data);
-        return ERROR_INT(ECONNRESET);
+                msg->hdr.msg_id, msg->in_size, ret);
+        free(msg->in_data);
+        msg->in_data = NULL;
+        return ERROR_INT(EINVAL);
     }
 
-    *datap = data;
     return 0;
 }
 
 static int
-tran_sock_reply(vfu_ctx_t *vfu_ctx, uint16_t msg_id,
-                struct iovec *iovecs, size_t nr_iovecs,
-                int *fds, int count, int err)
+tran_sock_reply(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg, int err)
 {
+    struct iovec *iovecs;
+    size_t nr_iovecs;
     tran_sock_t *ts;
+    int ret;
 
     assert(vfu_ctx != NULL);
     assert(vfu_ctx->tran_data != NULL);
+    assert(msg != NULL);
 
     ts = vfu_ctx->tran_data;
 
+    /* First iovec entry is for msg header. */
+    nr_iovecs = (msg->nr_out_iovecs != 0) ? (msg->nr_out_iovecs + 1) : 2;
+    iovecs = calloc(nr_iovecs, sizeof(*iovecs));
+
+    if (iovecs == NULL) {
+        return -1;
+    }
+
+    if (msg->out_iovecs != NULL) {
+        bcopy(msg->out_iovecs, iovecs + 1,
+              msg->nr_out_iovecs * sizeof(*iovecs));
+    } else {
+        iovecs[1].iov_base = msg->out_data;
+        iovecs[1].iov_len = msg->out_size;
+    }
+
     // FIXME: SPEC: should the reply include the command? I'd say yes?
-    return tran_sock_send_iovec(ts->conn_fd, msg_id, true, 0,
-                                iovecs, nr_iovecs, fds, count, err);
+    ret = tran_sock_send_iovec(ts->conn_fd, msg->hdr.msg_id, true, 0,
+                               iovecs, nr_iovecs,
+                               msg->out_fds, msg->nr_out_fds, err);
+
+    free(iovecs);
+
+    return ret;
 }
 
 static int
@@ -904,7 +916,7 @@ struct transport_ops tran_sock_ops = {
     .init = tran_sock_init,
     .get_poll_fd = tran_sock_get_poll_fd,
     .attach = tran_sock_attach,
-    .get_request = tran_sock_get_request,
+    .get_request_header = tran_sock_get_request_header,
     .recv_body = tran_sock_recv_body,
     .reply = tran_sock_reply,
     .send_msg = tran_sock_send_msg,
diff --git a/samples/client.c b/samples/client.c
index 0647328..e27a323 100644
--- a/samples/client.c
+++ b/samples/client.c
@@ -434,7 +434,7 @@ static int
 access_region(int sock, int region, bool is_write, uint64_t offset,
             void *data, size_t data_len)
 {
-    static int msg_id = -1;
+    static int msg_id = 0xf00f;
     struct vfio_user_region_access send_region_access = {
         .offset = offset,
         .region = region,
diff --git a/test/mocks.c b/test/mocks.c
index 3361c5a..be9767f 100644
--- a/test/mocks.c
+++ b/test/mocks.c
@@ -59,11 +59,10 @@ static struct function funcs[] = {
     { .name = "device_is_stopped_and_copying" },
     { .name = "device_is_stopped" },
     { .name = "dma_controller_add_region" },
-    { .name = "dma_controller_unmap_region" },
     { .name = "dma_controller_remove_region" },
-    { .name = "dma_map_region" },
+    { .name = "dma_controller_unmap_region" },
     { .name = "exec_command" },
-    { .name = "get_next_command" },
+    { .name = "get_request_header" },
     { .name = "handle_dirty_pages" },
     { .name = "process_request" },
     { .name = "should_exec_command" },
@@ -156,6 +155,24 @@ dma_controller_unmap_region(dma_controller_t *dma,
     check_expected(region);
 }
 
+int
+tran_sock_send_iovec(int sock, uint16_t msg_id, bool is_reply,
+                     enum vfio_user_command cmd,
+                     struct iovec *iovecs, size_t nr_iovecs,
+                     int *fds, int count, int err)
+{
+    check_expected(sock);
+    check_expected(msg_id);
+    check_expected(is_reply);
+    check_expected(cmd);
+    check_expected(iovecs);
+    check_expected(nr_iovecs);
+    check_expected(fds);
+    check_expected(count);
+    check_expected(err);
+    return mock();
+}
+
 bool
 device_is_stopped(struct migration *migration)
 {
@@ -167,57 +184,22 @@ device_is_stopped(struct migration *migration)
 }
 
 int
-get_next_command(vfu_ctx_t *vfu_ctx, struct vfio_user_header *hdr,
-                 int *fds, size_t *nr_fds)
+get_request_header(vfu_ctx_t *vfu_ctx, vfu_msg_t **msgp)
 {
     check_expected(vfu_ctx);
-    check_expected(hdr);
-    check_expected(fds);
-    check_expected(nr_fds);
+    check_expected(msgp);
     return mock();
 }
 
 int
-exec_command(vfu_ctx_t *vfu_ctx, struct vfio_user_header *hdr,
-             size_t size, int *fds, size_t nr_fds, int **fds_out,
-             size_t *nr_fds_out, struct iovec *_iovecs,
-             struct iovec **iovecs, size_t *nr_iovecs,
-             bool *free_iovec_data)
+exec_command(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg)
 {
     if (!is_patched("exec_command")) {
-        return __real_exec_command(vfu_ctx, hdr, size, fds, nr_fds, fds_out,
-                                   nr_fds_out, _iovecs, iovecs, nr_iovecs,
-                                   free_iovec_data);
+        return __real_exec_command(vfu_ctx, msg);
     }
     check_expected(vfu_ctx);
-    check_expected(hdr);
-    check_expected(size);
-    check_expected(fds);
-    check_expected(nr_fds);
-    check_expected(fds_out);
-    check_expected(nr_fds_out);
-    check_expected(_iovecs);
-    check_expected(iovecs);
-    check_expected(nr_iovecs);
-    check_expected(free_iovec_data);
-    return mock();
-}
-
-int
-tran_sock_send_iovec(int sock, uint16_t msg_id, bool is_reply,
-                     enum vfio_user_command cmd,
-                     struct iovec *iovecs, size_t nr_iovecs,
-                     int *fds, int count, int err)
-{
-    check_expected(sock);
-    check_expected(msg_id);
-    check_expected(is_reply);
-    check_expected(cmd);
-    check_expected(iovecs);
-    check_expected(nr_iovecs);
-    check_expected(fds);
-    check_expected(count);
-    check_expected(err);
+    check_expected(msg);
+    errno = mock();
     return mock();
 }
 
@@ -265,19 +247,14 @@ should_exec_command(vfu_ctx_t *vfu_ctx, uint16_t cmd)
 }
 
 int
-handle_dirty_pages(vfu_ctx_t *vfu_ctx, uint32_t size,
-                   struct iovec **iovecs, size_t *nr_iovecs,
-                   struct vfio_iommu_type1_dirty_bitmap *dirty_bitmap)
+handle_dirty_pages(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg)
 {
     if (!is_patched("handle_dirty_pages")) {
-        return __real_handle_dirty_pages(vfu_ctx, size, iovecs, nr_iovecs,
-                                         dirty_bitmap);
+        return __real_handle_dirty_pages(vfu_ctx, msg);
     }
     check_expected(vfu_ctx);
-    check_expected(size);
-    check_expected(iovecs);
-    check_expected(nr_iovecs);
-    check_expected(dirty_bitmap);
+    check_expected(msg);
+    errno = mock();
     return mock();
 }
 
diff --git a/test/unit-tests.c b/test/unit-tests.c
index 9aa1106..28eec74 100644
--- a/test/unit-tests.c
+++ b/test/unit-tests.c
@@ -42,57 +42,115 @@
 #include <sys/param.h>
 
 #include "dma.h"
+#include "irq.h"
 #include "libvfio-user.h"
-#include "pci.h"
-#include "private.h"
 #include "migration.h"
+#include "migration_priv.h"
 #include "mocks.h"
+#include "pci.h"
+#include "private.h"
 #include "tran_sock.h"
-#include "migration_priv.h"
+
+#define DMACSIZE (sizeof(dma_controller_t) + sizeof(dma_memory_region_t) * 5)
+
+/*
+ * These globals are used in the unit tests; they're re-initialized each time by
+ * setup(), but having them as globals makes for significantly less
+ * boiler-plate.
+ */
+static char dmacbuf[DMACSIZE];
+static vfu_ctx_t vfu_ctx;
+static vfu_msg_t msg;
+static size_t nr_fds;
+static int fds[2];
+static int ret;
+
+static vfu_msg_t *
+mkmsg(enum vfio_user_command cmd, void *data, size_t size)
+{
+    msg.hdr.cmd = cmd;
+    msg.in_data = data;
+    msg.in_size = size;
+
+    if (nr_fds != 0) {
+        msg.in_fds = fds;
+        msg.nr_in_fds = nr_fds;
+    } else {
+        msg.in_fds = NULL;
+        msg.nr_in_fds = 0;
+    }
+
+    return &msg;
+}
+
+/*
+ * FIXME we shouldn't have to specify a setup function explicitly for each unit
+ * test, cmocka should provide that. E.g. cmocka_run_group_tests enables us to
+ * run a function before/after ALL unit tests have finished, we can extend it
+ * and provide a function to execute before and after each unit test.
+ */
+static int
+setup(void **state UNUSED)
+{
+    memset(&vfu_ctx, 0, sizeof(vfu_ctx));
+
+    vfu_ctx.client_max_fds = 10;
+
+    memset(dmacbuf, 0, DMACSIZE);
+
+    vfu_ctx.dma = (void *)dmacbuf;
+    vfu_ctx.dma->max_regions = 10;
+    vfu_ctx.dma->vfu_ctx = &vfu_ctx;
+
+    memset(&msg, 0, sizeof(msg));
+
+    msg.hdr.flags.type = VFIO_USER_F_TYPE_COMMAND;
+    msg.hdr.msg_size = sizeof(msg.hdr);
+
+    fds[0] = fds[1] = -1;
+    nr_fds = 0;
+    ret = 0;
+
+    unpatch_all();
+    return 0;
+}
 
 static void
 test_dma_map_without_dma(void **state UNUSED)
 {
-    vfu_ctx_t vfu_ctx = { 0 };
-    size_t size = sizeof(struct vfio_user_dma_region);
     struct vfio_user_dma_region dma_region = {
         .flags = VFIO_USER_F_DMA_REGION_MAPPABLE
     };
-    int fd;
 
-    assert_int_equal(0, handle_dma_map_or_unmap(&vfu_ctx, size, true, &fd, 1, &dma_region));
+    vfu_ctx.dma = NULL;
+
+    ret = handle_dma_map_or_unmap(&vfu_ctx, mkmsg(VFIO_USER_DMA_MAP,
+                                  &dma_region, sizeof(dma_region)));
+    assert_int_equal(0, ret);
 }
 
 static void
 test_dma_map_mappable_without_fd(void **state UNUSED)
 {
-    dma_controller_t dma = { 0 };
-    vfu_ctx_t vfu_ctx = { .dma = &dma };
-    size_t size = sizeof(struct vfio_user_dma_region);
     struct vfio_user_dma_region dma_region = {
         .flags = VFIO_USER_F_DMA_REGION_MAPPABLE
     };
-    int fd;
 
-    assert_int_equal(-1, handle_dma_map_or_unmap(&vfu_ctx, size, true, &fd, 0, &dma_region));
+    ret = handle_dma_map_or_unmap(&vfu_ctx, mkmsg(VFIO_USER_DMA_MAP,
+                                  &dma_region, sizeof(dma_region)));
+    assert_int_equal(-1, ret);
     assert_int_equal(errno, EINVAL);
 }
 
 static void
 test_dma_map_without_fd(void **state UNUSED)
 {
-    dma_controller_t dma = { 0 };
-    vfu_ctx_t vfu_ctx = { .dma = &dma };
-    dma.vfu_ctx = &vfu_ctx;
-    size_t size = sizeof(struct vfio_user_dma_region);
-
     struct vfio_user_dma_region r = {
         .addr = 0xdeadbeef,
         .size = 0xcafebabe,
         .offset = 0x8badf00d,
         .prot = PROT_NONE
     };
-    int fd;
 
     patch("dma_controller_add_region");
     will_return(dma_controller_add_region, 0);
@@ -103,7 +161,9 @@ test_dma_map_without_fd(void **state UNUSED)
     expect_value(dma_controller_add_region, fd, -1);
     expect_value(dma_controller_add_region, offset, r.offset);
     expect_value(dma_controller_add_region, prot, r.prot);
-    assert_int_equal(0, handle_dma_map_or_unmap(&vfu_ctx, size, true, &fd, 0, &r));
+    ret = handle_dma_map_or_unmap(&vfu_ctx, mkmsg(VFIO_USER_DMA_MAP,
+                                  &r, sizeof(r)));
+    assert_int_equal(0, ret);
 }
 
 static int
@@ -133,11 +193,7 @@ check_dma_info(const LargestIntegralType value,
 static void
 test_dma_add_regions_mixed(void **state UNUSED)
 {
-    dma_controller_t *dma = alloca(sizeof(*dma) + sizeof(dma_memory_region_t) * 2);
     size_t count = 0;
-    vfu_ctx_t vfu_ctx = { .dma = dma , .dma_register = mock_dma_register,
-                          .pvt = &count };
-    dma->vfu_ctx = &vfu_ctx;
     struct vfio_user_dma_region r[2] = {
         [0] = {
             .addr = 0xdeadbeef,
@@ -153,15 +209,20 @@ test_dma_add_regions_mixed(void **state UNUSED)
             .prot = PROT_READ | PROT_WRITE
         }
     };
-    int fd = 0x0badf00d;
 
-    memset(dma, 0, sizeof(*dma) + sizeof(dma_memory_region_t) * 2);
-    dma->nregions = 2;
-    dma->regions[0].info.mapping.iov_base = (void *)0x123456789;
-    dma->regions[0].info.prot = r[0].prot;
-    dma->regions[1].info.mapping.iov_base = (void *)0x987654321;
-    dma->regions[1].info.vaddr = (void *)0x987654321;
-    dma->regions[1].info.prot = r[1].prot;
+    vfu_ctx.dma_register = mock_dma_register;
+    vfu_ctx.pvt = &count;
+
+    fds[0] = 0x0badf00d;
+    nr_fds = 1;
+
+    vfu_ctx.dma->regions[0].info.iova.iov_base = (void *)r[0].addr;
+    vfu_ctx.dma->regions[0].info.iova.iov_len = r[0].size;
+    vfu_ctx.dma->regions[0].info.prot = r[0].prot;
+    vfu_ctx.dma->regions[1].info.iova.iov_base = (void *)r[1].addr;
+    vfu_ctx.dma->regions[1].info.iova.iov_len = r[1].size;
+    vfu_ctx.dma->regions[1].info.prot = r[1].prot;
+    vfu_ctx.dma->nregions = 2;
 
     patch("dma_controller_add_region");
     /* 1st region */
@@ -175,21 +236,24 @@ test_dma_add_regions_mixed(void **state UNUSED)
     expect_value(dma_controller_add_region, prot, r[0].prot);
     expect_value(mock_dma_register, vfu_ctx, &vfu_ctx);
     expect_check(mock_dma_register, info, check_dma_info,
-                 &dma->regions[0].info);
+                 &vfu_ctx.dma->regions[0].info);
     /* 2nd region */
     will_return(dma_controller_add_region, 0);
     will_return(dma_controller_add_region, 1);
     expect_value(dma_controller_add_region, dma, vfu_ctx.dma);
     expect_value(dma_controller_add_region, dma_addr, r[1].addr);
     expect_value(dma_controller_add_region, size, r[1].size);
-    expect_value(dma_controller_add_region, fd, fd);
+    expect_value(dma_controller_add_region, fd, fds[0]);
     expect_value(dma_controller_add_region, offset, r[1].offset);
     expect_value(dma_controller_add_region, prot, r[1].prot);
     expect_value(mock_dma_register, vfu_ctx, &vfu_ctx);
     expect_check(mock_dma_register, info, check_dma_info,
-                 &dma->regions[1].info);
+                 &vfu_ctx.dma->regions[1].info);
 
-    assert_int_equal(0, handle_dma_map_or_unmap(&vfu_ctx, sizeof(r), true, &fd, 1, r));
+    ret = handle_dma_map_or_unmap(&vfu_ctx, mkmsg(VFIO_USER_DMA_MAP,
+                                  r, sizeof(r)));
+    assert_int_equal(0, ret);
+    assert_int_equal(-1, fds[0]);
 }
 
 /*
@@ -199,9 +263,6 @@ test_dma_add_regions_mixed(void **state UNUSED)
 static void
 test_dma_add_regions_mixed_partial_failure(void **state UNUSED)
 {
-    dma_controller_t dma = { 0 };
-    vfu_ctx_t vfu_ctx = { .dma = &dma };
-    dma.vfu_ctx = &vfu_ctx;
     struct vfio_user_dma_region r[3] = {
         [0] = {
             .addr = 0xdeadbeef,
@@ -223,48 +284,50 @@ test_dma_add_regions_mixed_partial_failure(void **state UNUSED)
             .prot = PROT_READ|PROT_WRITE
         }
     };
-    int fds[] = {0xa, 0xb};
+
+    fds[0] = 0xa;
+    fds[1] = 0xb;
+    nr_fds = 2;
 
     patch("dma_controller_add_region");
 
     /* 1st region */
+    will_return(dma_controller_add_region, 0);
+    will_return(dma_controller_add_region, 0);
     expect_value(dma_controller_add_region, dma, vfu_ctx.dma);
     expect_value(dma_controller_add_region, dma_addr, r[0].addr);
     expect_value(dma_controller_add_region, size, r[0].size);
     expect_value(dma_controller_add_region, fd, -1);
     expect_value(dma_controller_add_region, offset, r[0].offset);
     expect_value(dma_controller_add_region, prot, r[0].prot);
-    will_return(dma_controller_add_region, 0);
-    will_return(dma_controller_add_region, 0);
 
     /* 2nd region */
+    will_return(dma_controller_add_region, 0);
+    will_return(dma_controller_add_region, 1);
     expect_value(dma_controller_add_region, dma, vfu_ctx.dma);
     expect_value(dma_controller_add_region, dma_addr, r[1].addr);
     expect_value(dma_controller_add_region, size, r[1].size);
     expect_value(dma_controller_add_region, fd, fds[0]);
     expect_value(dma_controller_add_region, offset, r[1].offset);
     expect_value(dma_controller_add_region, prot, r[1].prot);
-    will_return(dma_controller_add_region, 0);
-    will_return(dma_controller_add_region, 0);
 
     /* 3rd region */
+    will_return(dma_controller_add_region, EREMOTEIO);
+    will_return(dma_controller_add_region, -1);
     expect_value(dma_controller_add_region, dma, vfu_ctx.dma);
     expect_value(dma_controller_add_region, dma_addr, r[2].addr);
     expect_value(dma_controller_add_region, size, r[2].size);
     expect_value(dma_controller_add_region, fd, fds[1]);
     expect_value(dma_controller_add_region, offset, r[2].offset);
     expect_value(dma_controller_add_region, prot, r[2].prot);
-    will_return(dma_controller_add_region, EREMOTEIO);
-    will_return(dma_controller_add_region, -1);
 
     patch("close");
     expect_value(close, fd, 0xb);
     will_return(close, 0);
 
-    assert_int_equal(-1,
-                     handle_dma_map_or_unmap(&vfu_ctx,
-                                             ARRAY_SIZE(r) * sizeof(struct vfio_user_dma_region),
-                                             true, fds, 2, r));
+    ret = handle_dma_map_or_unmap(&vfu_ctx, mkmsg(VFIO_USER_DMA_MAP,
+                                  r, sizeof(r)));
+    assert_int_equal(-1, ret);
     assert_int_equal(EREMOTEIO, errno);
 }
 
@@ -279,7 +342,6 @@ test_dma_map_return_value(void **state UNUSED)
     vfu_ctx_t vfu_ctx = { .dma = &dma };
     dma.vfu_ctx = &vfu_ctx;
     struct vfio_user_dma_region r = { 0 };
-    int fd = 0;
 
     patch("dma_controller_add_region");
     expect_value(dma_controller_add_region, dma, vfu_ctx.dma);
@@ -291,9 +353,8 @@ test_dma_map_return_value(void **state UNUSED)
     will_return(dma_controller_add_region, 0);
     will_return(dma_controller_add_region, 2);
 
-    assert_int_equal(0,
-        handle_dma_map_or_unmap(&vfu_ctx, sizeof(struct vfio_user_dma_region),
-            true, &fd, 0, &r));
+    assert_int_equal(0, handle_dma_map_or_unmap(&vfu_ctx,
+                     mkmsg(VFIO_USER_DMA_MAP, &r, sizeof(r))));
 }
 
 /*
@@ -302,68 +363,52 @@ test_dma_map_return_value(void **state UNUSED)
 static void
 test_handle_dma_unmap(void **state UNUSED)
 {
-    size_t size = sizeof(dma_controller_t) * sizeof(dma_memory_region_t) * 3;
-    dma_controller_t *d = alloca(size);
-    vfu_ctx_t v = {
-        .dma = d,
-    };
     struct vfio_user_dma_region r = {
         .addr = 0x1000, .size = 0x1000
     };
-    int ret;
 
-    memset(d, 0, size);
+    vfu_ctx.dma->nregions = 3;
+    vfu_ctx.dma->regions[0].info.iova.iov_base = (void *)0x1000;
+    vfu_ctx.dma->regions[0].info.iova.iov_len = 0x1000;
+    vfu_ctx.dma->regions[0].fd = -1;
+    vfu_ctx.dma->regions[1].info.iova.iov_base = (void *)0x4000;
+    vfu_ctx.dma->regions[1].info.iova.iov_len = 0x2000;
+    vfu_ctx.dma->regions[1].fd = -1;
+    vfu_ctx.dma->regions[2].info.iova.iov_base = (void *)0x8000;
+    vfu_ctx.dma->regions[2].info.iova.iov_len = 0x3000;
+    vfu_ctx.dma->regions[2].fd = -1;
 
-    d->nregions = 3;
-    d->regions[0].info.iova.iov_base = (void *)0x1000;
-    d->regions[0].info.iova.iov_len = 0x1000;
-    d->regions[0].fd = -1;
-    d->regions[1].info.iova.iov_base = (void *)0x4000;
-    d->regions[1].info.iova.iov_len = 0x2000;
-    d->regions[1].fd = -1;
-    d->regions[2].info.iova.iov_base = (void *)0x8000;
-    d->regions[2].info.iova.iov_len = 0x3000;
-    d->regions[2].fd = -1;
+    vfu_ctx.dma_unregister = mock_dma_unregister;
 
-    v.dma_unregister = mock_dma_unregister;
-
-    expect_value(mock_dma_unregister, vfu_ctx, &v);
+    expect_value(mock_dma_unregister, vfu_ctx, &vfu_ctx);
     expect_check(mock_dma_unregister, info, check_dma_info,
-                 &d->regions[0].info);
+                 &vfu_ctx.dma->regions[0].info);
     will_return(mock_dma_unregister, 0);
 
-    ret = handle_dma_map_or_unmap(&v, sizeof(r), false, NULL, 0, &r);
+    ret = handle_dma_map_or_unmap(&vfu_ctx, mkmsg(VFIO_USER_DMA_UNMAP,
+                                  &r, sizeof(r)));
 
     assert_int_equal(0, ret);
-    assert_int_equal(2, d->nregions);
-    assert_int_equal(0x4000, d->regions[0].info.iova.iov_base);
-    assert_int_equal(0x2000, d->regions[0].info.iova.iov_len);
-    assert_int_equal(0x8000, d->regions[1].info.iova.iov_base);
-    assert_int_equal(0x3000, d->regions[1].info.iova.iov_len);
+    assert_int_equal(2, vfu_ctx.dma->nregions);
+    assert_int_equal(0x4000, vfu_ctx.dma->regions[0].info.iova.iov_base);
+    assert_int_equal(0x2000, vfu_ctx.dma->regions[0].info.iova.iov_len);
+    assert_int_equal(0x8000, vfu_ctx.dma->regions[1].info.iova.iov_base);
+    assert_int_equal(0x3000, vfu_ctx.dma->regions[1].info.iova.iov_len);
 }
 
 static void
 test_dma_controller_add_region_no_fd(void **state UNUSED)
 {
-    vfu_ctx_t vfu_ctx = { 0 };
-    dma_controller_t *dma = alloca(sizeof(*dma) + sizeof(dma_memory_region_t));
-    void *dma_addr = (void *)0xdeadbeef;
+    vfu_dma_addr_t dma_addr = (void *)0xdeadbeef;
+    dma_memory_region_t *r;
+    off_t offset = 0;
     size_t size = 0;
     int fd = -1;
-    off_t offset = 0;
-    dma_memory_region_t *r;
-
-    memset(dma, 0, sizeof(*dma) + sizeof(dma_memory_region_t));
-
-    dma->vfu_ctx = &vfu_ctx;
-    dma->max_regions = 1;
-
-    assert_int_equal(0,
-                     dma_controller_add_region(dma, dma_addr, size, fd,
-                        offset, PROT_NONE));
 
-    assert_int_equal(1, dma->nregions);
-    r = &dma->regions[0];
+    assert_int_equal(0, dma_controller_add_region(vfu_ctx.dma, dma_addr,
+                                                  size, fd, offset, PROT_NONE));
+    assert_int_equal(1, vfu_ctx.dma->nregions);
+    r = &vfu_ctx.dma->regions[0];
     assert_ptr_equal(NULL, r->info.vaddr);
     assert_ptr_equal(NULL, r->info.mapping.iov_base);
     assert_int_equal(0, r->info.mapping.iov_len);
@@ -379,153 +424,143 @@ test_dma_controller_add_region_no_fd(void **state UNUSED)
 static void
 test_dma_controller_remove_region_mapped(void **state UNUSED)
 {
-    vfu_ctx_t v = { 0 };
-    size_t size = sizeof(dma_controller_t) + sizeof(dma_memory_region_t);
-    dma_controller_t *d = alloca(size);
-    memset(d, 0, size);
-
-    d->vfu_ctx = &v;
-    d->max_regions = d->nregions = 1;
-    d->regions[0].info.iova.iov_base = (void *)0xdeadbeef;
-    d->regions[0].info.iova.iov_len = 0x100;
-    d->regions[0].info.mapping.iov_base = (void *)0xcafebabe;
-    d->regions[0].info.mapping.iov_len = 0x1000;
-    d->regions[0].info.vaddr = (void *)0xcafebabe;
-    expect_value(mock_dma_unregister, vfu_ctx, &v);
+    vfu_ctx.dma->nregions = 1;
+    vfu_ctx.dma->regions[0].info.iova.iov_base = (void *)0xdeadbeef;
+    vfu_ctx.dma->regions[0].info.iova.iov_len = 0x100;
+    vfu_ctx.dma->regions[0].info.mapping.iov_base = (void *)0xcafebabe;
+    vfu_ctx.dma->regions[0].info.mapping.iov_len = 0x1000;
+    vfu_ctx.dma->regions[0].info.vaddr = (void *)0xcafebabe;
+
+    expect_value(mock_dma_unregister, vfu_ctx, &vfu_ctx);
     expect_check(mock_dma_unregister, info, check_dma_info,
-                 &d->regions[0].info);
+                 &vfu_ctx.dma->regions[0].info);
     /* FIXME add unit test when dma_unregister fails */
     will_return(mock_dma_unregister, 0);
     patch("dma_controller_unmap_region");
-    expect_value(dma_controller_unmap_region, dma, d);
-    expect_value(dma_controller_unmap_region, region, &d->regions[0]);
+    expect_value(dma_controller_unmap_region, dma, vfu_ctx.dma);
+    expect_value(dma_controller_unmap_region, region, &vfu_ctx.dma->regions[0]);
     assert_int_equal(0,
-        dma_controller_remove_region(d, (void *)0xdeadbeef, 0x100,
-            mock_dma_unregister, &v));
+        dma_controller_remove_region(vfu_ctx.dma, (void *)0xdeadbeef, 0x100,
+            mock_dma_unregister, &vfu_ctx));
 }
 
 static void
 test_dma_controller_remove_region_unmapped(void **state UNUSED)
 {
-    vfu_ctx_t v = { 0 };
-    size_t size = sizeof(dma_controller_t) + sizeof(dma_memory_region_t);
-    dma_controller_t *d = alloca(size);
-    memset(d, 0, size);
-
-    d->vfu_ctx = &v;
-    d->max_regions = d->nregions = 1;
-    d->regions[0].info.iova.iov_base = (void *)0xdeadbeef;
-    d->regions[0].info.iova.iov_len = 0x100;
-    d->regions[0].fd = -1;
-    expect_value(mock_dma_unregister, vfu_ctx, &v);
+    vfu_ctx.dma->nregions = 1;
+    vfu_ctx.dma->regions[0].info.iova.iov_base = (void *)0xdeadbeef;
+    vfu_ctx.dma->regions[0].info.iova.iov_len = 0x100;
+    vfu_ctx.dma->regions[0].fd = -1;
+
+    expect_value(mock_dma_unregister, vfu_ctx, &vfu_ctx);
     expect_check(mock_dma_unregister, info, check_dma_info,
-                 &d->regions[0].info);
+                 &vfu_ctx.dma->regions[0].info);
     will_return(mock_dma_unregister, 0);
     patch("dma_controller_unmap_region");
     assert_int_equal(0,
-        dma_controller_remove_region(d, (void *)0xdeadbeef, 0x100,
-            mock_dma_unregister, &v));
+        dma_controller_remove_region(vfu_ctx.dma, (void *)0xdeadbeef, 0x100,
+            mock_dma_unregister, &vfu_ctx));
 }
 
-static int fds[] = { 0xab, 0xcd };
-
-static int
-set_fds(const long unsigned int value, const long unsigned int data)
+static void
+test_dma_map_sg(void **state UNUSED)
 {
-    assert(value != 0);
-    if ((void*)data == &get_next_command) {
-        memcpy((int*)value, fds, ARRAY_SIZE(fds) * sizeof(int));
-    } else if ((void*)data == &exec_command) {
-        ((int*)value)[0] = -1;
-    }
-    return 1;
-}
+    dma_sg_t sg = { .region = 1 };
+    struct iovec iovec = { 0 };
 
-static int
-set_nr_fds(const long unsigned int value,
-           const long unsigned int data UNUSED)
-{
-    int *nr_fds = (int*)value;
-    assert(nr_fds != NULL);
-    *nr_fds = ARRAY_SIZE(fds);
-    return 1;
-}
+    /* bad region */
+    assert_int_equal(-1, dma_map_sg(vfu_ctx.dma, &sg, &iovec, 1));
+    assert_int_equal(EINVAL, errno);
 
-typedef struct {
-    int fd;
-    int conn_fd;
-} tran_sock_t;
+    vfu_ctx.dma->nregions = 1;
+
+    /* w/o fd */
+    sg.region = 0;
+    assert_int_equal(-1, dma_map_sg(vfu_ctx.dma, &sg, &iovec, 1));
+    assert_int_equal(EFAULT, errno);
+
+    /* w/ fd */
+    vfu_ctx.dma->regions[0].info.vaddr = (void *)0xdead0000;
+
+    sg.offset = 0x0000beef;
+    sg.length = 0xcafebabe;
+    assert_int_equal(0, dma_map_sg(vfu_ctx.dma, &sg, &iovec, 1));
+    assert_int_equal(0xdeadbeef, iovec.iov_base);
+    assert_int_equal((int)0x00000000cafebabe, iovec.iov_len);
+}
 
-/*
- * Tests that if if exec_command fails then process_request frees passed file
- * descriptors.
- */
 static void
-test_process_command_free_passed_fds(void **state UNUSED)
+test_dma_addr_to_sg(void **state UNUSED)
 {
-    tran_sock_t ts = { .fd = 23, .conn_fd = 24 };
-    vfu_ctx_t vfu_ctx = {
-        .client_max_fds = ARRAY_SIZE(fds),
-        .migration = (struct migration *)0x8badf00d,
-        .tran = &tran_sock_ops,
-        .tran_data = &ts
-    };
+    dma_memory_region_t *r;
+    dma_sg_t sg;
+    int ret;
 
-    patch("get_next_command");
-    expect_value(get_next_command, vfu_ctx, &vfu_ctx);
-    expect_any(get_next_command, hdr);
-    expect_check(get_next_command, fds, &set_fds, &get_next_command);
-    expect_check(get_next_command, nr_fds, &set_nr_fds, NULL);
-    will_return(get_next_command, 0x0000beef);
+    vfu_ctx.dma->nregions = 1;
+    r = &vfu_ctx.dma->regions[0];
+    r->info.iova.iov_base = (void *)0x1000;
+    r->info.iova.iov_len = 0x4000;
+    r->info.vaddr = (void *)0xdeadbeef;
 
-    patch("exec_command");
-    expect_value(exec_command, vfu_ctx, &vfu_ctx);
-    expect_any(exec_command, hdr);
-    expect_value(exec_command, size, 0x0000beef);
-    expect_check(exec_command, fds, &set_fds, &exec_command);
-    expect_any(exec_command, nr_fds);
-    expect_any(exec_command, fds_out);
-    expect_any(exec_command, nr_fds_out);
-    expect_any(exec_command, _iovecs);
-    expect_any(exec_command, iovecs);
-    expect_any(exec_command, nr_iovecs);
-    expect_any(exec_command, free_iovec_data);
-    will_return(exec_command, -0x1234);
+    /* fast path, region hint hit */
+    r->info.prot = PROT_WRITE;
+    ret = dma_addr_to_sg(vfu_ctx.dma, (vfu_dma_addr_t)0x2000,
+                         0x400, &sg, 1, PROT_READ);
+    assert_int_equal(1, ret);
+    assert_int_equal(r->info.iova.iov_base, sg.dma_addr);
+    assert_int_equal(0, sg.region);
+    assert_int_equal(0x2000 - (unsigned long long)r->info.iova.iov_base,
+                     sg.offset);
+    assert_int_equal(0x400, sg.length);
+    assert_true(sg.mappable);
 
-    patch("close");
-    expect_value(close, fd, 0xcd);
-    will_return(close, 0);
+    errno = 0;
+    r->info.prot = PROT_WRITE;
+    ret = dma_addr_to_sg(vfu_ctx.dma, (vfu_dma_addr_t)0x6000,
+                         0x400, &sg, 1, PROT_READ);
+    assert_int_equal(-1, ret);
+    assert_int_equal(ENOENT, errno);
 
-    patch("tran_sock_send_iovec");
-    expect_value(tran_sock_send_iovec, sock, ts.conn_fd);
-    expect_any(tran_sock_send_iovec, msg_id);
-    expect_value(tran_sock_send_iovec, is_reply, true);
-    expect_any(tran_sock_send_iovec, cmd);
-    expect_any(tran_sock_send_iovec, iovecs);
-    expect_any(tran_sock_send_iovec, nr_iovecs);
-    expect_any(tran_sock_send_iovec, fds);
-    expect_any(tran_sock_send_iovec, count);
-    expect_any(tran_sock_send_iovec, err);
-    will_return(tran_sock_send_iovec, 0);
+    r->info.prot = PROT_READ;
+    ret = dma_addr_to_sg(vfu_ctx.dma, (vfu_dma_addr_t)0x2000,
+                         0x400, &sg, 1, PROT_WRITE);
+    assert_int_equal(-1, ret);
+    assert_int_equal(EACCES, errno);
 
-    assert_int_equal(0, process_request(&vfu_ctx));
+    r->info.prot = PROT_READ|PROT_WRITE;
+    ret = dma_addr_to_sg(vfu_ctx.dma, (vfu_dma_addr_t)0x2000,
+                         0x400, &sg, 1, PROT_READ);
+    assert_int_equal(1, ret);
+
+    /* TODO test more scenarios */
+}
+
+static void
+test_vfu_setup_device_dma(void **state UNUSED)
+{
+    vfu_ctx_t vfu_ctx = { 0 };
+
+    assert_int_equal(0, vfu_setup_device_dma(&vfu_ctx, NULL, NULL));
+    assert_non_null(vfu_ctx.dma);
+    free(vfu_ctx.dma);
 }
 
 static void
 test_realize_ctx(void **state UNUSED)
 {
-    vfu_reg_info_t *cfg_reg;
     vfu_reg_info_t reg_info[VFU_PCI_DEV_NUM_REGIONS + 1] = { { 0 } };
-    vfu_ctx_t vfu_ctx = {
-        .reg_info = reg_info,
-        .nr_regions = VFU_PCI_DEV_NUM_REGIONS + 1
-    };
+    vfu_reg_info_t *cfg_reg;
+
+    vfu_ctx.reg_info = reg_info;
+    vfu_ctx.nr_regions = VFU_PCI_DEV_NUM_REGIONS + 1;
 
     assert_int_equal(0, vfu_realize_ctx(&vfu_ctx));
     assert_true(vfu_ctx.realized);
+
     cfg_reg = &vfu_ctx.reg_info[VFU_PCI_DEV_CFG_REGION_IDX];
     assert_int_equal(VFU_REGION_FLAG_RW, cfg_reg->flags);
     assert_int_equal(PCI_CFG_SPACE_SIZE, cfg_reg->size);
+
     assert_non_null(vfu_ctx.pci.config_space);
     assert_non_null(vfu_ctx.irqs);
     assert_int_equal(0, vfu_ctx.pci.nr_caps);
@@ -535,6 +570,11 @@ test_realize_ctx(void **state UNUSED)
     free(vfu_ctx.pci.config_space);
 }
 
+typedef struct {
+    int fd;
+    int conn_fd;
+} tran_sock_t;
+
 static int
 dummy_attach(vfu_ctx_t *vfu_ctx)
 {
@@ -549,9 +589,8 @@ test_attach_ctx(void **state UNUSED)
     struct transport_ops transport_ops = {
         .attach = &dummy_attach,
     };
-    vfu_ctx_t vfu_ctx = {
-        .tran = &transport_ops,
-    };
+
+    vfu_ctx.tran = &transport_ops;
 
     assert_int_equal(0, vfu_attach_ctx(&vfu_ctx));
 }
@@ -559,9 +598,7 @@ test_attach_ctx(void **state UNUSED)
 static void
 test_run_ctx(UNUSED void **state)
 {
-    vfu_ctx_t vfu_ctx = {
-        .realized = false,
-    };
+    vfu_ctx.realized = false;
 
     // device un-realized
     assert_int_equal(-1, vfu_run_ctx(&vfu_ctx));
@@ -588,7 +625,7 @@ test_run_ctx(UNUSED void **state)
 static void
 test_get_region_info(UNUSED void **state)
 {
-    struct iovec iov = { .iov_base = (void*)0x8badf00, .iov_len = 0x0d15ea5e };
+    struct iovec iov = { .iov_base = (void *)0x8badf00, .iov_len = 0x0d15ea5e };
     vfu_reg_info_t reg_info[VFU_PCI_DEV_NUM_REGIONS] = {
         {
             .size = 0xcadebabe
@@ -604,97 +641,129 @@ test_get_region_info(UNUSED void **state)
             .fd = -1
         }
     };
-    vfu_ctx_t vfu_ctx = {
-        .client_max_fds = 1,
-        .nr_regions = ARRAY_SIZE(reg_info),
-        .reg_info = reg_info,
+    struct vfio_region_info_cap_sparse_mmap *sparse;
+    struct vfio_region_info_cap_type *type;
+    struct vfio_region_info in_info = {
+        .index = 0
     };
-    uint32_t index = 0;
-    uint32_t argsz = 0;
-    struct vfio_region_info *vfio_reg;
-    int *fds = NULL;
-    size_t nr_fds;
+    struct vfio_region_info *out_info;
+    int ret;
+
+    vfu_ctx.nr_regions = ARRAY_SIZE(reg_info);
+    vfu_ctx.reg_info = reg_info;
 
     /* bad argsz */
-    assert_int_equal(-1,
-                     dev_get_reginfo(&vfu_ctx, index, argsz, &vfio_reg,
-                                     &fds, &nr_fds));
+    ret = handle_device_get_region_info(&vfu_ctx,
+                                        mkmsg(VFIO_USER_DEVICE_GET_REGION_INFO,
+                                        &in_info, 0));
+    assert_int_equal(-1, ret);
     assert_int_equal(EINVAL, errno);
 
     /* bad region */
-    index = vfu_ctx.nr_regions;
-    argsz = sizeof(struct vfio_region_info);
-    assert_int_equal(-1,
-                     dev_get_reginfo(&vfu_ctx, index, argsz, &vfio_reg,
-                                     &fds, &nr_fds));
+    in_info.index = vfu_ctx.nr_regions;
+    in_info.argsz = sizeof(struct vfio_region_info);
+
+    ret = handle_device_get_region_info(&vfu_ctx,
+                                        mkmsg(VFIO_USER_DEVICE_GET_REGION_INFO,
+                                        &in_info, in_info.argsz));
+    assert_int_equal(-1, ret);
     assert_int_equal(EINVAL, errno);
 
     /* no region caps */
-    index = 1;
-    assert_int_equal(0,
-                     dev_get_reginfo(&vfu_ctx, index, argsz, &vfio_reg,
-                                     &fds, &nr_fds));
-    assert_int_equal(sizeof(struct vfio_region_info), vfio_reg->argsz);
+    in_info.index = 1;
+
+    ret = handle_device_get_region_info(&vfu_ctx,
+                                        mkmsg(VFIO_USER_DEVICE_GET_REGION_INFO,
+                                        &in_info, in_info.argsz));
+    assert_int_equal(0, ret);
+
+    out_info = msg.out_data;
+
+    assert_int_equal(sizeof(struct vfio_region_info), out_info->argsz);
     assert_int_equal(VFIO_REGION_INFO_FLAG_READ | VFIO_REGION_INFO_FLAG_WRITE |
-                     VFIO_REGION_INFO_FLAG_MMAP, vfio_reg->flags);
-    assert_int_equal(1, vfio_reg->index);
-    assert_int_equal(0x10000000000, vfio_reg->offset);
-    assert_int_equal(0xdeadbeef, vfio_reg->size);
-    assert_int_equal(0, nr_fds);
+                     VFIO_REGION_INFO_FLAG_MMAP, out_info->flags);
+    assert_int_equal(1, out_info->index);
+    assert_int_equal(0x10000000000, out_info->offset);
+    assert_int_equal(0xdeadbeef, out_info->size);
+    assert_int_equal(0, msg.nr_out_fds);
 
-    free(vfio_reg);
+    free(msg.out_data);
+    msg.out_data = NULL;
 
     /* regions caps (sparse mmap) but argsz too small */
     vfu_ctx.reg_info[1].mmap_areas = &iov;
     vfu_ctx.reg_info[1].nr_mmap_areas = 1;
 
-    assert_int_equal(0,
-                     dev_get_reginfo(&vfu_ctx, index, argsz, &vfio_reg,
-                                     &fds, &nr_fds));
-    assert_int_equal(argsz + sizeof(struct vfio_region_info_cap_sparse_mmap) + sizeof(struct vfio_region_sparse_mmap_area),
-                     vfio_reg->argsz);
+    ret = handle_device_get_region_info(&vfu_ctx,
+                                        mkmsg(VFIO_USER_DEVICE_GET_REGION_INFO,
+                                        &in_info, in_info.argsz));
+    assert_int_equal(0, ret);
+
+    out_info = msg.out_data;
+
+    assert_int_equal(in_info.argsz +
+                     sizeof(struct vfio_region_info_cap_sparse_mmap) +
+                     sizeof(struct vfio_region_sparse_mmap_area),
+                     out_info->argsz);
     assert_int_equal(VFIO_REGION_INFO_FLAG_READ | VFIO_REGION_INFO_FLAG_WRITE |
                      VFIO_REGION_INFO_FLAG_MMAP | VFIO_REGION_INFO_FLAG_CAPS,
-                     vfio_reg->flags);
-    assert_int_equal(0, nr_fds);
+                     out_info->flags);
 
-    free(vfio_reg);
+    free(msg.out_data);
+    msg.out_data = NULL;
 
     /* region caps and argsz large enough */
-    argsz += sizeof(struct vfio_region_info_cap_sparse_mmap) + sizeof(struct vfio_region_sparse_mmap_area);
-    assert_int_equal(0,
-                     dev_get_reginfo(&vfu_ctx, index, argsz, &vfio_reg,
-                                     &fds, &nr_fds));
-    struct vfio_region_info_cap_sparse_mmap *sparse = (struct vfio_region_info_cap_sparse_mmap*)(vfio_reg + 1);
+    in_info.argsz += sizeof(struct vfio_region_info_cap_sparse_mmap) +
+                     sizeof(struct vfio_region_sparse_mmap_area);
+
+    ret = handle_device_get_region_info(&vfu_ctx,
+                                        mkmsg(VFIO_USER_DEVICE_GET_REGION_INFO,
+                                        &in_info, in_info.argsz));
+    assert_int_equal(0, ret);
+
+    out_info = msg.out_data;
+
+    sparse = (struct vfio_region_info_cap_sparse_mmap *)(out_info + 1);
     assert_int_equal(VFIO_REGION_INFO_CAP_SPARSE_MMAP, sparse->header.id);
     assert_int_equal(1, sparse->header.version);
     assert_int_equal(0, sparse->header.next);
     assert_int_equal(1, sparse->nr_areas);
-    assert_non_null(fds);
-    assert_int_equal(1, nr_fds);
-    assert_int_equal(0x12345, fds[0]);
 
-    free(vfio_reg);
-    free(fds);
+    assert_int_equal(1, msg.nr_out_fds);
+    assert_int_equal(0x12345, msg.out_fds[0]);
+
+    free(msg.out_fds);
+    msg.out_fds = NULL;
+    msg.nr_out_fds = 0;
+
+    free(msg.out_data);
+    msg.out_data = NULL;
 
     /* migration cap */
-    fds = NULL;
     vfu_ctx.reg_info[1].mmap_areas = NULL;
     vfu_ctx.reg_info[1].nr_mmap_areas = 0;
-    argsz = sizeof(struct vfio_region_info) + sizeof(struct vfio_region_info_cap_type);
-    assert_int_equal(0,
-                     dev_get_reginfo(&vfu_ctx, VFU_PCI_DEV_MIGR_REGION_IDX,
-                                     argsz, &vfio_reg, &fds, &nr_fds));
+
+    in_info.index = VFU_PCI_DEV_MIGR_REGION_IDX;
+    in_info.argsz = sizeof(in_info) + sizeof(struct vfio_region_info_cap_type);
+
+    ret = handle_device_get_region_info(&vfu_ctx,
+                                        mkmsg(VFIO_USER_DEVICE_GET_REGION_INFO,
+                                        &in_info, in_info.argsz));
+    assert_int_equal(0, ret);
+
+    out_info = msg.out_data;
+
     assert_int_equal(VFIO_REGION_INFO_FLAG_READ | VFIO_REGION_INFO_FLAG_WRITE |
-                     VFIO_REGION_INFO_FLAG_CAPS,
-                     vfio_reg->flags);
-    struct vfio_region_info_cap_type *type = (struct vfio_region_info_cap_type*)(vfio_reg + 1);
+                     VFIO_REGION_INFO_FLAG_CAPS, out_info->flags);
+    type = (struct vfio_region_info_cap_type *)(out_info + 1);
     assert_int_equal(VFIO_REGION_INFO_CAP_TYPE, type->header.id);
     assert_int_equal(VFIO_REGION_TYPE_MIGRATION , type->type);
     assert_int_equal(VFIO_REGION_SUBTYPE_MIGRATION, type->subtype);
-    assert_null(fds);
-    assert_int_equal(0, nr_fds);
-    free(vfio_reg);
+    assert_null(msg.out_fds);
+    assert_int_equal(0, msg.nr_out_fds);
+
+    free(msg.out_data);
+    msg.out_data = NULL;
 
     /* FIXME add check  for multiple sparse areas */
 }
@@ -751,7 +820,7 @@ test_vfu_ctx_create(void **state UNUSED)
     vfu_destroy_ctx(vfu_ctx);
 }
 
-bool pci_caps_writing = true;
+static bool pci_caps_writing = true;
 
 static ssize_t
 test_pci_caps_region_cb(vfu_ctx_t *vfu_ctx, char *buf, size_t count,
@@ -778,12 +847,6 @@ static void
 test_pci_caps(void **state UNUSED)
 {
     vfu_pci_config_space_t config_space;
-    vfu_reg_info_t reg_info[VFU_PCI_DEV_NUM_REGIONS] = {
-        [VFU_PCI_DEV_CFG_REGION_IDX] = { .size = PCI_CFG_SPACE_SIZE },
-    };
-    vfu_ctx_t vfu_ctx = { .pci.config_space = &config_space,
-                          .reg_info = reg_info,
-    };
     struct vsc *vsc1 = alloca(sizeof(*vsc1) + 3);
     struct vsc *vsc2 = alloca(sizeof(*vsc2) + 13);
     struct vsc *vsc3 = alloca(sizeof(*vsc3) + 13);
@@ -801,6 +864,8 @@ test_pci_caps(void **state UNUSED)
     ssize_t ret;
     char buf[256];
 
+    vfu_ctx.pci.config_space = &config_space;
+
     memset(&config_space, 0, sizeof(config_space));
 
     vfu_ctx.reg_info = calloc(VFU_PCI_DEV_NUM_REGIONS,
@@ -988,13 +1053,6 @@ static void
 test_pci_ext_caps(void **state UNUSED)
 {
     uint8_t config_space[PCI_CFG_SPACE_EXP_SIZE] = { 0, };
-    vfu_reg_info_t reg_info[VFU_PCI_DEV_NUM_REGIONS] = {
-        [VFU_PCI_DEV_CFG_REGION_IDX] = { .size = PCI_CFG_SPACE_EXP_SIZE },
-    };
-    vfu_ctx_t vfu_ctx = { .pci.config_space = (void *)&config_space,
-                          .reg_info = reg_info,
-    };
-
     struct pcie_ext_cap_hdr *hdr;
     size_t explens[] = {
         sizeof(struct pcie_ext_cap_vsc_hdr) + 5,
@@ -1018,11 +1076,10 @@ test_pci_ext_caps(void **state UNUSED)
     ssize_t ret;
     char buf[512];
 
+    vfu_ctx.pci.config_space = (void *)&config_space;
     vfu_ctx.pci.type = VFU_PCI_TYPE_EXPRESS;
-
     vfu_ctx.reg_info = calloc(VFU_PCI_DEV_NUM_REGIONS,
                               sizeof(*vfu_ctx.reg_info));
-
     vfu_ctx.reg_info[VFU_PCI_DEV_CFG_REGION_IDX].cb = test_pci_ext_caps_region_cb;
     vfu_ctx.reg_info[VFU_PCI_DEV_CFG_REGION_IDX].size = PCI_CFG_SPACE_EXP_SIZE;
 
@@ -1180,54 +1237,45 @@ test_pci_ext_caps(void **state UNUSED)
 static void
 test_device_get_info(void **state UNUSED)
 {
-    vfu_ctx_t vfu_ctx = { .nr_regions = 0xdeadbeef };
-    struct vfio_device_info d_in = { .argsz = sizeof(d_in) };
-    struct vfio_device_info d_out;
+    struct vfio_device_info d_in = { .argsz = sizeof(d_in) + 1 };
+    struct vfio_device_info *d_out;
+    int ret;
 
-    assert_int_equal(0, handle_device_get_info(&vfu_ctx, sizeof(d_in),
-                                               &d_in, &d_out));
-    assert_int_equal(sizeof(d_out), d_out.argsz);
-    assert_int_equal(VFIO_DEVICE_FLAGS_PCI | VFIO_DEVICE_FLAGS_RESET,
-                     d_out.flags);
-    assert_int_equal(vfu_ctx.nr_regions, d_out.num_regions);
-    assert_int_equal(VFU_DEV_NUM_IRQS, d_out.num_irqs);
-}
+    vfu_ctx.nr_regions = 0xdeadbeef;
 
-/*
- * Checks that handle_device_get_info handles correctly struct vfio_device_info
- * with more fields.
- */
-static void
-test_device_get_info_compat(void **state UNUSED)
-{
-    vfu_ctx_t vfu_ctx = { .nr_regions = 0xdeadbeef };
-    struct vfio_device_info d_in = { .argsz = sizeof(d_in) + 1 };
-    struct vfio_device_info d_out;
+    ret = handle_device_get_info(&vfu_ctx, mkmsg(VFIO_USER_DEVICE_GET_INFO,
+                                 &d_in, sizeof (d_in)));
+
+    assert_int_equal(0, ret);
 
-    assert_int_equal(0, handle_device_get_info(&vfu_ctx, sizeof(d_in) + 1,
-                                               &d_in, &d_out));
-    assert_int_equal(sizeof(d_out), d_out.argsz);
+
+    d_out = msg.out_data;
+    assert_int_equal(sizeof(*d_out), d_out->argsz);
     assert_int_equal(VFIO_DEVICE_FLAGS_PCI | VFIO_DEVICE_FLAGS_RESET,
-                     d_out.flags);
-    assert_int_equal(vfu_ctx.nr_regions, d_out.num_regions);
-    assert_int_equal(VFU_DEV_NUM_IRQS, d_out.num_irqs);
-
-    /* fewer fields */
-    assert_int_equal(-1,
-                     handle_device_get_info(&vfu_ctx, (sizeof(d_in)) - 1,
-                                            &d_in, &d_out));
+                     d_out->flags);
+    assert_int_equal(vfu_ctx.nr_regions, d_out->num_regions);
+    assert_int_equal(VFU_DEV_NUM_IRQS, d_out->num_irqs);
+
+    free(msg.out_data);
+    msg.out_data = NULL;
+    msg.out_size = 0;
+
+    /* bad size */
+    ret = handle_device_get_info(&vfu_ctx, mkmsg(VFIO_USER_DEVICE_GET_INFO,
+                                 &d_in, sizeof (d_in) - 1));
+    assert_int_equal(-1, ret);
     assert_int_equal(EINVAL, errno);
+    assert_ptr_equal(NULL, msg.out_data);
+    assert_ptr_equal(0, msg.out_size);
 }
 
-
 /*
  * Performs various checks when adding sparse memory regions.
  */
 static void
 test_setup_sparse_region(void **state UNUSED)
 {
-    vfu_reg_info_t reg_info;
-    vfu_ctx_t vfu_ctx = { .reg_info = &reg_info };
+    vfu_reg_info_t reg_info = { 0 };
     struct iovec mmap_areas[2] = {
         [0] = {
             .iov_base = (void*)0x0,
@@ -1238,7 +1286,8 @@ test_setup_sparse_region(void **state UNUSED)
             .iov_len = 0x1000
         }
     };
-    int ret;
+
+    vfu_ctx.reg_info = &reg_info;
 
     /* invalid mappable settings */
     ret = vfu_setup_region(&vfu_ctx, VFU_PCI_DEV_BAR0_REGION_IDX,
@@ -1275,88 +1324,240 @@ test_setup_sparse_region(void **state UNUSED)
 }
 
 static void
-test_dma_map_sg(void **state UNUSED)
+test_dirty_pages_without_dma(UNUSED void **state)
 {
-    vfu_ctx_t vfu_ctx = { 0 };
-    size_t size = sizeof(dma_controller_t) + sizeof(dma_memory_region_t);
-    dma_controller_t *dma = alloca(size);
-    dma_sg_t sg = { .region = 1 };
-    struct iovec iovec = { 0 };
+    int ret;
 
-    memset(dma, 0, size);
-    dma->vfu_ctx = &vfu_ctx;
-    dma->nregions = 1;
+    /* with DMA controller */
 
-    /* bad region */
-    assert_int_equal(-1, dma_map_sg(dma, &sg, &iovec, 1));
-    assert_int_equal(EINVAL, errno);
+    patch("handle_dirty_pages");
 
-    /* w/o fd */
-    sg.region = 0;
-    assert_int_equal(-1, dma_map_sg(dma, &sg, &iovec, 1));
-    assert_int_equal(EFAULT, errno);
+    expect_value(handle_dirty_pages, vfu_ctx, &vfu_ctx);
+    expect_any(handle_dirty_pages, msg);
+    will_return(handle_dirty_pages, EREMOTEIO);
+    will_return(handle_dirty_pages, -1);
 
-    /* w/ fd */
-    dma->regions[0].info.vaddr = (void *)0xdead0000;
-    sg.offset = 0x0000beef;
-    sg.length = 0xcafebabe;
-    assert_int_equal(0, dma_map_sg(dma, &sg, &iovec, 1));
-    assert_int_equal(0xdeadbeef, iovec.iov_base);
-    assert_int_equal((int)0x00000000cafebabe, iovec.iov_len);
+    ret = exec_command(&vfu_ctx, mkmsg(VFIO_USER_DIRTY_PAGES, NULL, 0));
+    assert_int_equal(-1, ret);
+    assert_int_equal(EREMOTEIO, errno);
+
+    /* without DMA controller */
+
+    vfu_ctx.dma = NULL;
+
+    ret = exec_command(&vfu_ctx, mkmsg(VFIO_USER_DIRTY_PAGES, NULL, 0));
+    assert_int_equal(0, ret);
 
 }
 
 static void
-test_dma_addr_to_sg(void **state UNUSED)
+test_device_set_irqs(UNUSED void **state)
 {
-    dma_controller_t *dma = alloca(sizeof(dma_controller_t) + sizeof(dma_memory_region_t));
-    dma_sg_t sg;
-    dma_memory_region_t *r;
+    vfu_irqs_t *irqs = alloca(sizeof (*irqs) + sizeof (int));
+    struct vfio_irq_set irq_set = { 0, };
+    //int fd = 0xdead;
 
-    dma->nregions = 1;
-    r = &dma->regions[0];
-    r->info.iova.iov_base = (void *)0x1000;
-    r->info.iova.iov_len = 0x4000;
-    r->info.vaddr = (void *)0xdeadbeef;
+    vfu_ctx.irq_count[VFU_DEV_MSIX_IRQ] = 2048;
+    vfu_ctx.irq_count[VFU_DEV_ERR_IRQ] = 1;
+    vfu_ctx.irq_count[VFU_DEV_REQ_IRQ] = 1;
+    vfu_ctx.irqs = irqs;
 
-    /* fast path, region hint hit */
-    r->info.prot = PROT_WRITE;
-    assert_int_equal(1,
-        dma_addr_to_sg(dma, (vfu_dma_addr_t)0x2000, 0x400, &sg, 1, PROT_READ));
-    assert_int_equal(r->info.iova.iov_base, sg.dma_addr);
-    assert_int_equal(0, sg.region);
-    assert_int_equal(0x2000 - (unsigned long long)r->info.iova.iov_base, sg.offset);
-    assert_int_equal(0x400, sg.length);
-    assert_true(sg.mappable);
+    memset(irqs, 0, sizeof (*irqs) + sizeof (int));
 
-    errno = 0;
-    r->info.prot = PROT_WRITE;
-    assert_int_equal(-1,
-        dma_addr_to_sg(dma, (vfu_dma_addr_t)0x6000, 0x400, &sg, 1, PROT_READ));
-    assert_int_equal(ENOENT, errno);
+    irq_set.argsz = sizeof (irq_set);
 
-    r->info.prot = PROT_READ;
-    assert_int_equal(-1,
-        dma_addr_to_sg(dma, (vfu_dma_addr_t)0x2000, 0x400, &sg, 1, PROT_WRITE));
-    assert_int_equal(EACCES, errno);
+    /*
+     * Validation tests.
+     */
 
-    r->info.prot = PROT_READ|PROT_WRITE;
-    assert_int_equal(1,
-        dma_addr_to_sg(dma, (vfu_dma_addr_t)0x2000, 0x400, &sg, 1, PROT_READ));
+    /* bad message size */
+    ret = handle_device_set_irqs(&vfu_ctx, mkmsg(VFIO_USER_DEVICE_SET_IRQS,
+                                 &irq_set, 0));
+    assert_int_equal(-1, ret);
+    assert_int_equal(EINVAL, errno);
 
-    /* TODO test more scenarios */
-}
+    /* bad .argsz */
+    irq_set.argsz = 3;
 
-static void
-test_vfu_setup_device_dma(void **state UNUSED)
-{
-    vfu_ctx_t vfu_ctx = { 0 };
+    ret = handle_device_set_irqs(&vfu_ctx, mkmsg(VFIO_USER_DEVICE_SET_IRQS,
+                                 &irq_set, sizeof (irq_set)));
+    assert_int_equal(-1, ret);
+    assert_int_equal(EINVAL, errno);
 
-    assert_int_equal(0, vfu_setup_device_dma(&vfu_ctx, NULL, NULL));
-    assert_non_null(vfu_ctx.dma);
-    free(vfu_ctx.dma);
-}
+    /* bad .index */
+    irq_set.argsz = sizeof (irq_set);
+    irq_set.index = VFU_DEV_NUM_IRQS;
+
+    ret = handle_device_set_irqs(&vfu_ctx, mkmsg(VFIO_USER_DEVICE_SET_IRQS,
+                                 &irq_set, sizeof (irq_set)));
+    assert_int_equal(-1, ret);
+    assert_int_equal(EINVAL, errno);
+
+    /* bad flags, MASK and UNMASK */
+    irq_set.index = VFU_DEV_MSIX_IRQ;
+    irq_set.flags = VFIO_IRQ_SET_ACTION_MASK | VFIO_IRQ_SET_ACTION_UNMASK;
+
+    ret = handle_device_set_irqs(&vfu_ctx, mkmsg(VFIO_USER_DEVICE_SET_IRQS,
+                                 &irq_set, sizeof (irq_set)));
+    assert_int_equal(-1, ret);
+    assert_int_equal(EINVAL, errno);
+
+    /* bad flags, DATA_NONE and DATA_BOOL */
+    irq_set.flags = VFIO_IRQ_SET_ACTION_MASK | VFIO_IRQ_SET_DATA_NONE |
+                    VFIO_IRQ_SET_DATA_BOOL;
+
+    ret = handle_device_set_irqs(&vfu_ctx, mkmsg(VFIO_USER_DEVICE_SET_IRQS,
+                                 &irq_set, sizeof (irq_set)));
+    assert_int_equal(-1, ret);
+    assert_int_equal(EINVAL, errno);
+
+    /* bad start, count range */
+    irq_set.flags = VFIO_IRQ_SET_ACTION_MASK | VFIO_IRQ_SET_DATA_NONE;
+    irq_set.start = 2047;
+    irq_set.count = 2;
+
+    ret = handle_device_set_irqs(&vfu_ctx, mkmsg(VFIO_USER_DEVICE_SET_IRQS,
+                                 &irq_set, sizeof (irq_set)));
+    assert_int_equal(-1, ret);
+    assert_int_equal(EINVAL, errno);
+
+    /* bad start, count range */
+    irq_set.start = 2049;
+    irq_set.count = 1;
+
+    ret = handle_device_set_irqs(&vfu_ctx, mkmsg(VFIO_USER_DEVICE_SET_IRQS,
+                                 &irq_set, sizeof (irq_set)));
+    assert_int_equal(-1, ret);
+    assert_int_equal(EINVAL, errno);
+
+    /* bad action for err irq */
+    irq_set.start = 0;
+    irq_set.count = 1;
+    irq_set.index = VFU_DEV_ERR_IRQ;
+
+    ret = handle_device_set_irqs(&vfu_ctx, mkmsg(VFIO_USER_DEVICE_SET_IRQS,
+                                 &irq_set, sizeof (irq_set)));
+    assert_int_equal(-1, ret);
+    assert_int_equal(EINVAL, errno);
+
+    /* bad action for req irq */
+    irq_set.index = VFU_DEV_REQ_IRQ;
+
+    ret = handle_device_set_irqs(&vfu_ctx, mkmsg(VFIO_USER_DEVICE_SET_IRQS,
+                                 &irq_set, sizeof (irq_set)));
+    assert_int_equal(-1, ret);
+    assert_int_equal(EINVAL, errno);
+
+    /* bad start for count == 0 */
+    irq_set.start = 1;
+    irq_set.count = 0;
+    irq_set.index = VFU_DEV_MSIX_IRQ;
+
+    ret = handle_device_set_irqs(&vfu_ctx, mkmsg(VFIO_USER_DEVICE_SET_IRQS,
+                                 &irq_set, sizeof (irq_set)));
+    assert_int_equal(-1, ret);
+    assert_int_equal(EINVAL, errno);
+
+    /* bad action for count == 0 */
+    irq_set.flags = VFIO_IRQ_SET_ACTION_MASK | VFIO_IRQ_SET_DATA_NONE;
+    irq_set.count = 0;
+    irq_set.start = 0;
+
+    ret = handle_device_set_irqs(&vfu_ctx, mkmsg(VFIO_USER_DEVICE_SET_IRQS,
+                                 &irq_set, sizeof (irq_set)));
+    assert_int_equal(-1, ret);
+    assert_int_equal(EINVAL, errno);
+
+    /* bad action and data type for count == 0 */
+    irq_set.flags = VFIO_IRQ_SET_ACTION_TRIGGER | VFIO_IRQ_SET_DATA_BOOL;
+    irq_set.count = 0;
+    irq_set.start = 0;
+
+    ret = handle_device_set_irqs(&vfu_ctx, mkmsg(VFIO_USER_DEVICE_SET_IRQS,
+                                 &irq_set, sizeof (irq_set)));
+    assert_int_equal(-1, ret);
+    assert_int_equal(EINVAL, errno);
+
+    /* bad fds for DATA_BOOL */
+    irq_set.flags = VFIO_IRQ_SET_ACTION_TRIGGER | VFIO_IRQ_SET_DATA_BOOL;
+    irq_set.count = 1;
+    irq_set.start = 0;
+    nr_fds = 1;
+    fds[0] = 0xbeef;
+
+    ret = handle_device_set_irqs(&vfu_ctx, mkmsg(VFIO_USER_DEVICE_SET_IRQS,
+                                 &irq_set, sizeof (irq_set)));
+    assert_int_equal(-1, ret);
+    assert_int_equal(EINVAL, errno);
+
+    /* bad fds for DATA_NONE */
+    irq_set.flags = VFIO_IRQ_SET_ACTION_TRIGGER | VFIO_IRQ_SET_DATA_NONE;
+    irq_set.count = 1;
+    irq_set.start = 0;
+
+    ret = handle_device_set_irqs(&vfu_ctx, mkmsg(VFIO_USER_DEVICE_SET_IRQS,
+                                 &irq_set, sizeof (irq_set)));
+    assert_int_equal(-1, ret);
+    assert_int_equal(EINVAL, errno);
+
+    /* bad fds for count == 2 */
+    irq_set.flags = VFIO_IRQ_SET_ACTION_TRIGGER | VFIO_IRQ_SET_DATA_EVENTFD;
+    irq_set.count = 2;
+    irq_set.start = 0;
+
+    ret = handle_device_set_irqs(&vfu_ctx, mkmsg(VFIO_USER_DEVICE_SET_IRQS,
+                                 &irq_set, sizeof (irq_set)));
+    assert_int_equal(-1, ret);
+    assert_int_equal(EINVAL, errno);
+
+    irqs->err_efd = irqs->req_efd = -1;
+
+    /*
+     * Basic disable functionality.
+     */
+
+    nr_fds = 0;
 
+    irq_set.index = VFU_DEV_REQ_IRQ;
+    irq_set.flags = VFIO_IRQ_SET_ACTION_TRIGGER | VFIO_IRQ_SET_DATA_NONE;
+    irq_set.count = 0;
+    irq_set.start = 0;
+
+    ret = handle_device_set_irqs(&vfu_ctx, mkmsg(VFIO_USER_DEVICE_SET_IRQS,
+                                 &irq_set, sizeof (irq_set)));
+    assert_int_equal(0, ret);
+
+    irq_set.index = VFU_DEV_REQ_IRQ;
+    irq_set.flags = VFIO_IRQ_SET_ACTION_TRIGGER | VFIO_IRQ_SET_DATA_EVENTFD;
+    irq_set.count = 1;
+    irq_set.start = 0;
+
+    ret = handle_device_set_irqs(&vfu_ctx, mkmsg(VFIO_USER_DEVICE_SET_IRQS,
+                                 &irq_set, sizeof (irq_set)));
+    assert_int_equal(0, ret);
+
+    /*
+     * Basic enable functionality.
+     */
+
+    irq_set.index = VFU_DEV_MSIX_IRQ;
+    vfu_ctx.irq_count[VFU_DEV_MSIX_IRQ] = 1;
+    irqs->efds[0] = -1;
+
+    nr_fds = 1;
+    fds[0] = 0xbeef;
+
+    irq_set.index = VFU_DEV_MSIX_IRQ;
+    irq_set.flags = VFIO_IRQ_SET_ACTION_TRIGGER | VFIO_IRQ_SET_DATA_EVENTFD;
+    irq_set.count = 1;
+    irq_set.start = 0;
+
+    ret = handle_device_set_irqs(&vfu_ctx, mkmsg(VFIO_USER_DEVICE_SET_IRQS,
+                                 &irq_set, sizeof (irq_set)));
+    assert_int_equal(0, ret);
+    assert_int_equal(0xbeef, irqs->efds[0]);
+
+}
 static void
 test_migration_state_transitions(void **state UNUSED)
 {
@@ -1423,19 +1624,6 @@ test_migration_state_transitions(void **state UNUSED)
     }
 }
 
-/*
- * FIXME we shouldn't have to specify a setup function explicitly for each unit
- * test, cmocka should provide that. E.g. cmocka_run_group_tests enables us to
- * run a function before/after ALL unit tests have finished, we can extend it
- * and provide a function to execute before and after each unit test.
- */
-static int
-setup(void **state UNUSED)
-{
-    unpatch_all();
-    return 0;
-}
-
 static struct test_setup_migr_reg_dat {
     vfu_ctx_t *v;
     size_t rs; /* migration registers size */
@@ -1585,8 +1773,6 @@ test_setup_migration_callbacks(void **state)
 static void
 test_device_is_stopped_and_copying(UNUSED void **state)
 {
-    vfu_ctx_t vfu_ctx = { .migration = NULL };
-
     assert_false(device_is_stopped_and_copying(vfu_ctx.migration));
     assert_false(device_is_stopped(vfu_ctx.migration));
 
@@ -1633,7 +1819,8 @@ static void
 test_should_exec_command(UNUSED void **state)
 {
     struct migration migration = { { 0 } };
-    vfu_ctx_t vfu_ctx = { .migration = &migration };
+
+    vfu_ctx.migration = &migration;
 
     patch("device_is_stopped_and_copying");
     patch("cmd_allowed_when_stopped_and_copying");
@@ -1669,270 +1856,88 @@ test_should_exec_command(UNUSED void **state)
 }
 
 static int
-recv_body(UNUSED vfu_ctx_t *vfu_ctx, UNUSED const struct vfio_user_header *hdr,
-          UNUSED void **datap)
+check_request_header_msg(const LargestIntegralType value,
+                         const LargestIntegralType cvalue UNUSED)
 {
-    /* hack to avoid having to refactor the rest of exec_command */
-    errno = ENOBUFS;
-    return -1;
-}
+    vfu_msg_t **msgp = (vfu_msg_t **)value;
 
-static void
-test_exec_command(UNUSED void **state)
-{
-    vfu_ctx_t vfu_ctx = { 0 };
-    struct vfio_user_header hdr = {
-        .cmd = 0xbeef,
-        .flags.type = VFIO_USER_F_TYPE_COMMAND,
-        .msg_size = sizeof(hdr) + 1
-    };
-    size_t size = sizeof(hdr);
-    int fds = 0;
-    struct iovec _iovecs = { 0 };
-    struct iovec *iovecs = NULL;
-    size_t nr_iovecs = 0;
-    bool free_iovec_data = false;
-    int r;
-
-    /* XXX should NOT execute command */
-    patch("should_exec_command");
-    will_return(should_exec_command, false);
-    expect_value(should_exec_command, vfu_ctx, &vfu_ctx);
-    expect_value(should_exec_command, cmd, 0xbeef);
-    r = exec_command(&vfu_ctx, &hdr, size, &fds, 0, NULL, NULL, &_iovecs,
-                     &iovecs, &nr_iovecs, &free_iovec_data);
-    assert_int_equal(-1, r);
-    assert_int_equal(EINVAL, errno);
+    *msgp = malloc(sizeof(msg));
 
-    /* XXX should execute command */
-    struct transport_ops tran = { .recv_body = recv_body };
-    vfu_ctx.tran = &tran;
-    will_return(should_exec_command, true);
-    expect_value(should_exec_command, vfu_ctx, &vfu_ctx);
-    expect_value(should_exec_command, cmd, 0xbeef);
-    r = exec_command(&vfu_ctx, &hdr, size, &fds, 0, NULL, NULL, &_iovecs,
-                     &iovecs, &nr_iovecs, &free_iovec_data);
-    assert_int_equal(-1, r);
-    assert_int_equal(ENOBUFS, errno);
+    assert_non_null(*msgp);
+
+    memcpy(*msgp, &msg, sizeof(msg));
+
+    return 1;
 }
 
-static void
-test_dirty_pages_without_dma(UNUSED void **state)
+static int
+check_exec_command_msg(const LargestIntegralType value,
+                       const LargestIntegralType cvalue UNUSED)
 {
-    vfu_ctx_t vfu_ctx = { .migration = NULL };
-    struct vfio_user_header hdr = {
-        .cmd = VFIO_USER_DIRTY_PAGES,
-        .flags = {
-            .type = VFIO_USER_F_TYPE_COMMAND
-        },
-        .msg_size = sizeof(hdr)
-    };
-    size_t size = sizeof(hdr);
-    int fds = 0;
-    struct iovec _iovecs = { 0 };
-    struct iovec *iovecs = NULL;
-    size_t nr_iovecs = 0;
-    bool free_iovec_data = false;
-    int r;
+    vfu_msg_t *cmsg = (vfu_msg_t *)value;
 
+    int ret = cmsg->nr_in_fds == ARRAY_SIZE(fds) &&
+              cmsg->in_fds[0] == fds[0] &&
+              cmsg->in_fds[1] == fds[1] &&
+              cmsg->in_data == NULL &&
+              cmsg->in_size == 0 &&
+              memcmp(&cmsg->hdr, &msg.hdr, sizeof (msg.hdr)) == 0;
 
-    patch("handle_dirty_pages");
+    consume_fd(cmsg->in_fds, cmsg->nr_in_fds, 0);
 
-    /* XXX w/o DMA controller */
-    r = exec_command(&vfu_ctx, &hdr, size, &fds, 0, NULL, NULL,
-                     &_iovecs, &iovecs, &nr_iovecs, &free_iovec_data);
-    assert_int_equal(0, r);
-
-    /* XXX w/ DMA controller */
-    vfu_ctx.dma = (void*)0xdeadbeef;
-    expect_value(handle_dirty_pages, vfu_ctx, &vfu_ctx);
-    expect_value(handle_dirty_pages, size, 0);
-    expect_value(handle_dirty_pages, iovecs, &iovecs);
-    expect_value(handle_dirty_pages, nr_iovecs, &nr_iovecs);
-    expect_value(handle_dirty_pages, dirty_bitmap, NULL);
-    will_return(handle_dirty_pages, 0xabcd);
-    r = exec_command(&vfu_ctx, &hdr, size, &fds, 0, NULL, NULL,
-                     &_iovecs, &iovecs, &nr_iovecs, &free_iovec_data);
-    assert_int_equal(0xabcd, r);
+    return ret;
 }
 
+/*
+ * Tests that if if exec_command fails then process_request() frees passed file
+ * descriptors.
+ */
 static void
-test_device_set_irqs(UNUSED void **state)
+test_process_request_free_passed_fds(void **state UNUSED)
 {
-    vfu_irqs_t *irqs = alloca(sizeof (*irqs) + sizeof (int));
-    struct vfio_irq_set irq_set = { 0, };
-    vfu_ctx_t vfu_ctx = { 0, };
-    int fd = 0xdead;
-    int ret;
-
-    vfu_ctx.irq_count[VFU_DEV_MSIX_IRQ] = 2048;
-    vfu_ctx.irq_count[VFU_DEV_ERR_IRQ] = 1;
-    vfu_ctx.irq_count[VFU_DEV_REQ_IRQ] = 1;
-    vfu_ctx.irqs = irqs;
-
-    /* validation tests */
-
-    irq_set.argsz = sizeof (irq_set);
-
-    ret = handle_device_set_irqs(&vfu_ctx, 0, NULL, 0, &irq_set);
-    /* bad message size */
-    assert_int_equal(-1, ret);
-    assert_int_equal(EINVAL, errno);
-
-    irq_set.argsz = 3;
-
-    ret = handle_device_set_irqs(&vfu_ctx, sizeof (irq_set), NULL, 0, &irq_set);
-    /* bad .argsz */
-    assert_int_equal(-1, ret);
-    assert_int_equal(EINVAL, errno);
-
-    irq_set.argsz = sizeof (irq_set);
-    irq_set.index = VFU_DEV_NUM_IRQS;
-
-    ret = handle_device_set_irqs(&vfu_ctx, sizeof (irq_set), NULL, 0, &irq_set);
-    /* bad .index */
-    assert_int_equal(-1, ret);
-    assert_int_equal(EINVAL, errno);
-
-    irq_set.index = VFU_DEV_MSIX_IRQ;
-    irq_set.flags = VFIO_IRQ_SET_ACTION_MASK | VFIO_IRQ_SET_ACTION_UNMASK;
-
-    ret = handle_device_set_irqs(&vfu_ctx, sizeof (irq_set), NULL, 0, &irq_set);
-    /* bad flags, MASK and UNMASK */
-    assert_int_equal(-1, ret);
-    assert_int_equal(EINVAL, errno);
-
-    irq_set.flags = VFIO_IRQ_SET_ACTION_MASK | VFIO_IRQ_SET_DATA_NONE |
-                    VFIO_IRQ_SET_DATA_BOOL;
-
-    ret = handle_device_set_irqs(&vfu_ctx, sizeof (irq_set), NULL, 0, &irq_set);
-    /* bad flags, DATA_NONE and DATA_BOOL */
-    assert_int_equal(-1, ret);
-    assert_int_equal(EINVAL, errno);
-
-    irq_set.flags = VFIO_IRQ_SET_ACTION_MASK | VFIO_IRQ_SET_DATA_NONE;
-    irq_set.start = 2047;
-    irq_set.count = 2;
-
-    ret = handle_device_set_irqs(&vfu_ctx, sizeof (irq_set), NULL, 0, &irq_set);
-    /* bad start, count range */
-    assert_int_equal(-1, ret);
-    assert_int_equal(EINVAL, errno);
-
-    irq_set.start = 2049;
-    irq_set.count = 1;
-
-    ret = handle_device_set_irqs(&vfu_ctx, sizeof (irq_set), NULL, 0, &irq_set);
-    /* bad start, count range */
-    assert_int_equal(-1, ret);
-    assert_int_equal(EINVAL, errno);
-
-    irq_set.start = 0;
-    irq_set.count = 1;
-    irq_set.index = VFU_DEV_ERR_IRQ;
-
-    ret = handle_device_set_irqs(&vfu_ctx, sizeof (irq_set), NULL, 0, &irq_set);
-    /* bad action for err irq */
-    assert_int_equal(-1, ret);
-    assert_int_equal(EINVAL, errno);
-
-    irq_set.index = VFU_DEV_REQ_IRQ;
-
-    ret = handle_device_set_irqs(&vfu_ctx, sizeof (irq_set), NULL, 0, &irq_set);
-    /* bad action for req irq */
-    assert_int_equal(-1, ret);
-    assert_int_equal(EINVAL, errno);
-
-    irq_set.start = 1;
-    irq_set.count = 0;
-    irq_set.index = VFU_DEV_MSIX_IRQ;
-
-    ret = handle_device_set_irqs(&vfu_ctx, sizeof (irq_set), NULL, 0, &irq_set);
-    /* bad start for count == 0 */
-    assert_int_equal(-1, ret);
-    assert_int_equal(EINVAL, errno);
-
-    irq_set.flags = VFIO_IRQ_SET_ACTION_MASK | VFIO_IRQ_SET_DATA_NONE;
-    irq_set.count = 0;
-    irq_set.start = 0;
-
-    ret = handle_device_set_irqs(&vfu_ctx, sizeof (irq_set), NULL, 0, &irq_set);
-    /* bad action for count == 0 */
-    assert_int_equal(-1, ret);
-    assert_int_equal(EINVAL, errno);
-
-    irq_set.flags = VFIO_IRQ_SET_ACTION_TRIGGER | VFIO_IRQ_SET_DATA_BOOL;
-    irq_set.count = 0;
-    irq_set.start = 0;
-
-    ret = handle_device_set_irqs(&vfu_ctx, sizeof (irq_set), NULL, 0, &irq_set);
-    /* bad action and data type for count == 0 */
-    assert_int_equal(-1, ret);
-    assert_int_equal(EINVAL, errno);
-
-    irq_set.flags = VFIO_IRQ_SET_ACTION_TRIGGER | VFIO_IRQ_SET_DATA_BOOL;
-    irq_set.count = 1;
-    irq_set.start = 0;
-
-    ret = handle_device_set_irqs(&vfu_ctx, sizeof (irq_set), &fd, 1, &irq_set);
-    /* bad fds for DATA_BOOL */
-    assert_int_equal(-1, ret);
-    assert_int_equal(EINVAL, errno);
-
-    irq_set.flags = VFIO_IRQ_SET_ACTION_TRIGGER | VFIO_IRQ_SET_DATA_NONE;
-    irq_set.count = 1;
-    irq_set.start = 0;
-
-    ret = handle_device_set_irqs(&vfu_ctx, sizeof (irq_set), &fd, 1, &irq_set);
-    /* bad fds for DATA_NONE */
-    assert_int_equal(-1, ret);
-    assert_int_equal(EINVAL, errno);
-
-    irq_set.flags = VFIO_IRQ_SET_ACTION_TRIGGER | VFIO_IRQ_SET_DATA_EVENTFD;
-    irq_set.count = 2;
-    irq_set.start = 0;
-
-    ret = handle_device_set_irqs(&vfu_ctx, sizeof (irq_set), &fd, 1, &irq_set);
-    /* bad fds for count == 2 */
-    assert_int_equal(-1, ret);
-    assert_int_equal(EINVAL, errno);
-
-    irqs->err_efd = irqs->req_efd = -1;
-
-    /* Basic disable functionality. */
-
-    irq_set.index = VFU_DEV_REQ_IRQ;
-    irq_set.flags = VFIO_IRQ_SET_ACTION_TRIGGER | VFIO_IRQ_SET_DATA_NONE;
-    irq_set.count = 0;
-    irq_set.start = 0;
+    tran_sock_t ts = { .fd = 23, .conn_fd = 24 };
 
-    ret = handle_device_set_irqs(&vfu_ctx, sizeof (irq_set), NULL, 0, &irq_set);
-    assert_int_equal(0, ret);
+    mkmsg(VFIO_USER_DMA_MAP, NULL, 0);
 
-    irq_set.index = VFU_DEV_REQ_IRQ;
-    irq_set.flags = VFIO_IRQ_SET_ACTION_TRIGGER | VFIO_IRQ_SET_DATA_EVENTFD;
-    irq_set.count = 1;
-    irq_set.start = 0;
+    fds[0] = 0xab;
+    fds[1] = 0xcd;
+    msg.nr_in_fds = 2;
+    msg.in_fds = malloc(sizeof(int) * msg.nr_in_fds);
+    assert_non_null(msg.in_fds);
+    msg.in_fds[0] = fds[0];
+    msg.in_fds[1] = fds[1];
 
-    ret = handle_device_set_irqs(&vfu_ctx, sizeof (irq_set), NULL, 0, &irq_set);
-    assert_int_equal(0, ret);
+    vfu_ctx.tran = &tran_sock_ops;
+    vfu_ctx.tran_data = &ts;
 
-    /* Basic enable. */
+    patch("get_request_header");
+    expect_value(get_request_header, vfu_ctx, &vfu_ctx);
+    expect_check(get_request_header, msgp, check_request_header_msg, NULL);
+    will_return(get_request_header, 0);
 
-    irq_set.index = VFU_DEV_MSIX_IRQ;
-    vfu_ctx.irq_count[VFU_DEV_MSIX_IRQ] = 1;
-    irqs->efds[0] = -1;
-    fd = 0xbeef;
+    patch("exec_command");
+    expect_value(exec_command, vfu_ctx, &vfu_ctx);
+    expect_check(exec_command, msg, check_exec_command_msg, NULL);
+    will_return(exec_command, -1);
+    will_return(exec_command, EREMOTEIO);
 
-    irq_set.index = VFU_DEV_MSIX_IRQ;
-    irq_set.flags = VFIO_IRQ_SET_ACTION_TRIGGER | VFIO_IRQ_SET_DATA_EVENTFD;
-    irq_set.count = 1;
-    irq_set.start = 0;
+    patch("close");
+    expect_value(close, fd, fds[1]);
+    will_return(close, 0);
 
-    ret = handle_device_set_irqs(&vfu_ctx, sizeof (irq_set), &fd, 1, &irq_set);
-    assert_int_equal(0, ret);
-    assert_int_equal(0xbeef, irqs->efds[0]);
+    patch("tran_sock_send_iovec");
+    expect_value(tran_sock_send_iovec, sock, ts.conn_fd);
+    expect_any(tran_sock_send_iovec, msg_id);
+    expect_value(tran_sock_send_iovec, is_reply, true);
+    expect_any(tran_sock_send_iovec, cmd);
+    expect_any(tran_sock_send_iovec, iovecs);
+    expect_any(tran_sock_send_iovec, nr_iovecs);
+    expect_any(tran_sock_send_iovec, fds);
+    expect_any(tran_sock_send_iovec, count);
+    expect_any(tran_sock_send_iovec, err);
+    will_return(tran_sock_send_iovec, 0);
 
+    assert_int_equal(0, process_request(&vfu_ctx));
 }
 
 int
@@ -1944,25 +1949,25 @@ main(void)
         cmocka_unit_test_setup(test_dma_map_without_fd, setup),
         cmocka_unit_test_setup(test_dma_add_regions_mixed, setup),
         cmocka_unit_test_setup(test_dma_add_regions_mixed_partial_failure, setup),
+        cmocka_unit_test_setup(test_dma_map_return_value, setup),
+        cmocka_unit_test_setup(test_handle_dma_unmap, setup),
         cmocka_unit_test_setup(test_dma_controller_add_region_no_fd, setup),
         cmocka_unit_test_setup(test_dma_controller_remove_region_mapped, setup),
         cmocka_unit_test_setup(test_dma_controller_remove_region_unmapped, setup),
-        cmocka_unit_test_setup(test_handle_dma_unmap, setup),
-        cmocka_unit_test_setup(test_process_command_free_passed_fds, setup),
+        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, setup),
         cmocka_unit_test_setup(test_realize_ctx, setup),
         cmocka_unit_test_setup(test_attach_ctx, setup),
         cmocka_unit_test_setup(test_run_ctx, setup),
+        cmocka_unit_test_setup(test_get_region_info, setup),
         cmocka_unit_test_setup(test_vfu_ctx_create, setup),
         cmocka_unit_test_setup(test_pci_caps, setup),
         cmocka_unit_test_setup(test_pci_ext_caps, setup),
         cmocka_unit_test_setup(test_device_get_info, setup),
-        cmocka_unit_test_setup(test_device_get_info_compat, setup),
-        cmocka_unit_test_setup(test_get_region_info, setup),
         cmocka_unit_test_setup(test_setup_sparse_region, setup),
-        cmocka_unit_test_setup(test_dma_map_return_value, setup),
-        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, setup),
+        cmocka_unit_test_setup(test_dirty_pages_without_dma, setup),
+        cmocka_unit_test_setup(test_device_set_irqs, 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,
@@ -1991,9 +1996,7 @@ main(void)
         cmocka_unit_test_setup(test_device_is_stopped_and_copying, setup),
         cmocka_unit_test_setup(test_cmd_allowed_when_stopped_and_copying, setup),
         cmocka_unit_test_setup(test_should_exec_command, setup),
-        cmocka_unit_test_setup(test_exec_command, setup),
-        cmocka_unit_test_setup(test_dirty_pages_without_dma, setup),
-        cmocka_unit_test_setup(test_device_set_irqs, setup),
+        cmocka_unit_test_setup(test_process_request_free_passed_fds, setup),
     };
 
     return cmocka_run_group_tests(tests, NULL, NULL);