From 2e8ec2e17a5252d29bae849eb4ccd7ca6bab216a Mon Sep 17 00:00:00 2001
From: Mattias Nissler <122288598+mnissler-rivos@users.noreply.github.com>
Date: Thu, 31 Aug 2023 11:23:52 +0200
Subject: [PATCH] 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>
---
 test/py/libvfio_user.py | 77 ++++++++++++++++++++++++++++++-----------
 1 file changed, 56 insertions(+), 21 deletions(-)

diff --git a/test/py/libvfio_user.py b/test/py/libvfio_user.py
index f29b50ea..1a4cb240 100644
--- a/test/py/libvfio_user.py
+++ b/test/py/libvfio_user.py
@@ -180,8 +180,11 @@ def is_32bit():
 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