fixes for VFIO_USER_DIRTY_PAGES (#537)

- we should only accept one range, not multiple ones
- clearly define and implement argsz behaviour
- we need to check if migration is configured
- add proper test coverage; move existing testing to python

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

6 files changed, 474 insertions, 123 deletions

diff --git a/test/mocks.c b/test/mocks.c
index 4f202d5..6e34d4a 100644
--- a/test/mocks.c
+++ b/test/mocks.c
@@ -62,11 +62,9 @@ static struct function funcs[] = {
     { .name = "dma_controller_add_region" },
     { .name = "dma_controller_remove_region" },
     { .name = "dma_controller_unmap_region" },
-    { .name = "handle_dirty_pages" },
     { .name = "process_request" },
     { .name = "should_exec_command" },
     { .name = "migration_region_access_registers" },
-    { .name = "handle_dirty_pages_get" },
     { .name = "handle_device_state" },
     {. name = "vfio_migr_state_transition_is_valid" },
     { .name = "state_trans_notify" },
@@ -214,18 +212,6 @@ should_exec_command(vfu_ctx_t *vfu_ctx, uint16_t cmd)
     return mock();
 }
 
-int
-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, msg);
-    }
-    check_expected(vfu_ctx);
-    check_expected(msg);
-    errno = mock();
-    return mock();
-}
-
 ssize_t
 migration_region_access_registers(vfu_ctx_t *vfu_ctx, char *buf, size_t count,
                                   loff_t pos, bool is_write)
@@ -243,23 +229,6 @@ migration_region_access_registers(vfu_ctx_t *vfu_ctx, char *buf, size_t count,
     return mock();
 }
 
