pyverbs: Update MW rkey after type 2 bind

set_bind_wr() assigns the incremented rkey to the bind WR but leaves MW.rkey
unchanged. After the bind succeeds, applications that publish MW.rkey to the
peer send the old rkey, while the NIC has already bound the MW with the new
rkey. Remote access using the stale rkey consequently fails with a remote
access error.

Synchronize MW.rkey with the rkey of successfully posted IBV_WR_BIND_MW
requests. When posting a WR list partially fails, only Bind WRs before bad_wr
are updated, since those are the requests accepted by ibv_post_send().

Signed-off-by: Zhiwei Zhang <zhangzhiwei@bitintelligence.io>
Signed-off-by: Leon Romanovsky <leonro@nvidia.com>

pyverbs/qp.pyx

diff --git a/pyverbs/qp.pyx b/pyverbs/qp.pyx
index a14fd93..abf8635 100644
--- a/pyverbs/qp.pyx
+++ b/pyverbs/qp.pyx
@@ -33,6 +33,19 @@
     unsigned long htobe32(unsigned long host_32bits)
 
 
+cdef void update_mw_rkeys(v.ibv_send_wr *wr, v.ibv_send_wr *bad_wr):
+    """
+    Update MW rkeys for successfully posted bind WRs.
+
+    ibv_post_send() returns the first WR that was not posted in bad_wr, so
+    only the WRs preceding it were accepted when posting a list fails.
+    """
+    while wr != bad_wr:
+        if wr.opcode == e.IBV_WR_BIND_MW:
+            wr.bind_mw.mw.rkey = wr.bind_mw.rkey
+        wr = wr.next
+
+
 cdef class QPCap(PyverbsObject):
     def __init__(self, max_send_wr=1, max_recv_wr=10, max_send_sge=1,
                  max_recv_sge=1, max_inline_data=0):
@@ -1286,8 +1299,9 @@
         """
         # In order to provide a pointer to a pointer, use a temporary cdef'ed
         # variable.
-        cdef v.ibv_send_wr *my_bad_wr
+        cdef v.ibv_send_wr *my_bad_wr = NULL
         rc = v.ibv_post_send(self.qp, &wr.send_wr, &my_bad_wr)
+        update_mw_rkeys(&wr.send_wr, my_bad_wr)
         if rc != 0:
             if (bad_wr):
                 memcpy(&bad_wr.send_wr, my_bad_wr, sizeof(bad_wr.send_wr))

tests/test_mr.py

diff --git a/tests/test_mr.py b/tests/test_mr.py
index ef75f4b..b7ef04e 100644
--- a/tests/test_mr.py
+++ b/tests/test_mr.py
@@ -317,9 +317,11 @@
         # Poll the bind MW WR.
         u.poll_cq(self.server.cq)
         u.poll_cq(self.client.cq)
-        self.server.rkey = client_send_wr.rkey
+        self.assertEqual(self.client.mw.rkey, client_send_wr.rkey)
+        self.assertEqual(self.server.mw.rkey, server_send_wr.rkey)
+        self.server.rkey = self.client.mw.rkey
         self.server.raddr = self.client.mr.buf
-        self.client.rkey = server_send_wr.rkey
+        self.client.rkey = self.server.mw.rkey
         self.client.raddr = self.server.mr.buf
 
     def invalidate_mw_type1(self):