restore argsz for DMA map/unmap (#523)

use DMA map/unmap format similar to VFIO's

Using a DMA map/unmap format similar to VFIO's (vfio_iommu_type1_dma_map / vfio_iommu_type1_dma_unmap) makes it easier to adapt to future changes. Consequently we also honor the passed argsz.

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

6 files changed, 226 insertions, 172 deletions

diff --git a/docs/vfio-user.rst b/docs/vfio-user.rst
index b5825d8..57d689b 100644
--- a/docs/vfio-user.rst
+++ b/docs/vfio-user.rst
@@ -531,40 +531,49 @@ or if the client uses a vIOMMU.
 Request
 ^^^^^^^
 
-The request payload for this message is an entry of the following format:
+The request payload for this message is a structure of the following format:
 
 +-------------+--------+-------------+
 | Name        | Offset | Size        |
 +=============+========+=============+
-| Address     | 0      | 8           |
+| argsz       | 0      | 4           |
 +-------------+--------+-------------+
-| Size        | 8      | 8           |
-+-------------+--------+-------------+
-| Offset      | 16     | 8           |
-+-------------+--------+-------------+
-| Protections | 24     | 4           |
-+-------------+--------+-------------+
-| Flags       | 28     | 4           |
+| flags       | 4      | 4           |
 +-------------+--------+-------------+
 |             | +-----+------------+ |
 |             | | Bit | Definition | |
 |             | +=====+============+ |
-|             | | 0   | Mappable   | |
+|             | | 0   | readable   | |
 |             | +-----+------------+ |
+|             | | 1   | writeable  | |
+|             | +-----+------------+ |
+|             | | 2   | mappable   | |
+|             | +-----+------------+ |
++-------------+--------+-------------+
+| offset      | 8      | 8           |
++-------------+--------+-------------+
+| address     | 16     | 8           |
++-------------+--------+-------------+
+| size        | 24     | 8           |
 +-------------+--------+-------------+
 
-* *Address* is the base DMA address of the region.
-* *Size* is the size of the region.
-* *Offset* is the file offset of the region with respect to the associated file
-  descriptor, or zero if the region is not mappable
-* *Protections* are the region's protection attributes as encoded in
-  ``<sys/mman.h>``; only ``PROT_READ`` and ``PROT_WRITE`` are supported
-* *Flags* contains the following region attributes:
+* *argsz* is the size of the above structure. Note there is no reply payload,
+  so this field differs from other message types.
+* *flags* contains the following region attributes:
 
-  * *Mappable* indicates that the region can be mapped via the mmap() system
+  * *readable* indicates that the region can be read from.
+
+  * *writeable* indicates that the region can be written to.
+
+  * *mappable* indicates that the region can be mapped via the mmap() system
     call using the file descriptor provided in the message meta-data.
 
-This structure is 32 bytes in size, so the message size is 16 + 32 == 48 bytes.
+* *offset* is the file offset of the region with respect to the associated file
+  descriptor, or zero if the region is not mappable
+* *address* is the base DMA address of the region.
+* *size* is the size of the region.
+
+This structure is 32 bytes in size, so the message size is 16 + 32 bytes.
 
 If the DMA region being added can be directly mapped by the server, a file
 descriptor must be sent as part of the message meta-data. On ``AF_UNIX``
@@ -592,44 +601,45 @@ described by the following structure:
 Request
 ^^^^^^^
 
-The request payload for this message is an entry of the following format:
-
-+--------------+--------+---------------------------------------+
-| Name         | Offset | Size                                  |
-+==============+========+=======================================+
-| Address      | 0      | 8                                     |
-+--------------+--------+---------------------------------------+
-| Size         | 8      | 8                                     |
-+--------------+--------+---------------------------------------+
-| Offset       | 16     | 8                                     |
-+--------------+--------+---------------------------------------+
-| Protections  | 24     | 4                                     |
-+--------------+--------+---------------------------------------+
-| Flags        | 28     | 4                                     |
-+--------------+--------+---------------------------------------+
-|              | +-----+--------------------------------------+ |
-|              | | Bit | Definition                           | |
-|              | +=====+======================================+ |
-|              | | 0   | VFIO_DMA_UNMAP_FLAG_GET_DIRTY_BITMAP | |
-|              | +-----+--------------------------------------+ |
-+--------------+--------+---------------------------------------+
-
-* *Address* is the base DMA address of the DMA region.
-* *Size* is the size of the DMA region.
-* *Offset* must be zero (ignored)
-* *Protections* must be zero (ignored)
-* *Flags* contains the following DMA region attributes:
-
-  * ``VFIO_DMA_UNMAP_FLAG_GET_DIRTY_BITMAP`` indicates that a dirty page bitmap
-    must be populated before unmapping the DMA region. The client must provide
-    a ``struct vfio_user_bitmap`` immediately following this table entry.
-
-If ``VFIO_DMA_UNMAP_FLAG_GET_DIRTY_BITMAP`` is set in Flags, the size of the
-total request message is: 16 + 32 + 16.
-Otherwise, if ``VFIO_DMA_UNMAP_FLAG_GET_DIRTY_BITMAP`` is not set in Flags, the
-size of the total request message is: 16 + 32.
-
-.. _VFIO bitmap format:
+The request payload for this message is a structure of the following format:
+
++--------------+--------+------------------------+
+| Name         | Offset | Size                   |
++==============+========+========================+
+| argsz        | 0      | 4                      |
++--------------+--------+------------------------+
+| flags        | 4      | 4                      |
++--------------+--------+------------------------+
+|              | +-----+-----------------------+ |
+|              | | Bit | Definition            | |
+|              | +=====+=======================+ |
+|              | | 0   | get dirty page bitmap | |
+|              | +-----+-----------------------+ |
++--------------+--------+------------------------+
+| address      | 8      | 8                      |
++--------------+--------+------------------------+
+| size         | 16     | 8                      |
++--------------+--------+------------------------+
+
+* *argsz* is the maximum size of the reply payload.
+* *flags* contains the following DMA region attributes:
+
+  * *get dirty page bitmap* indicates that a dirty page bitmap must be
+    populated before unmapping the DMA region. The client must provide a
+    `VFIO bitmap`_ structure, explained below, immediately following this
+    entry.
+
+* *address* is the base DMA address of the DMA region.
+* *size* is the size of the DMA region.
+
+The size of request message depends on whether or not the
+*get dirty page bitmap* bit is set in Flags:
+
+* If not set, the size of the total request message is: 16 + 24.
+
+* If set, the size of the total request message is: 16 + 24 + 16.
+
+.. _VFIO bitmap:
 
 VFIO bitmap format
 """"""""""""""""""
@@ -653,18 +663,14 @@ mapped then the server must release all references to that DMA region before
 replying, which potentially includes in-flight DMA transactions. Removing a
 portion of a DMA region is possible.
 
-The payload depends on whether or not ``VFIO_DMA_UNMAP_FLAG_GET_DIRTY_BITMAP``
-is set in Flags in the request:
-
-* ``VFIO_DMA_UNMAP_FLAG_GET_DIRTY_BITMAP`` is set in Flags in the request:
-  The server responds with a payload that contains the original
-  ``struct vfio_user_bitmap``  sent in the request, followed by the
-  corresponding dirty page bitmap, where each bit represents one page of size
-  ``struct vfio_user_bitmap.pgsize``.  The size of the total reply message is:
-  16 + 16 + size of dirty page bitmap.
+The server responds with the original DMA entry in the request. If the
+*get dirty page bitmap* bit is set in flags in the request, then
+the server also includes the `VFIO bitmap`_ structure sent in the request,
+followed by the corresponding dirty page bitmap, where each bit represents
+one page of size *pgsize* in `VFIO bitmap`_ .
 
-* ``VFIO_DMA_UNMAP_FLAG_GET_DIRTY_BITMAP`` is not set in Flags in the request:
-   There is no payload. The size of the reply message is: 16.
+The total size of the total reply message is:
+16 + 24 + (16 + *size* in `VFIO bitmap`_ if *get dirty page bitmap* is set).
 
 ``VFIO_USER_DEVICE_GET_INFO``
 -----------------------------
@@ -1542,8 +1548,7 @@ VFIO Dirty Bitmap Get Format
 
 * *size* is the size of the IOVA region
 
-* *bitmap* is the VFIO bitmap (``struct vfio_bitmap``). This field is explained
-  in `VFIO bitmap format`_.
+* *bitmap* is the VFIO bitmap explained in `VFIO bitmap`_.
 
 Reply
 ^^^^^
diff --git a/include/vfio-user.h b/include/vfio-user.h
index fa3aba8..ef3159a 100644
--- a/include/vfio-user.h
+++ b/include/vfio-user.h
@@ -114,18 +114,29 @@ struct vfio_user_bitmap {
     char data[];
 } __attribute__((packed));
 
-struct vfio_user_dma_region {
-    uint64_t                addr;
-    uint64_t                size;
-    uint64_t                offset;
-    uint32_t                prot;
+/* based on struct vfio_iommu_type1_dma_map */
+struct vfio_user_dma_map {
+    uint32_t argsz;
+#define VFIO_USER_F_DMA_REGION_READ     (1 << 0)
+#define VFIO_USER_F_DMA_REGION_WRITE    (1 << 1)
+#define VFIO_USER_F_DMA_REGION_MAPPABLE (1 << 2)
+    uint32_t flags;
+    uint64_t offset;
+    uint64_t addr;
+    uint64_t size;
+} __attribute__((packed));
+
+/* based on struct vfio_iommu_type1_dma_unmap */
+struct vfio_user_dma_unmap {
+    uint32_t argsz;
 #ifndef VFIO_DMA_UNMAP_FLAG_GET_DIRTY_BITMAP
 #define VFIO_DMA_UNMAP_FLAG_GET_DIRTY_BITMAP (1 << 0)
 #endif
-#define VFIO_USER_F_DMA_REGION_MAPPABLE (1 << 0)
-    uint32_t                flags;
+    uint32_t flags;
+    uint64_t addr;
+    uint64_t size;
     struct vfio_user_bitmap bitmap[];
-} __attribute__((packed));
+};
 
 struct vfio_user_region_access {
     uint64_t    offset;
diff --git a/lib/libvfio-user.c b/lib/libvfio-user.c
index 307f909..99a115d 100644
--- a/lib/libvfio-user.c
+++ b/lib/libvfio-user.c
@@ -468,30 +468,40 @@ consume_fd(int *fds, size_t nr_fds, size_t index)
 
 int
 handle_dma_map(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg,
-               struct vfio_user_dma_region *region)
+               struct vfio_user_dma_map *dma_map)
 {
     char rstr[1024];
     int fd = -1;
     int ret;
+    uint32_t prot = 0;
 
     assert(vfu_ctx != NULL);
     assert(msg != NULL);
-    assert(region != NULL);
+    assert(dma_map != NULL);
 
-    if (msg->in_size != sizeof(*region)) {
-        vfu_log(vfu_ctx, LOG_ERR, "bad size of DMA map region %zu",
-                msg->in_size);
+    if (msg->in_size < sizeof(*dma_map) || dma_map->argsz < sizeof(*dma_map)) {
+        vfu_log(vfu_ctx, LOG_ERR, "bad DMA map region size=%zu argsz=%u",
+                msg->in_size, dma_map->argsz);
         return ERROR_INT(EINVAL);
     }
 
-    snprintf(rstr, sizeof(rstr),
-             "[%#lx, %#lx) offset=%#lx prot=%#x flags=%#x",
-             region->addr, region->addr + region->size, region->offset,
-             region->prot, region->flags);
+    snprintf(rstr, sizeof(rstr), "[%#lx, %#lx) offset=%#lx flags=%#x",
+             dma_map->addr, dma_map->addr + dma_map->size, dma_map->offset,
+             dma_map->flags);
 
     vfu_log(vfu_ctx, LOG_DEBUG, "adding DMA region %s", rstr);
 
-    if (region->flags == VFIO_USER_F_DMA_REGION_MAPPABLE) {
+    if (dma_map->flags & VFIO_USER_F_DMA_REGION_READ) {
+        prot |= PROT_READ;
+        dma_map->flags &= ~VFIO_USER_F_DMA_REGION_READ;
+    }
+
+    if (dma_map->flags & VFIO_USER_F_DMA_REGION_WRITE) {
+        prot |= PROT_WRITE;
+        dma_map->flags &= ~VFIO_USER_F_DMA_REGION_WRITE;
+    }
+
+    if (dma_map->flags == VFIO_USER_F_DMA_REGION_MAPPABLE) {
         fd = consume_fd(msg->in_fds, msg->nr_in_fds, 0);
         if (fd < 0) {
             vfu_log(vfu_ctx, LOG_ERR,
@@ -499,11 +509,13 @@ handle_dma_map(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg,
                     rstr);
             return -1;
         }
+    } else if (dma_map->flags != 0) {
+        vfu_log(vfu_ctx, LOG_ERR, "bad flags=%#x", dma_map->flags);
     }
 
-    ret = dma_controller_add_region(vfu_ctx->dma, (void *)region->addr,
-                                    region->size, fd, region->offset,
-                                    region->prot);
+    ret = dma_controller_add_region(vfu_ctx->dma, (void *)dma_map->addr,
+                                    dma_map->size, fd, dma_map->offset,
+                                    prot);
     if (ret < 0) {
         ret = errno;
         vfu_log(vfu_ctx, LOG_ERR, "failed to add DMA region %s: %m", rstr);
@@ -521,31 +533,41 @@ handle_dma_map(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg,
 
 int
 handle_dma_unmap(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg,
-                 struct vfio_user_dma_region *region)
+                 struct vfio_user_dma_unmap *dma_unmap)
 {
     int ret;
-    char *bitmap;
+    char *bitmap = NULL;
     char rstr[1024];
 
     assert(vfu_ctx != NULL);
     assert(msg != NULL);
-    assert(region != NULL);
+    assert(dma_unmap != NULL);
 
-    if (msg->in_size < sizeof(*region)) {
-        vfu_log(vfu_ctx, LOG_ERR, "bad size of DMA unmap region %zu",
-                msg->in_size);
+    if (msg->in_size < sizeof(*dma_unmap) || dma_unmap->argsz < sizeof(*dma_unmap)) {
+        vfu_log(vfu_ctx, LOG_ERR, "bad DMA unmap region size=%zu argsz=%u",
+                msg->in_size, dma_unmap->argsz);
         return ERROR_INT(EINVAL);
     }
 
-    snprintf(rstr, sizeof(rstr), "[%#lx, %#lx) offset=%#lx prot=%#x flags=%#x",
-             region->addr, region->addr + region->size,
-             region->offset, region->prot, region->flags);
+    snprintf(rstr, sizeof(rstr), "[%#lx, %#lx) flags=%#x",
+             dma_unmap->addr, dma_unmap->addr + dma_unmap->size, dma_unmap->flags);
 
     vfu_log(vfu_ctx, LOG_DEBUG, "removing DMA region %s", rstr);
 
-    if (region->flags == VFIO_DMA_UNMAP_FLAG_GET_DIRTY_BITMAP) {
-        if (msg->in_size < sizeof(*region) + sizeof(struct vfio_user_bitmap)) {
-            vfu_log(vfu_ctx, LOG_ERR, "bad message size %#lx", msg->in_size);
+    msg->out_size = sizeof(*dma_unmap);
+
+    if (dma_unmap->flags == VFIO_DMA_UNMAP_FLAG_GET_DIRTY_BITMAP) {
+        if (msg->in_size < sizeof(*dma_unmap) + sizeof(*dma_unmap->bitmap)
+            || dma_unmap->argsz < sizeof(*dma_unmap) + sizeof(*dma_unmap->bitmap) + dma_unmap->bitmap->size) {
+            vfu_log(vfu_ctx, LOG_ERR, "bad message size=%#lx argsz=%#x",
+                    msg->in_size, dma_unmap->argsz);
+
+            /*
+             * Ideally we should set argsz in the reply and fail the request
+             * with a struct vfio_user_dma_unmap payload, however this isn't
+             * currently supported. Instead, we simply fail the request,
+             * that's what VFIO does anyway.
+             */
             return ERROR_INT(EINVAL);
         }
         /*
@@ -554,28 +576,35 @@ handle_dma_unmap(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg,
          * the DMA controller.
          */
         ret = dma_controller_dirty_page_get(vfu_ctx->dma,
-                                            (vfu_dma_addr_t)region->addr,
-                                            region->size,
-                                            region->bitmap->pgsize,
-                                            region->bitmap->size,
+                                            (vfu_dma_addr_t)dma_unmap->addr,
+                                            dma_unmap->size,
+                                            dma_unmap->bitmap->pgsize,
+                                            dma_unmap->bitmap->size,
                                             &bitmap);
         if (ret < 0) {
             vfu_log(vfu_ctx, LOG_ERR, "failed to get dirty page bitmap: %m");
             return -1;
         }
-        msg->out_data = malloc(region->bitmap->size);
-        if (msg->out_data == NULL) {
-            return ERROR_INT(ENOMEM);
-        }
-        memcpy(msg->out_data, bitmap, region->bitmap->size);
-        msg->out_size = region->bitmap->size;
-    } else if (region->flags != 0) {
-        vfu_log(vfu_ctx, LOG_ERR, "bad flags=%#x", region->flags);
+        msg->out_size += sizeof(*dma_unmap->bitmap) + dma_unmap->bitmap->size;
+    } else if (dma_unmap->flags != 0) {
+        vfu_log(vfu_ctx, LOG_ERR, "bad flags=%#x", dma_unmap->flags);
         return ERROR_INT(ENOTSUP);
     }
+
+    msg->out_data = malloc(msg->out_size);
+    if (msg->out_data == NULL) {
+        return ERROR_INT(ENOMEM);
+    }
+    memcpy(msg->out_data, dma_unmap, sizeof(*dma_unmap));
+
+    if (dma_unmap->flags & VFIO_DMA_UNMAP_FLAG_GET_DIRTY_BITMAP) {
+        memcpy(msg->out_data + sizeof(*dma_unmap), dma_unmap->bitmap, sizeof(*dma_unmap->bitmap));
+        memcpy(msg->out_data + sizeof(*dma_unmap) + sizeof(*dma_unmap->bitmap), bitmap, dma_unmap->bitmap->size);
+    }
+
     ret = dma_controller_remove_region(vfu_ctx->dma,
-                                       (void *)region->addr,
-                                       region->size,
+                                       (void *)dma_unmap->addr,
+                                       dma_unmap->size,
                                        vfu_ctx->dma_unregister,
                                        vfu_ctx);
     if (ret < 0) {
diff --git a/lib/private.h b/lib/private.h
index 2698e4c..4d9e562 100644
--- a/lib/private.h
+++ b/lib/private.h
@@ -181,11 +181,11 @@ exec_command(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg);
 
 int
 handle_dma_map(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg,
-               struct vfio_user_dma_region *dma_regions);
+               struct vfio_user_dma_map *dma_map);
 
 int
 handle_dma_unmap(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg,
-                 struct vfio_user_dma_region *dma_regions);
+                 struct vfio_user_dma_unmap *dma_unmap);
 
 int
 handle_device_get_region_info(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg);
diff --git a/samples/client.c b/samples/client.c
index 0de23cf..82cb12b 100644
--- a/samples/client.c
+++ b/samples/client.c
@@ -540,7 +540,7 @@ wait_for_irq(int irq_fd)
 }
 
 static void
-handle_dma_write(int sock, struct vfio_user_dma_region *dma_regions,
+handle_dma_write(int sock, struct vfio_user_dma_map *dma_regions,
                  int nr_dma_regions, int *dma_region_fds)
 {
     struct vfio_user_dma_region_access dma_access;
@@ -589,7 +589,7 @@ handle_dma_write(int sock, struct vfio_user_dma_region *dma_regions,
 }
 
 static void
-handle_dma_read(int sock, struct vfio_user_dma_region *dma_regions,
+handle_dma_read(int sock, struct vfio_user_dma_map *dma_regions,
                 int nr_dma_regions, int *dma_region_fds)
 {
     struct vfio_user_dma_region_access dma_access, *response;
@@ -632,7 +632,7 @@ handle_dma_read(int sock, struct vfio_user_dma_region *dma_regions,
 }
 
 static void
-handle_dma_io(int sock, struct vfio_user_dma_region *dma_regions,
+handle_dma_io(int sock, struct vfio_user_dma_map *dma_regions,
               int nr_dma_regions, int *dma_region_fds)
 {
     handle_dma_write(sock, dma_regions, nr_dma_regions, dma_region_fds);
@@ -640,7 +640,7 @@ handle_dma_io(int sock, struct vfio_user_dma_region *dma_regions,
 }
 
 static void
-get_dirty_bitmaps(int sock, struct vfio_user_dma_region *dma_regions,
+get_dirty_bitmaps(int sock, struct vfio_user_dma_map *dma_regions,
                   UNUSED int nr_dma_regions)
 {
     struct vfio_iommu_type1_dirty_bitmap dirty_bitmap = { 0 };
@@ -1039,7 +1039,7 @@ migrate_to(char *old_sock_path, int *server_max_fds,
 }
 
 static void
-map_dma_regions(int sock, struct vfio_user_dma_region *dma_regions,
+map_dma_regions(int sock, struct vfio_user_dma_map *dma_regions,
                 int *dma_region_fds, int nr_dma_regions)
 {
     int i, ret;
@@ -1065,7 +1065,7 @@ map_dma_regions(int sock, struct vfio_user_dma_region *dma_regions,
 int main(int argc, char *argv[])
 {
 	int ret, sock, irq_fd;
-    struct vfio_user_dma_region *dma_regions;
+    struct vfio_user_dma_map *dma_regions;
     struct vfio_user_device_info client_dev_info = {0};
     int *dma_region_fds;
     int i;
@@ -1154,11 +1154,11 @@ int main(int argc, char *argv[])
     dma_region_fds = alloca(sizeof(*dma_region_fds) * nr_dma_regions);
 
     for (i = 0; i < nr_dma_regions; i++) {
+        dma_regions[i].argsz = sizeof(struct vfio_user_dma_map);
         dma_regions[i].addr = i * sysconf(_SC_PAGESIZE);
         dma_regions[i].size = sysconf(_SC_PAGESIZE);
         dma_regions[i].offset = dma_regions[i].addr;
-        dma_regions[i].prot = PROT_READ | PROT_WRITE;
-        dma_regions[i].flags = VFIO_USER_F_DMA_REGION_MAPPABLE;
+        dma_regions[i].flags = VFIO_USER_F_DMA_REGION_READ | VFIO_USER_F_DMA_REGION_WRITE | VFIO_USER_F_DMA_REGION_MAPPABLE;
         dma_region_fds[i] = fileno(fp);
     }
 
@@ -1212,19 +1212,18 @@ int main(int argc, char *argv[])
      *
      * unmap the first group of the DMA regions
      */
-    {
-        struct vfio_user_dma_region r[server_max_fds];
-        memcpy(r, dma_regions, sizeof(r));
-        for (i = 0; i < (int)ARRAY_SIZE(r); i++) {
-            r[i].flags = 0;
-            ret = tran_sock_msg(sock, 7, VFIO_USER_DMA_UNMAP, &r[i],
-                                sizeof(struct vfio_user_dma_region),
-                                NULL, NULL, 0);
+    for (i = 0; i < server_max_fds; i++) {
+        struct vfio_user_dma_unmap r = {
+            .argsz = sizeof(r),
+            .addr = dma_regions[i].addr,
+            .size = dma_regions[i].size
+        };
+        ret = tran_sock_msg(sock, 7, VFIO_USER_DMA_UNMAP, &r, sizeof(r),
+                            NULL, &r, sizeof(r));
+        if (ret < 0) {
+            err(EXIT_FAILURE, "failed to unmap DMA region");
         }
     }
-    if (ret < 0) {
-        err(EXIT_FAILURE, "failed to unmap DMA regions");
-    }
 
     /*
      * Schedule an interrupt in 3 seconds from now in the old server and then
diff --git a/test/unit-tests.c b/test/unit-tests.c
index 7525bb4..c083d05 100644
--- a/test/unit-tests.c
+++ b/test/unit-tests.c
@@ -121,13 +121,14 @@ setup(void **state UNUSED)
 static void
 test_dma_map_mappable_without_fd(void **state UNUSED)
 {
-    struct vfio_user_dma_region dma_region = {
+    struct vfio_user_dma_map dma_map = {
+        .argsz = sizeof(dma_map),
         .flags = VFIO_USER_F_DMA_REGION_MAPPABLE
     };
 
     ret = handle_dma_map(&vfu_ctx,
-                         mkmsg(VFIO_USER_DMA_MAP, &dma_region, sizeof(dma_region)),
-                         &dma_region);
+                         mkmsg(VFIO_USER_DMA_MAP, &dma_map, sizeof(dma_map)),
+                         &dma_map);
     assert_int_equal(-1, ret);
     assert_int_equal(errno, EINVAL);
 }
@@ -135,23 +136,26 @@ test_dma_map_mappable_without_fd(void **state UNUSED)
 static void
 test_dma_map_without_fd(void **state UNUSED)
 {
-    struct vfio_user_dma_region r = {
+    struct vfio_user_dma_map dma_map = {
+        .argsz = sizeof(dma_map),
         .addr = 0xdeadbeef,
         .size = 0xcafebabe,
         .offset = 0x8badf00d,
-        .prot = PROT_NONE
+        .flags = 0
     };
 
     patch("dma_controller_add_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.addr);
-    expect_value(dma_controller_add_region, size, r.size);
+    expect_value(dma_controller_add_region, dma_addr, dma_map.addr);
+    expect_value(dma_controller_add_region, size, dma_map.size);
     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);
-    ret = handle_dma_map(&vfu_ctx, mkmsg(VFIO_USER_DMA_MAP, &r, sizeof(r)), &r);
+    expect_value(dma_controller_add_region, offset, dma_map.offset);
+    expect_value(dma_controller_add_region, prot, PROT_NONE);
+    ret = handle_dma_map(&vfu_ctx,
+                         mkmsg(VFIO_USER_DMA_MAP, &dma_map, sizeof(dma_map)),
+                         &dma_map);
     assert_int_equal(0, ret);
 }
 
@@ -181,22 +185,24 @@ test_dma_map_return_value(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 = { 0 };
+    struct vfio_user_dma_map dma_map = {
+        .argsz = sizeof(dma_map)
+    };
 
     patch("dma_controller_add_region");
     expect_value(dma_controller_add_region, dma, vfu_ctx.dma);
-    expect_value(dma_controller_add_region, dma_addr, r.addr);
-    expect_value(dma_controller_add_region, size, r.size);
+    expect_value(dma_controller_add_region, dma_addr, dma_map.addr);
+    expect_value(dma_controller_add_region, size, dma_map.size);
     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);
+    expect_value(dma_controller_add_region, offset, dma_map.offset);
+    expect_value(dma_controller_add_region, prot, PROT_NONE);
     will_return(dma_controller_add_region, 0);
     will_return(dma_controller_add_region, 2);
 
     assert_int_equal(0, handle_dma_map(&vfu_ctx,
-                                       mkmsg(VFIO_USER_DMA_MAP, &r,
-                                             sizeof(r)),
-                                       &r));
+                                       mkmsg(VFIO_USER_DMA_MAP, &dma_map,
+                                             sizeof(dma_map)),
+                                       &dma_map));
 }
 
 /*
@@ -205,8 +211,10 @@ test_dma_map_return_value(void **state UNUSED)
 static void
 test_handle_dma_unmap(void **state UNUSED)
 {
-    struct vfio_user_dma_region r = {
-        .addr = 0x1000, .size = 0x1000
+    struct vfio_user_dma_unmap dma_unmap = {
+        .argsz = sizeof(dma_unmap),
+        .addr = 0x1000,
+        .size = 0x1000
     };
 
     vfu_ctx.dma->nregions = 3;
@@ -228,8 +236,9 @@ test_handle_dma_unmap(void **state UNUSED)
     will_return(mock_dma_unregister, 0);
 
     ret = handle_dma_unmap(&vfu_ctx,
-                           mkmsg(VFIO_USER_DMA_UNMAP, &r, sizeof(r)),
-                           &r);
+                           mkmsg(VFIO_USER_DMA_UNMAP, &dma_unmap,
+                                 sizeof(dma_unmap)),
+                           &dma_unmap);
 
     assert_int_equal(0, ret);
     assert_int_equal(2, vfu_ctx.dma->nregions);
@@ -237,20 +246,21 @@ test_handle_dma_unmap(void **state UNUSED)
     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);
+    free(msg.out_data);
 }
 
 static void
 test_handle_dma_unmap_dirty(void **state UNUSED)
 {
     uint64_t bitmap = 0xdeadbeef;
-    size_t size = sizeof(struct vfio_user_dma_region) + sizeof(struct vfio_user_bitmap);
-    struct vfio_user_dma_region *r = alloca(size);
-    r->addr= 0x0;
-    r->size = 0x1000;
-    r->offset = r->prot = 0; /* silence valgrind */
-    r->flags = VFIO_DMA_UNMAP_FLAG_GET_DIRTY_BITMAP;
-    r->bitmap->pgsize = 0x1000;
-    r->bitmap->size = 8;
+    size_t size = sizeof(struct vfio_user_dma_unmap) + sizeof(struct vfio_user_bitmap);
+    struct vfio_user_dma_unmap *dma_unmap = alloca(size);
+    dma_unmap->argsz = size + sizeof(bitmap);
+    dma_unmap->addr = 0x0;
+    dma_unmap->size = 0x1000;
+    dma_unmap->flags = VFIO_DMA_UNMAP_FLAG_GET_DIRTY_BITMAP;
+    dma_unmap->bitmap->pgsize = 0x1000;
+    dma_unmap->bitmap->size = sizeof(bitmap);
 
     vfu_ctx.dma->nregions = 1;
     vfu_ctx.dma->regions[0].info.iova.iov_base = (void *)0x0;
@@ -272,13 +282,13 @@ test_handle_dma_unmap_dirty(void **state UNUSED)
     will_return(mock_dma_unregister, 0);
 
     ret = handle_dma_unmap(&vfu_ctx,
-                           mkmsg(VFIO_USER_DMA_UNMAP, &r, size),
-                           r);
+                           mkmsg(VFIO_USER_DMA_UNMAP, &dma_unmap, size),
+                           dma_unmap);
 
     assert_int_equal(0, ret);
     assert_int_equal(0, vfu_ctx.dma->nregions);
-    assert_int_equal(sizeof(uint64_t), msg.out_size);
-    assert_int_equal(0xdeadbeef, *(uint64_t *)msg.out_data);
+    assert_int_equal(size + sizeof(bitmap), msg.out_size);
+    assert_int_equal(0xdeadbeef, *(uint64_t *)(msg.out_data + size));
     free(msg.out_data);
 }