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>

4 files changed, 474 insertions, 45 deletions

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)