move PCI capability testing to Python (#453)

This also adds a couple of additional tests (as well as more fully testing the
region access path).

Signed-off-by: John Levon <john.levon@nutanix.com>
Reviewed-by: Swapnil Ingle <swapnil.ingle@nutanix.com>

1 files changed, 146 insertions, 19 deletions

diff --git a/test/py/libvfio_user.py b/test/py/libvfio_user.py
index 2d64269..60a93a8 100644
--- a/test/py/libvfio_user.py
+++ b/test/py/libvfio_user.py
@@ -40,7 +40,32 @@ import socket
 import struct
 import syslog
 
-SOCK_PATH = b"/tmp/vfio-user.sock.%d" % os.getpid()
+# from linux/pci_regs.h and linux/pci_defs.h
+
+PCI_HEADER_TYPE_NORMAL = 0
+
+PCI_STD_HEADER_SIZEOF = 64
+
+PCI_BARS_NR = 6
+
+PCI_PM_SIZEOF = 8
+
+PCI_CFG_SPACE_SIZE = 256
+PCI_CFG_SPACE_EXP_SIZE = 4096
+
+PCI_CAP_LIST_NEXT = 1
+
+PCI_CAP_ID_PM = 0x1
+PCI_CAP_ID_VNDR = 0x9
+
+PCI_EXT_CAP_ID_DSN = 0x03
+PCI_EXT_CAP_ID_VNDR = 0x0b
+
+PCI_EXT_CAP_DSN_SIZEOF = 12
+
+PCI_EXT_CAP_VNDR_HDR_SIZEOF = 8
+
+# libvfio-user defines
 
 VFU_TRANS_SOCK = 0
 LIBVFIO_USER_FLAG_ATTACH_NB = (1 << 0)
@@ -94,8 +119,6 @@ VFU_REGION_FLAG_WRITE = 2
 VFU_REGION_FLAG_RW = (VFU_REGION_FLAG_READ | VFU_REGION_FLAG_WRITE)
 VFU_REGION_FLAG_MEM   = 4
 
-PCI_BARS_NR = 6
-
 # enum vfu_dev_irq_type
 VFU_DEV_INTX_IRQ = 0
 VFU_DEV_MSI_IRQ  = 1
@@ -110,13 +133,11 @@ VFU_PCI_TYPE_PCI_X_1      = 1
 VFU_PCI_TYPE_PCI_X_2      = 2
 VFU_PCI_TYPE_EXPRESS      = 3
 
-PCI_HEADER_TYPE_NORMAL = 0
-
-PCI_STD_HEADER_SIZEOF = 64
-
-PCI_CFG_SPACE_SIZE = 256
-PCI_CAP_ID_PM = b'\1'
+VFU_CAP_FLAG_EXTENDED = (1 << 0)
+VFU_CAP_FLAG_CALLBACK = (1 << 1)
+VFU_CAP_FLAG_READONLY = (1 << 2)
 
+SOCK_PATH = b"/tmp/vfio-user.sock.%d" % os.getpid()
 
 topdir = os.path.realpath(os.path.dirname(__file__) + "/../..")
 build_type = os.getenv("BUILD_TYPE", default="dbg")
@@ -130,14 +151,22 @@ lib.vfu_realize_ctx.argtypes = (c.c_void_p,)
 lib.vfu_attach_ctx.argtypes = (c.c_void_p,)
 lib.vfu_run_ctx.argtypes = (c.c_void_p,)
 lib.vfu_destroy_ctx.argtypes = (c.c_void_p,)
-lib.vfu_setup_region.argtypes = (c.c_void_p, c.c_int, c.c_long, c.c_void_p,
-                                 c.c_int, c.c_void_p, c.c_int, c.c_int)
+vfu_region_access_cb_t = c.CFUNCTYPE(c.c_int, c.c_void_p, c.POINTER(c.c_char),
+                                     c.c_ulong, c.c_long, c.c_bool)
+lib.vfu_setup_region.argtypes = (c.c_void_p, c.c_int, c.c_long,
+                                 vfu_region_access_cb_t, c.c_int, c.c_void_p,
+                                 c.c_uint32, c.c_int)
 lib.vfu_pci_get_config_space.argtypes = (c.c_void_p,)
 lib.vfu_pci_get_config_space.restype = (c.c_void_p)
-lib.vfu_setup_device_nr_irqs.argtypes = (c.c_void_p, c.c_int, c.c_int)
+lib.vfu_setup_device_nr_irqs.argtypes = (c.c_void_p, c.c_int, c.c_uint32)
 lib.vfu_pci_init.argtypes = (c.c_void_p, c.c_int, c.c_int, c.c_int)
-lib.vfu_pci_add_capability.argtypes = (c.c_void_p, c.c_long, c.c_int,
+lib.vfu_pci_add_capability.argtypes = (c.c_void_p, c.c_ulong, c.c_int,
                                        c.POINTER(c.c_byte))
+lib.vfu_pci_find_capability.argtypes = (c.c_void_p, c.c_bool, c.c_int)
+lib.vfu_pci_find_capability.restype = (c.c_ulong)
+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)
 
 msg_id = 1
 
@@ -187,25 +216,111 @@ class vfu_pci_hdr_t(c.Structure):
 # Util functions
 #
 
+def to_byte(val):
+    """Cast an int to a byte value."""
+    return val.to_bytes(1, 'little')
+
+def ext_cap_hdr(buf, offset):
+    """Read an extended cap header."""
+
+    # struct pcie_ext_cap_hdr
+    cap_id, cap_next = struct.unpack('HH', buf[offset:offset+4])
+    cap_next >>= 4
+    return cap_id, cap_next
+
+def skip(data, fmt):
+    """Return the data remaining after skipping the given elements."""
+    return data[struct.calcsize(fmt):]
+
+def parse_json(json_str):
+    """Parse JSON into an object with attributes (instead of using a dict)."""
+    return json.loads(json_str, object_hook=lambda d: SimpleNamespace(**d))
+
 def connect_sock():
     sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
     sock.connect(SOCK_PATH)
     return sock
 
+def connect_client(ctx):
+    sock = connect_sock()
+
+    json = b'{ "capabilities": { "max_fds": 8 } }'
+    # struct vfio_user_version
+    payload = struct.pack("HH%dsc" % len(json), LIBVFIO_USER_MAJOR,
+                          LIBVFIO_USER_MINOR, json, b'\0')
+    hdr = vfio_user_header(VFIO_USER_VERSION, size=len(payload))
+    sock.send(hdr + payload)
+    vfu_attach_ctx(ctx, expect=0)
+    payload = get_reply(sock, expect=0)
+    return sock
+
+def disconnect_client(ctx, sock):
+    sock.close()
+
+    # notice client closed connection
+    vfu_run_ctx(ctx)
+
 def get_reply(sock, expect=0):
     buf = sock.recv(4096)
     (msg_id, cmd, msg_size, flags, errno) = struct.unpack("HHIII", buf[0:16])
+    assert (flags & VFIO_USER_F_TYPE_REPLY) != 0
     assert errno == expect
     return buf[16:]
 
-def parse_json(json_str):
-    """Parse JSON into an object with attributes (instead of using a dict)."""
-    return json.loads(json_str, object_hook=lambda d: SimpleNamespace(**d))
-
 def get_pci_header(ctx):
     ptr = lib.vfu_pci_get_config_space(ctx)
     return c.cast(ptr, c.POINTER(vfu_pci_hdr_t)).contents
 
+def get_pci_cfg_space(ctx):
+    ptr = lib.vfu_pci_get_config_space(ctx)
+    return c.cast(ptr, c.POINTER(c.c_char))[0:PCI_CFG_SPACE_SIZE]
+
+def get_pci_ext_cfg_space(ctx):
+    ptr = lib.vfu_pci_get_config_space(ctx)
+    return c.cast(ptr, c.POINTER(c.c_char))[0:PCI_CFG_SPACE_EXP_SIZE]
+
+def read_pci_cfg_space(ctx, buf, count, offset, extended=False):
+    space = get_pci_ext_cfg_space(ctx) if extended else get_pci_cfg_space(ctx)
+
+    for i in range(count):
+        buf[i] = space[offset+i]
+    return count
+
+def write_pci_cfg_space(ctx, buf, count, offset, extended=False):
+    max_offset = PCI_CFG_SPACE_EXP_SIZE if extended else PCI_CFG_SPACE_SIZE
+
+    assert offset + count <= max_offset
+
+    space = c.cast(lib.vfu_pci_get_config_space(ctx), c.POINTER(c.c_char))
+
+    for i in range(count):
+        space[offset+i] = buf[i]
+    return count
+
+def access_region(ctx, sock, is_write, region, offset, count,
+                  data=None, expect=0):
+    # struct vfio_user_region_access
+    payload = struct.pack("QII", offset, region, count)
+    if is_write:
+        payload += data
+
+    cmd = VFIO_USER_REGION_WRITE if is_write else VFIO_USER_REGION_READ
+    hdr = vfio_user_header(cmd, size=len(payload))
+    sock.send(hdr + payload)
+    vfu_run_ctx(ctx)
+    result = get_reply(sock, expect=expect)
+
+    if is_write:
+        return None
+
+    return skip(result, "QII")
+
+def write_region(ctx, sock, region, offset, count, data, expect=0):
+    access_region(ctx, sock, True, region, offset, count, data, expect=expect)
+
+def read_region(ctx, sock, region, offset, count, expect=0):
+    return access_region(ctx, sock, False, region, offset, count, expect=expect)
+
 #
 # Library wrappers
 #
@@ -257,9 +372,11 @@ def vfu_destroy_ctx(ctx):
     if os.path.exists(SOCK_PATH):
         os.remove(SOCK_PATH)
 
-def vfu_setup_region(ctx, index, size, flags=0):
+def vfu_setup_region(ctx, index, size, flags=0, cb=None):
     assert ctx != None
-    ret = lib.vfu_setup_region(ctx, index, size, None, flags, None, 0, -1)
+    ret = lib.vfu_setup_region(ctx, index, size,
+                               c.cast(cb, vfu_region_access_cb_t),
+                               flags, None, 0, -1)
     return ret
 
 def vfu_setup_device_nr_irqs(ctx, irqtype, count):
@@ -276,3 +393,13 @@ def vfu_pci_add_capability(ctx, pos, flags, data):
 
     databuf = (c.c_byte * len(data)).from_buffer(bytearray(data))
     return lib.vfu_pci_add_capability(ctx, pos, flags, databuf)
+
+def vfu_pci_find_capability(ctx, extended, cap_id):
+    assert ctx != None
+
+    return lib.vfu_pci_find_capability(ctx, extended, cap_id)
+
+def vfu_pci_find_next_capability(ctx, extended, offset, cap_id):
+    assert ctx != None
+
+    return lib.vfu_pci_find_next_capability(ctx, extended, offset, cap_id)