Prepare python test helpers for receiving commands (#774)

Thus far, the python test code has only ever sent messages of type
commands to the server and processed the corresponding replies. For the
twin-socket feature, the tests will exercise flows where DMA access
commands must be received, processed, and replied to by the client.

This change refactors the message handling python test code to provide
functions to handle server-to-client commands, reusing existing code as
appropriate.

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

1 files changed, 56 insertions, 21 deletions

diff --git a/test/py/libvfio_user.py b/test/py/libvfio_user.py
index f29b50e..1a4cb24 100644
--- a/test/py/libvfio_user.py
+++ b/test/py/libvfio_user.py
@@ -180,8 +180,11 @@ VFIO_USER_DEVICE_RESET = 13
 VFIO_USER_DIRTY_PAGES = 14
 VFIO_USER_MAX = 15
 
+VFIO_USER_F_TYPE = 0xf
 VFIO_USER_F_TYPE_COMMAND = 0
 VFIO_USER_F_TYPE_REPLY = 1
+VFIO_USER_F_NO_REPLY = 0x10
+VFIO_USER_F_ERROR = 0x20
 
 SIZEOF_VFIO_USER_HEADER = 16
 
@@ -714,6 +717,23 @@ def get_reply(sock, expect=0):
     return buf[16:]
 
 
+def send_msg(sock, cmd, msg_type, payload=bytearray(), fds=None, msg_id=None,
+             error_no=0):
+    """
+    Sends a message on the given socket. Can be used on either end of the
+    socket to send commands and replies.
+    """
+    hdr = vfio_user_header(cmd, size=len(payload), msg_type=msg_type,
+                           msg_id=msg_id, error=error_no != 0,
+                           error_no=error_no)
+
+    if fds:
+        sock.sendmsg([hdr + payload], [(socket.SOL_SOCKET, socket.SCM_RIGHTS,
+                                        struct.pack("I" * len(fds), *fds))])
+    else:
+        sock.send(hdr + payload)
+
+
 def msg(ctx, sock, cmd, payload=bytearray(), expect=0, fds=None,
         rsp=True, busy=False):
     """
@@ -726,13 +746,7 @@ def msg(ctx, sock, cmd, payload=bytearray(), expect=0, fds=None,
     response: it can later be retrieved, post vfu_device_quiesced(), with
     get_reply().
     """
-    hdr = vfio_user_header(cmd, size=len(payload))
-
-    if fds:
-        sock.sendmsg([hdr + payload], [(socket.SOL_SOCKET, socket.SCM_RIGHTS,
-                                        struct.pack("I" * len(fds), *fds))])
-    else:
-        sock.send(hdr + payload)
+    send_msg(sock, cmd, VFIO_USER_F_TYPE_COMMAND, payload, fds)
 
     if busy:
         vfu_run_ctx(ctx, errno.EBUSY)
@@ -745,15 +759,17 @@ def msg(ctx, sock, cmd, payload=bytearray(), expect=0, fds=None,
     return get_reply(sock, expect=expect)
 
 
-def get_reply_fds(sock, expect=0):
-    """Receives a message from a socket and pulls the returned file descriptors
-       out of the message."""
+def get_msg_fds(sock, expect_msg_type, expect_errno=0):
+    """
+    Receives a message from a socket and pulls the returned file descriptors
+    out of the message.
+    """
     fds = array.array("i")
-    data, ancillary, flags, addr = sock.recvmsg(4096,
-                                            socket.CMSG_LEN(64 * fds.itemsize))
+    data, ancillary, flags, addr = sock.recvmsg(SERVER_MAX_MSG_SIZE,
+        socket.CMSG_LEN(64 * fds.itemsize))
     (msg_id, cmd, msg_size, msg_flags, errno) = struct.unpack("HHIII",
                                                               data[0:16])
-    assert errno == expect
+    assert errno == expect_errno
 
     cmsg_level, cmsg_type, packed_fd = ancillary[0] if len(ancillary) != 0 \
                                                     else (0, 0, [])
@@ -762,8 +778,18 @@ def get_reply_fds(sock, expect=0):
         [unpacked_fd] = struct.unpack_from("i", packed_fd, offset=i)
         unpacked_fds.append(unpacked_fd)
     assert len(packed_fd)/4 == len(unpacked_fds)
-    assert (msg_flags & VFIO_USER_F_TYPE_REPLY) != 0
-    return (unpacked_fds, data[16:])
+    assert (msg_flags & VFIO_USER_F_TYPE) == expect_msg_type
+    return (unpacked_fds, msg_id, cmd, data[16:])
+
+
+def get_reply_fds(sock, expect=0):
+    """
+    Receives a reply from a socket and returns the included file descriptors
+    and message payload data.
+    """
+    (unpacked_fds, _, _, data) = get_msg_fds(sock, VFIO_USER_F_TYPE_REPLY,
+                                             expect)
+    return (unpacked_fds, data)
 
 
 def msg_fds(ctx, sock, cmd, payload, expect=0, fds=None):
@@ -962,7 +988,7 @@ def prepare_ctx_for_dma(dma_register=__dma_register,
 #
 
 
-msg_id = 1
+next_msg_id = 1
 
 
 @c.CFUNCTYPE(None, c.c_void_p, c.c_int, c.c_char_p)
@@ -978,13 +1004,22 @@ def log(ctx, level, msg):
     print(lvl2str[level] + ": " + msg.decode("utf-8"))
 
 
-def vfio_user_header(cmd, size, no_reply=False, error=False, error_no=0):
-    global msg_id
+def vfio_user_header(cmd, size, msg_type=VFIO_USER_F_TYPE_COMMAND, msg_id=None,
+                     no_reply=False, error=False, error_no=0):
+    global next_msg_id
 
-    buf = struct.pack("HHIII", msg_id, cmd, SIZEOF_VFIO_USER_HEADER + size,
-                      VFIO_USER_F_TYPE_COMMAND, error_no)
+    if msg_id is None:
+        msg_id = next_msg_id
+        next_msg_id += 1
+
+    flags = msg_type
+    if no_reply:
+        flags |= VFIO_USER_F_NO_REPLY
+    if error:
+        flags |= VFIO_USER_F_ERROR
 
-    msg_id += 1
+    buf = struct.pack("HHIII", msg_id, cmd, SIZEOF_VFIO_USER_HEADER + size,
+                      flags, error_no)
 
     return buf