support VFIO_DMA_UNMAP_FLAG_GET_DIRTY_BITMAP (#521)

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

3 files changed, 80 insertions, 5 deletions

diff --git a/include/vfio-user.h b/include/vfio-user.h
index 3994abc..fa3aba8 100644
--- a/include/vfio-user.h
+++ b/include/vfio-user.h
@@ -119,6 +119,9 @@ struct vfio_user_dma_region {
     uint64_t                size;
     uint64_t                offset;
     uint32_t                prot;
+#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;
     struct vfio_user_bitmap bitmap[];
diff --git a/lib/libvfio-user.c b/lib/libvfio-user.c
index df216bf..307f909 100644
--- a/lib/libvfio-user.c
+++ b/lib/libvfio-user.c
@@ -524,25 +524,52 @@ handle_dma_unmap(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg,
                  struct vfio_user_dma_region *region)
 {
     int ret;
+    char *bitmap;
     char rstr[1024];
 
     assert(vfu_ctx != NULL);
     assert(msg != NULL);
     assert(region != NULL);
 
-    if (msg->in_size != sizeof(*region)) {
+    if (msg->in_size < sizeof(*region)) {
         vfu_log(vfu_ctx, LOG_ERR, "bad size of DMA unmap region %zu",
                 msg->in_size);
         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 prot=%#x flags=%#x",
+             region->addr, region->addr + region->size,
+             region->offset, region->prot, region->flags);
 
     vfu_log(vfu_ctx, LOG_DEBUG, "removing DMA region %s", rstr);
 
-    if (region->flags != 0) {
+    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);
+            return ERROR_INT(EINVAL);
+        }
+        /*
+         * TODO this could be a separate function, but the implementation is
+         * temporary anyway since we're moving dirty page tracking out of
+         * 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,
+                                            &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);
         return ERROR_INT(ENOTSUP);
     }
diff --git a/test/unit-tests.c b/test/unit-tests.c
index 4f39228..7525bb4 100644
--- a/test/unit-tests.c
+++ b/test/unit-tests.c
@@ -240,6 +240,50 @@ test_handle_dma_unmap(void **state UNUSED)
 }
 
 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;
+
+    vfu_ctx.dma->nregions = 1;
+    vfu_ctx.dma->regions[0].info.iova.iov_base = (void *)0x0;
+    vfu_ctx.dma->regions[0].info.iova.iov_len = 0x1000;
+    vfu_ctx.dma->regions[0].fd = -1;
+
+    /*
+     * TODO Hack to avoid mocking dma_controller_dirty_page_get since we're
+     * moving testing to Python.
+     */
+    vfu_ctx.dma->dirty_pgsize = 0x1000;
+    vfu_ctx.dma->regions[0].dirty_bitmap = (void *)&bitmap;
+
+    vfu_ctx.dma_unregister = mock_dma_unregister;
+
+    expect_value(mock_dma_unregister, vfu_ctx, &vfu_ctx);
+    expect_check(mock_dma_unregister, info, check_dma_info,
+                 &vfu_ctx.dma->regions[0].info);
+    will_return(mock_dma_unregister, 0);
+
+    ret = handle_dma_unmap(&vfu_ctx,
+                           mkmsg(VFIO_USER_DMA_UNMAP, &r, size),
+                           r);
+
+    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);
+    free(msg.out_data);
+}
+
+
+static void
 test_dma_controller_add_region_no_fd(void **state UNUSED)
 {
     vfu_dma_addr_t dma_addr = (void *)0xdeadbeef;
@@ -708,6 +752,7 @@ main(void)
         cmocka_unit_test_setup(test_dma_map_without_fd, 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_handle_dma_unmap_dirty, 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),