-int
-handle_dirty_pages_get(vfu_ctx_t *vfu_ctx,
-                       struct iovec **iovecs, size_t *nr_iovecs,
-                       struct vfio_user_bitmap_range *ranges, uint32_t size)
-{
-    if (!is_patched("handle_dirty_pages_get")) {
-        return __real_handle_dirty_pages_get(vfu_ctx, iovecs, nr_iovecs,
-                                             ranges, size);
-    }
-    check_expected(vfu_ctx);
-    check_expected(iovecs);
-    check_expected(nr_iovecs);
-    check_expected(ranges);
-    check_expected(size);
-    return mock();
-}
-
 ssize_t
 handle_device_state(vfu_ctx_t *vfu_ctx, struct migration *migr,
                     uint32_t device_state, bool notify) {
diff --git a/test/py/libvfio_user.py b/test/py/libvfio_user.py
index fb201dd..a664da0 100644
--- a/test/py/libvfio_user.py
+++ b/test/py/libvfio_user.py
@@ -35,6 +35,7 @@ from collections import namedtuple
 from types import SimpleNamespace
 import ctypes as c
 import json
+import mmap
 import os
 import pathlib
 import socket
@@ -146,6 +147,16 @@ VFU_REGION_FLAG_WRITE = 2
 VFU_REGION_FLAG_RW = (VFU_REGION_FLAG_READ | VFU_REGION_FLAG_WRITE)
 VFU_REGION_FLAG_MEM   = 4
 
+VFIO_USER_F_DMA_REGION_READ = (1 << 0)
+VFIO_USER_F_DMA_REGION_WRITE = (1 << 1)
+VFIO_USER_F_DMA_REGION_MAPPABLE = (1 << 2)
+
+VFIO_DMA_UNMAP_FLAG_GET_DIRTY_BITMAP = (1 << 0)
+
+VFIO_IOMMU_DIRTY_PAGES_FLAG_START = (1 << 0)
+VFIO_IOMMU_DIRTY_PAGES_FLAG_STOP = (1 << 1)
+VFIO_IOMMU_DIRTY_PAGES_FLAG_GET_BITMAP = (1 << 2)
+
 # enum vfu_dev_irq_type
 VFU_DEV_INTX_IRQ = 0
 VFU_DEV_MSI_IRQ  = 1
@@ -169,6 +180,8 @@ VFU_CAP_FLAG_EXTENDED = (1 << 0)
 VFU_CAP_FLAG_CALLBACK = (1 << 1)
 VFU_CAP_FLAG_READONLY = (1 << 2)
 
+VFU_MIGR_CALLBACKS_VERS = 1
+
 SOCK_PATH = b"/tmp/vfio-user.sock.%d" % os.getpid()
 
 topdir = os.path.realpath(os.path.dirname(__file__) + "/../..")
@@ -195,8 +208,8 @@ class Structure(c.Structure):
 class vfu_bar_t(c.Union):
     _pack_ = 1
     _fields_ = [
-        ("mem", c.c_int),
-        ("io", c.c_int)
+        ("mem", c.c_int32),
+        ("io", c.c_int32)
     ]
 
 class vfu_pci_hdr_intr_t(Structure):
@@ -209,9 +222,9 @@ class vfu_pci_hdr_intr_t(Structure):
 class vfu_pci_hdr_t(Structure):
     _pack_ = 1
     _fields_ = [
-        ("id", c.c_int),
-        ("cmd", c.c_short),
-        ("sts", c.c_short),
+        ("id", c.c_int32),
+        ("cmd", c.c_uint16),
+        ("sts", c.c_uint16),
         ("rid", c.c_byte),
         ("cc_pi", c.c_byte),
         ("cc_scc", c.c_byte),
@@ -221,9 +234,9 @@ class vfu_pci_hdr_t(Structure):
         ("htype", c.c_byte),
         ("bist", c.c_byte),
         ("bars", vfu_bar_t * PCI_BARS_NR),
-        ("ccptr", c.c_int),
-        ("ss", c.c_int),
-        ("erom", c.c_int),
+        ("ccptr", c.c_int32),
+        ("ss", c.c_int32),
+        ("erom", c.c_int32),
         ("cap", c.c_byte),
         ("res1", c.c_byte * 7),
         ("intr", vfu_pci_hdr_intr_t),
@@ -234,66 +247,143 @@ class vfu_pci_hdr_t(Structure):
 class iovec_t(Structure):
     _fields_ = [
         ("iov_base", c.c_void_p),
-        ("iov_len", c.c_int)
+        ("iov_len", c.c_int32)
     ]
 
 class vfio_irq_info(Structure):
+    _pack_ = 1
     _fields_ = [
-        ("argsz", c.c_uint),
-        ("flags", c.c_uint),
-        ("index", c.c_uint),
-        ("count", c.c_uint),
+        ("argsz", c.c_uint32),
+        ("flags", c.c_uint32),
+        ("index", c.c_uint32),
+        ("count", c.c_uint32),
     ]
 
 class vfio_irq_set(Structure):
+    _pack_ = 1
     _fields_ = [
-        ("argsz", c.c_uint),
-        ("flags", c.c_uint),
-        ("index", c.c_uint),
-        ("start", c.c_uint),
-        ("count", c.c_uint),
+        ("argsz", c.c_uint32),
+        ("flags", c.c_uint32),
+        ("index", c.c_uint32),
+        ("start", c.c_uint32),
+        ("count", c.c_uint32),
     ]
 
 class vfio_user_device_info(Structure):
+    _pack_ = 1
     _fields_ = [
-        ("argsz", c.c_uint),
-        ("flags", c.c_uint),
-        ("num_regions", c.c_uint),
-        ("num_irqs", c.c_uint),
+        ("argsz", c.c_uint32),
+        ("flags", c.c_uint32),
+        ("num_regions", c.c_uint32),
+        ("num_irqs", c.c_uint32),
     ]
 
 class vfio_region_info(Structure):
+    _pack_ = 1
     _fields_ = [
-        ("argsz", c.c_uint),
-        ("flags", c.c_uint),
-        ("index", c.c_uint),
-        ("cap_offset", c.c_uint),
-        ("size", c.c_ulong),
-        ("offset", c.c_ulong),
+        ("argsz", c.c_uint32),
+        ("flags", c.c_uint32),
+        ("index", c.c_uint32),
+        ("cap_offset", c.c_uint32),
+        ("size", c.c_uint64),
+        ("offset", c.c_uint64),
     ]
 
 class vfio_region_info_cap_type(Structure):
+    _pack_ = 1
     _fields_ = [
-        ("id", c.c_ushort),
-        ("version", c.c_ushort),
-        ("next", c.c_uint),
-        ("type", c.c_uint),
-        ("subtype", c.c_uint),
+        ("id", c.c_uint16),
+        ("version", c.c_uint16),
+        ("next", c.c_uint32),
+        ("type", c.c_uint32),
+        ("subtype", c.c_uint32),
     ]
 
 class vfio_region_info_cap_sparse_mmap(Structure):
+    _pack_ = 1
     _fields_ = [
-        ("id", c.c_ushort),
-        ("version", c.c_ushort),
-        ("next", c.c_uint),
-        ("nr_areas", c.c_uint),
-        ("reserved", c.c_uint),
+        ("id", c.c_uint16),
+        ("version", c.c_uint16),
+        ("next", c.c_uint32),
+        ("nr_areas", c.c_uint32),
+        ("reserved", c.c_uint32),
     ]
 
 class vfio_region_sparse_mmap_area(Structure):
+    _pack_ = 1
+    _fields_ = [
+        ("offset", c.c_uint64),
+        ("size", c.c_uint64),
+    ]
+
+class vfio_user_dma_map(Structure):
+    _pack_ = 1
+    _fields_ = [
+        ("argsz", c.c_uint32),
+        ("flags", c.c_uint32),
+        ("offset", c.c_uint64),
+        ("addr", c.c_uint64),
+        ("size", c.c_uint64),
+    ]
+
+class vfu_dma_info_t(Structure):
+    _fields_ = [
+        ("iova", iovec_t),
+        ("vaddr", c.c_void_p),
+        ("mapping", iovec_t),
+        ("page_size", c.c_size_t),
+        ("prot", c.c_uint32)
+    ]
+
+class vfio_user_dirty_pages(Structure):
+    _pack_ = 1
+    _fields_ = [
+        ("argsz", c.c_uint32),
+        ("flags", c.c_uint32)
+    ]
+
+class vfio_user_bitmap(Structure):
+    _pack_ = 1
     _fields_ = [
-        ("offset", c.c_ulong),
-        ("size", c.c_ulong),
+        ("pgsize", c.c_uint64),
+        ("size", c.c_uint64)
+    ]
+
+class vfio_user_bitmap_range(Structure):
+    _pack_ = 1
+    _fields_ = [
+        ("iova", c.c_uint64),
+        ("size", c.c_uint64),
+        ("bitmap", vfio_user_bitmap)
+    ]
+
+transition_cb_t = c.CFUNCTYPE(c.c_int, c.c_void_p, c.c_int)
+get_pending_bytes_cb_t = c.CFUNCTYPE(c.c_uint64, c.c_void_p)
+prepare_data_cb_t = c.CFUNCTYPE(c.c_void_p, c.POINTER(c.c_uint64),
+                                c.POINTER(c.c_uint64))
+read_data_cb_t = c.CFUNCTYPE(c.c_ssize_t, c.c_void_p, c.c_void_p,
+                             c.c_uint64, c.c_uint64)
+write_data_cb_t = c.CFUNCTYPE(c.c_ssize_t, c.c_void_p, c.c_uint64)
+data_written_cb_t = c.CFUNCTYPE(c.c_int, c.c_void_p, c.c_uint64)
+
+class vfu_migration_callbacks_t(Structure):
+    _fields_ = [
+        ("version", c.c_int),
+        ("transition", transition_cb_t),
+        ("get_pending_bytes", get_pending_bytes_cb_t),
+        ("prepare_data", prepare_data_cb_t),
+        ("read_data", read_data_cb_t),
+        ("write_data", write_data_cb_t),
+        ("data_written", data_written_cb_t),
+    ]
+
+class dma_sg_t(Structure):
+    _fields_ = [
+        ("dma_addr", c.c_void_p),
+        ("region", c.c_int),
+        ("length", c.c_uint64),
+        ("offset", c.c_uint64),
+        ("mappable", c.c_bool)
     ]
 
 #
@@ -327,7 +417,15 @@ lib.vfu_pci_find_next_capability.argtypes = (c.c_void_p, c.c_bool, c.c_ulong,
                                              c.c_int)
 lib.vfu_pci_find_next_capability.restype = (c.c_ulong)
 lib.vfu_irq_trigger.argtypes = (c.c_void_p, c.c_uint)
-
+vfu_dma_register_cb_t = c.CFUNCTYPE(None, c.c_void_p, c.POINTER(vfu_dma_info_t))
+vfu_dma_unregister_cb_t = c.CFUNCTYPE(c.c_int, c.c_void_p,
+                                      c.POINTER(vfu_dma_info_t))
+lib.vfu_setup_device_dma.argtypes = (c.c_void_p, vfu_dma_register_cb_t,
+                                     vfu_dma_unregister_cb_t)
+lib.vfu_setup_device_migration_callbacks.argtypes = (c.c_void_p,
+    c.POINTER(vfu_migration_callbacks_t), c.c_uint64)
+lib.vfu_addr_to_sg.argtypes = (c.c_void_p, c.c_void_p, c.c_size_t,
+                               c.POINTER(dma_sg_t), c.c_int, c.c_int)
 
 def to_byte(val):
     """Cast an int to a byte value."""
@@ -553,3 +651,38 @@ def vfu_irq_trigger(ctx, subindex):
     assert ctx != None
 
     return lib.vfu_irq_trigger(ctx, subindex)
+
+def vfu_setup_device_dma(ctx, register_cb=None, unregister_cb=None):
+    assert ctx != None
+
+    return lib.vfu_setup_device_dma(ctx, c.cast(register_cb,
+                                                vfu_dma_register_cb_t),
+                                         c.cast(unregister_cb,
+                                                vfu_dma_unregister_cb_t))
+
+def vfu_setup_device_migration_callbacks(ctx, cbs=None, offset=0):
+    assert ctx != None
+
+    @c.CFUNCTYPE(c.c_int)
+    def stub():
+        return 0
+
+    if not cbs:
+        cbs = vfu_migration_callbacks_t()
+        cbs.version = VFU_MIGR_CALLBACKS_VERS
+        cbs.transition = c.cast(stub, transition_cb_t)
+        cbs.get_pending_bytes = c.cast(stub, get_pending_bytes_cb_t)
+        cbs.prepare_data = c.cast(stub, prepare_data_cb_t)
+        cbs.read_data = c.cast(stub, read_data_cb_t)
+        cbs.write_data = c.cast(stub, write_data_cb_t)
+        cbs.data_written = c.cast(stub, data_written_cb_t)
+
+    return lib.vfu_setup_device_migration_callbacks(ctx, cbs, offset)
+
+def vfu_addr_to_sg(ctx, dma_addr, length, max_sg=1,
+                   prot=(mmap.PROT_READ | mmap.PROT_WRITE)):
+    assert ctx != None
+
+    sg = dma_sg_t()
+
+    return lib.vfu_addr_to_sg(ctx, dma_addr, length, sg, max_sg, prot)
diff --git a/test/py/test_dirty_pages.py b/test/py/test_dirty_pages.py
new file mode 100644
index 0000000..5d4f6db
--- /dev/null
+++ b/test/py/test_dirty_pages.py
@@ -0,0 +1,298 @@
+#
+# Copyright (c) 2021 Nutanix Inc. All rights reserved.
+#
+# Authors: John Levon <john.levon@nutanix.com>
+#
+#  Redistribution and use in source and binary forms, with or without
+#  modification, are permitted provided that the following conditions are met:
+#      * Redistributions of source code must retain the above copyright
+#        notice, this list of conditions and the following disclaimer.
+#      * Redistributions in binary form must reproduce the above copyright
+#        notice, this list of conditions and the following disclaimer in the
+#        documentation and/or other materials provided with the distribution.
+#      * Neither the name of Nutanix nor the names of its contributors may be
+#        used to endorse or promote products derived from this software without
+#        specific prior written permission.
+#
+#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+#  ARE DISCLAIMED. IN NO EVENT SHALL <COPYRIGHT HOLDER> BE LIABLE FOR ANY
+#  DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+#  (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+#  SERVICESLOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+#  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
+#  LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
+#  OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH
+#  DAMAGE.
+#
+
+from libvfio_user import *
+import errno
+import mmap
+import tempfile
+
+ctx = None
+
+@vfu_dma_register_cb_t
+def dma_register(ctx, info):
+    pass
+
+@vfu_dma_unregister_cb_t
+def dma_unregister(ctx, info):
+    pass
+    return 0
+
+def test_dirty_pages_setup():
+    global ctx, sock
+
+    ctx = vfu_create_ctx(flags=LIBVFIO_USER_FLAG_ATTACH_NB)
+    assert ctx != None
+
+    ret = vfu_pci_init(ctx)
+    assert ret == 0
+
+    ret = vfu_setup_device_dma(ctx, dma_register, dma_unregister)
+    assert ret == 0
+
+    f = tempfile.TemporaryFile()
+    f.truncate(0x2000)
+
+    mmap_areas = [ (0x1000, 0x1000) ]
+
+    ret = vfu_setup_region(ctx, index=VFU_PCI_DEV_MIGR_REGION_IDX, size=0x2000,
+                           flags=VFU_REGION_FLAG_RW, mmap_areas=mmap_areas,
+                           fd=f.fileno())
+    assert ret == 0
+
+    ret = vfu_realize_ctx(ctx)
+    assert ret == 0
+
+    sock = connect_client(ctx)
+
+    f = tempfile.TemporaryFile()
+    f.truncate(0x10000)
+
+    payload = vfio_user_dma_map(argsz=len(vfio_user_dma_map()),
+        flags=(VFIO_USER_F_DMA_REGION_READ |
+               VFIO_USER_F_DMA_REGION_WRITE |
+               VFIO_USER_F_DMA_REGION_MAPPABLE),
+        offset=0, addr=0x10000, size=0x10000)
+
+    hdr = vfio_user_header(VFIO_USER_DMA_MAP, size=len(payload))
+
+    sock.sendmsg([hdr + payload], [(socket.SOL_SOCKET, socket.SCM_RIGHTS,
+                 struct.pack("I", f.fileno()))])
+    vfu_run_ctx(ctx)
+    get_reply(sock)
+
+def test_dirty_pages_short_write():
+    payload = struct.pack("I", 8)
+
+    hdr = vfio_user_header(VFIO_USER_DIRTY_PAGES, size=len(payload))
+    sock.send(hdr + payload)
+    vfu_run_ctx(ctx)
+    get_reply(sock, expect=errno.EINVAL)
+
+def test_dirty_pages_bad_argsz():
+    payload = vfio_user_dirty_pages(argsz=4,
+        flags=VFIO_IOMMU_DIRTY_PAGES_FLAG_START)
+
+    hdr = vfio_user_header(VFIO_USER_DIRTY_PAGES, size=len(payload))
+    sock.send(hdr + payload)
+    vfu_run_ctx(ctx)
+    get_reply(sock, expect=errno.EINVAL)
+
+def test_dirty_pages_start_no_migration():
+    payload = vfio_user_dirty_pages(argsz=len(vfio_user_dirty_pages()),
+        flags=VFIO_IOMMU_DIRTY_PAGES_FLAG_START)
+
+    hdr = vfio_user_header(VFIO_USER_DIRTY_PAGES, size=len(payload))
+    sock.send(hdr + payload)
+    vfu_run_ctx(ctx)
+    get_reply(sock, expect=errno.ENOTSUP)
+
+def test_dirty_pages_start_bad_flags():
+    #
+    # This is a little cheeky, after vfu_realize_ctx(), but it works at the
+    # moment.
+    #
+    vfu_setup_device_migration_callbacks(ctx, offset=0x1000)
+
+    payload = vfio_user_dirty_pages(argsz=len(vfio_user_dirty_pages()),
+        flags=(VFIO_IOMMU_DIRTY_PAGES_FLAG_START |
+               VFIO_IOMMU_DIRTY_PAGES_FLAG_STOP))
+
+    hdr = vfio_user_header(VFIO_USER_DIRTY_PAGES, size=len(payload))
+    sock.send(hdr + payload)
+    vfu_run_ctx(ctx)
+    get_reply(sock, expect=errno.EINVAL)
+
+    payload = vfio_user_dirty_pages(argsz=len(vfio_user_dirty_pages()),
+        flags=(VFIO_IOMMU_DIRTY_PAGES_FLAG_START |
+               VFIO_IOMMU_DIRTY_PAGES_FLAG_GET_BITMAP))
+
+    hdr = vfio_user_header(VFIO_USER_DIRTY_PAGES, size=len(payload))
+    sock.send(hdr + payload)
+    vfu_run_ctx(ctx)
+    get_reply(sock, expect=errno.EINVAL)
+
+def test_dirty_pages_start():
+    payload = vfio_user_dirty_pages(argsz=len(vfio_user_dirty_pages()),
+                                    flags=VFIO_IOMMU_DIRTY_PAGES_FLAG_START)
+
+    hdr = vfio_user_header(VFIO_USER_DIRTY_PAGES, size=len(payload))
+    sock.send(hdr + payload)
+    vfu_run_ctx(ctx)
+    get_reply(sock)
+
+    # should be idempotent
+    sock.send(hdr + payload)
+    vfu_run_ctx(ctx)
+    get_reply(sock)
+
+def test_dirty_pages_get_short_read():
+    payload = vfio_user_dirty_pages(argsz=len(vfio_user_dirty_pages()),
+        flags=VFIO_IOMMU_DIRTY_PAGES_FLAG_GET_BITMAP)
+
+    hdr = vfio_user_header(VFIO_USER_DIRTY_PAGES, size=len(payload))
+    sock.send(hdr + payload)
+    vfu_run_ctx(ctx)
+    get_reply(sock, expect=errno.EINVAL)
+
+#
+# This should in fact work; update when it does.
+#
+def test_dirty_pages_get_sub_range():
+    dirty_pages = vfio_user_dirty_pages(argsz=len(vfio_user_dirty_pages()),
+        flags=VFIO_IOMMU_DIRTY_PAGES_FLAG_GET_BITMAP)
+    bitmap = vfio_user_bitmap(pgsize=0x1000, size=1)
+    br = vfio_user_bitmap_range(iova=0x11000, size=0x1000, bitmap=bitmap)
+
+    hdr = vfio_user_header(VFIO_USER_DIRTY_PAGES,
+                           size=len(dirty_pages) + len(br))
+    sock.send(hdr + dirty_pages + br)
+    vfu_run_ctx(ctx)
+    get_reply(sock, expect=errno.ENOTSUP)
+
+def test_dirty_pages_get_bad_page_size():
+    dirty_pages = vfio_user_dirty_pages(argsz=len(vfio_user_dirty_pages()),
+        flags=VFIO_IOMMU_DIRTY_PAGES_FLAG_GET_BITMAP)
+    bitmap = vfio_user_bitmap(pgsize=0x2000, size=8)
+    br = vfio_user_bitmap_range(iova=0x10000, size=0x10000, bitmap=bitmap)
+
+    hdr = vfio_user_header(VFIO_USER_DIRTY_PAGES,
+                           size=len(dirty_pages) + len(br))
+    sock.send(hdr + dirty_pages + br)
+    vfu_run_ctx(ctx)
+    get_reply(sock, expect=errno.EINVAL)
+
+def test_dirty_pages_get_bad_bitmap_size():
+    dirty_pages = vfio_user_dirty_pages(argsz=len(vfio_user_dirty_pages()),
+        flags=VFIO_IOMMU_DIRTY_PAGES_FLAG_GET_BITMAP)
+    bitmap = vfio_user_bitmap(pgsize=0x1000, size=1)
+    br = vfio_user_bitmap_range(iova=0x10000, size=0x10000, bitmap=bitmap)
+
+    hdr = vfio_user_header(VFIO_USER_DIRTY_PAGES,
+                           size=len(dirty_pages) + len(br))
+    sock.send(hdr + dirty_pages + br)
+    vfu_run_ctx(ctx)
+    get_reply(sock, expect=errno.EINVAL)
+
+def test_dirty_pages_get_short_reply():
+    dirty_pages = vfio_user_dirty_pages(argsz=len(vfio_user_dirty_pages()),
+        flags=VFIO_IOMMU_DIRTY_PAGES_FLAG_GET_BITMAP)
+    bitmap = vfio_user_bitmap(pgsize=0x1000, size=8)
+    br = vfio_user_bitmap_range(iova=0x10000, size=0x10000, bitmap=bitmap)
+
+    hdr = vfio_user_header(VFIO_USER_DIRTY_PAGES,
+                           size=len(dirty_pages) + len(br))
+    sock.send(hdr + dirty_pages + br)
+    vfu_run_ctx(ctx)
+    result = get_reply(sock)
+
+    assert len(result) == len(vfio_user_dirty_pages())
+
+    dirty_pages, _ = vfio_user_dirty_pages.pop_from_buffer(result)
+
+    argsz = len(vfio_user_dirty_pages()) + len(vfio_user_bitmap_range()) + 8
+
+    assert dirty_pages.argsz == argsz
+    assert dirty_pages.flags == VFIO_IOMMU_DIRTY_PAGES_FLAG_GET_BITMAP
+
+def test_dirty_pages_get_unmodified():
+    argsz = len(vfio_user_dirty_pages()) + len(vfio_user_bitmap_range()) + 8
+
+    dirty_pages = vfio_user_dirty_pages(argsz=argsz,
+        flags=VFIO_IOMMU_DIRTY_PAGES_FLAG_GET_BITMAP)
+    bitmap = vfio_user_bitmap(pgsize=0x1000, size=8)
+    br = vfio_user_bitmap_range(iova=0x10000, size=0x10000, bitmap=bitmap)
+
+    hdr = vfio_user_header(VFIO_USER_DIRTY_PAGES,
+                           size=len(dirty_pages) + len(br))
+    sock.send(hdr + dirty_pages + br)
+    vfu_run_ctx(ctx)
+    result = get_reply(sock)
+
+    assert len(result) == argsz
+
+    dirty_pages, result = vfio_user_dirty_pages.pop_from_buffer(result)
+
+    assert dirty_pages.argsz == argsz
+    assert dirty_pages.flags == VFIO_IOMMU_DIRTY_PAGES_FLAG_GET_BITMAP
+
+    br, result = vfio_user_bitmap_range.pop_from_buffer(result)
+
+    assert br.iova == 0x10000
+    assert br.size == 0x10000
+
+    assert br.bitmap.pgsize == 0x1000
+    assert br.bitmap.size == 8
+
+def test_dirty_pages_get_modified():
+    # sufficient to mark the region dirty
+    ret = vfu_addr_to_sg(ctx, dma_addr=0x10000, length=0x1000)
+    assert ret == 1
+
+    ret = vfu_addr_to_sg(ctx, dma_addr=0x14000, length=0x4000)
+    assert ret == 1
+
+    argsz = len(vfio_user_dirty_pages()) + len(vfio_user_bitmap_range()) + 8
+
+    dirty_pages = vfio_user_dirty_pages(argsz=argsz,
+        flags=VFIO_IOMMU_DIRTY_PAGES_FLAG_GET_BITMAP)
+    bitmap = vfio_user_bitmap(pgsize=0x1000, size=8)
+    br = vfio_user_bitmap_range(iova=0x10000, size=0x10000, bitmap=bitmap)
+
+    hdr = vfio_user_header(VFIO_USER_DIRTY_PAGES,
+                           size=len(dirty_pages) + len(br))
+    sock.send(hdr + dirty_pages + br)
+    vfu_run_ctx(ctx)
+    result = get_reply(sock)
+
+    dirty_pages, result = vfio_user_dirty_pages.pop_from_buffer(result)
+    br, result = vfio_user_bitmap_range.pop_from_buffer(result)
+    bitmap = struct.unpack("Q", result)[0]
+
+    assert bitmap == 0b11110001
+
+def test_dirty_pages_stop():
+    payload = vfio_user_dirty_pages(argsz=len(vfio_user_dirty_pages()),
+                                    flags=VFIO_IOMMU_DIRTY_PAGES_FLAG_STOP)
+
+    hdr = vfio_user_header(VFIO_USER_DIRTY_PAGES, size=len(payload))
+    sock.send(hdr + payload)
+    vfu_run_ctx(ctx)
+    get_reply(sock)
+
+    payload = vfio_user_dirty_pages(argsz=len(vfio_user_dirty_pages()),
+                                    flags=VFIO_IOMMU_DIRTY_PAGES_FLAG_STOP)
+
+    hdr = vfio_user_header(VFIO_USER_DIRTY_PAGES, size=len(payload))
+    sock.send(hdr + payload)
+    vfu_run_ctx(ctx)
+    get_reply(sock)
+
+def test_dirty_pages_cleanup():
+    disconnect_client(ctx, sock)
+    vfu_destroy_ctx(ctx)
diff --git a/test/py/test_pci_caps.py b/test/py/test_pci_caps.py
index 88a9b7c..be1914d 100644
--- a/test/py/test_pci_caps.py
+++ b/test/py/test_pci_caps.py
@@ -70,8 +70,7 @@ def test_pci_cap_bad_pos():
     assert pos == -1
     assert c.get_errno() == errno.EINVAL
 
-@c.CFUNCTYPE(c.c_int, c.c_void_p, c.POINTER(c.c_char),
-             c.c_long, c.c_long, c.c_int)
+@vfu_region_access_cb_t
 def pci_region_cb(ctx, buf, count, offset, is_write):
     if not is_write:
         return read_pci_cfg_space(ctx, buf, count, offset)
diff --git a/test/py/test_pci_ext_caps.py b/test/py/test_pci_ext_caps.py
index 70f8253..ab10e12 100644
--- a/test/py/test_pci_ext_caps.py
+++ b/test/py/test_pci_ext_caps.py
@@ -91,8 +91,7 @@ def test_pci_ext_cap_bad_pos():
     assert pos == -1
     assert c.get_errno() == errno.EINVAL
 
-@c.CFUNCTYPE(c.c_int, c.c_void_p, c.POINTER(c.c_char),
-             c.c_long, c.c_long, c.c_int)
+@vfu_region_access_cb_t
 def pci_region_cb(ctx, buf, count, offset, is_write):
     if not is_write:
         return read_pci_cfg_space(ctx, buf, count, offset, extended=True)
diff --git a/test/unit-tests.c b/test/unit-tests.c
index c083d05..d334244 100644
--- a/test/unit-tests.c
+++ b/test/unit-tests.c
@@ -448,33 +448,6 @@ typedef struct {
 } tran_sock_t;
 
 static void
-test_dirty_pages_without_dma(UNUSED void **state)
-{
-    int ret;
-
-    /* with DMA controller */
-
-    patch("handle_dirty_pages");
-
-    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);
-
-    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_migration_state_transitions(void **state UNUSED)
 {
     bool (*f)(uint32_t, uint32_t) = vfio_migr_state_transition_is_valid;
@@ -736,24 +709,6 @@ test_should_exec_command(UNUSED void **state)
     assert_true(should_exec_command(&vfu_ctx, 0xbeef));
 }
 
-static void
-test_dma_controller_dirty_page_get(void **state UNUSED)
-{
-    dma_memory_region_t *r;
-    uint64_t len = UINT32_MAX + (uint64_t)10;
-    char bp[0x20008]; /* must be QWORD aligned */
-
-    vfu_ctx.dma->nregions = 1;
-    r = &vfu_ctx.dma->regions[0];
-    r->info.iova.iov_base = (void *)0;
-    r->info.iova.iov_len = len;
-    r->info.vaddr = (void *)0xdeadbeef;
-    vfu_ctx.dma->dirty_pgsize = 4096;
-
-    assert_int_equal(0, dma_controller_dirty_page_get(vfu_ctx.dma, (void *)0,
-                     len, 4096, sizeof(bp), (char **)&bp));
-}
-
 int
 main(void)
 {
@@ -769,7 +724,6 @@ main(void)
         cmocka_unit_test_setup(test_dma_map_sg, setup),
         cmocka_unit_test_setup(test_dma_addr_to_sg, setup),
         cmocka_unit_test_setup(test_vfu_setup_device_dma, setup),
-        cmocka_unit_test_setup(test_dirty_pages_without_dma, setup),
         cmocka_unit_test_setup(test_migration_state_transitions, setup),
         cmocka_unit_test_setup_teardown(test_setup_migration_region_size_ok,
             setup_test_setup_migration_region,
@@ -789,7 +743,6 @@ 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_dma_controller_dirty_page_get, setup),
     };
 
     return cmocka_run_group_tests(tests, NULL, NULL);