changes from Thanos's review

Signed-off-by: William Henderson <william.henderson@nutanix.com>

4 files changed, 58 insertions, 63 deletions

diff --git a/include/libvfio-user.h b/include/libvfio-user.h
index 93c417c..570ae68 100644
--- a/include/libvfio-user.h
+++ b/include/libvfio-user.h
@@ -606,8 +606,8 @@ typedef struct {
     /*
      * Function that is called to read `count` bytes of migration data into
      * `buf`. The function must return the amount of data read or -1 on error,
-     * setting errno.
-     * 
+     * setting errno. The function may return less data than requested.
+     *
      * If the function returns zero, this is interpreted to mean that there is
      * no more migration data to read.
      */
@@ -616,7 +616,8 @@ typedef struct {
     /*
      * Function that is called for writing previously stored device state. The
      * function must return the amount of data written or -1 on error, setting
-     * errno.
+     * errno. Partial writes are not supported, so any return value other than
+     * `count` is invalid.
      */
     ssize_t (*write_data)(vfu_ctx_t *vfu_ctx, void *buf, uint64_t count);
 
diff --git a/lib/dma.c b/lib/dma.c
index 05aca6f..45a720a 100644
--- a/lib/dma.c
+++ b/lib/dma.c
@@ -611,11 +611,11 @@ dma_controller_dirty_page_get(dma_controller_t *dma, vfu_dma_addr_t addr,
     }
 
     if (pgsize == dma->dirty_pgsize) {
-        dirty_page_get_equal_pgsize(region, bitmap, bitmap_size);
+        dirty_page_get_same_pgsize(region, bitmap, bitmap_size);
     } else {
-        dirty_page_get_change_pgsize(region, bitmap, bitmap_size,
-                                     converted_bitmap_size, pgsize,
-                                     dma->dirty_pgsize);
+        dirty_page_get_different_pgsize(region, bitmap, bitmap_size,
+                                        converted_bitmap_size, pgsize,
+                                        dma->dirty_pgsize);
     }
 
 #ifdef DEBUG
@@ -625,38 +625,41 @@ dma_controller_dirty_page_get(dma_controller_t *dma, vfu_dma_addr_t addr,
     return 0;
 }
 
+static void
+dirty_page_exchange(uint8_t *outp, uint8_t *bitmap)
+{
+    /*
+     * If no bits are dirty, avoid the atomic exchange. This is obviously
+     * racy, but it's OK: if we miss a dirty bit being set, we'll catch it
+     * the next time around.
+     *
+     * Otherwise, atomically exchange the dirty bits with zero: as we use
+     * atomic or in _dma_mark_dirty(), this cannot lose set bits - we might
+     * miss a bit being set after, but again, we'll catch that next time
+     * around.
+     */
+    if (*bitmap == 0) {
+        *outp = 0;
+    } else {
+        uint8_t zero = 0;
+        __atomic_exchange(bitmap, &zero, outp, __ATOMIC_SEQ_CST);
+    }
+}
+
 void
-dirty_page_get_equal_pgsize(dma_memory_region_t *region, char *bitmap,
-                      size_t bitmap_size)
+dirty_page_get_same_pgsize(dma_memory_region_t *region, char *bitmap,
+                           size_t bitmap_size)
 {
     for (size_t i = 0; i < bitmap_size; i++) {
-        uint8_t val = region->dirty_bitmap[i];
-        uint8_t *outp = (uint8_t *)&bitmap[i];
-
-        /*
-         * If no bits are dirty, avoid the atomic exchange. This is obviously
-         * racy, but it's OK: if we miss a dirty bit being set, we'll catch it
-         * the next time around.
-         *
-         * Otherwise, atomically exchange the dirty bits with zero: as we use
-         * atomic or in _dma_mark_dirty(), this cannot lose set bits - we might
-         * miss a bit being set after, but again, we'll catch that next time
-         * around.
-         */
-        if (val == 0) {
-            *outp = 0;
-        } else {
-            uint8_t zero = 0;
-            __atomic_exchange(&region->dirty_bitmap[i], &zero,
-                              outp, __ATOMIC_SEQ_CST);
-        }
+        dirty_page_exchange((uint8_t *)&bitmap[i], &region->dirty_bitmap[i]);
     }
 }
 
 void
-dirty_page_get_change_pgsize(dma_memory_region_t *region, char *bitmap,
-                             size_t bitmap_size, size_t converted_bitmap_size,
-                             size_t pgsize, size_t converted_pgsize)
+dirty_page_get_different_pgsize(dma_memory_region_t *region, char *bitmap,
+                                size_t bitmap_size,
+                                size_t converted_bitmap_size, size_t pgsize,
+                                size_t converted_pgsize)
 {
     uint8_t bit = 0;
     size_t i;
@@ -668,26 +671,9 @@ dirty_page_get_change_pgsize(dma_memory_region_t *region, char *bitmap,
 
     for (i = 0; i < bitmap_size &&
          bit / CHAR_BIT < (size_t)converted_bitmap_size; i++) {
-        uint8_t val = region->dirty_bitmap[i];
         uint8_t out = 0;
 
-        /*
-         * If no bits are dirty, avoid the atomic exchange. This is obviously
-         * racy, but it's OK: if we miss a dirty bit being set, we'll catch it
-         * the next time around.
-         *
-         * Otherwise, atomically exchange the dirty bits with zero: as we use
-         * atomic or in _dma_mark_dirty(), this cannot lose set bits - we might
-         * miss a bit being set after, but again, we'll catch that next time
-         * around.
-         */
-        if (val == 0) {
-            out = 0;
-        } else {
-            uint8_t zero = 0;
-            __atomic_exchange(&region->dirty_bitmap[i], &zero,
-                                &out, __ATOMIC_SEQ_CST);
-        }
+        dirty_page_exchange(&out, &region->dirty_bitmap[i]);
 
         /*
          * Iterate through the bits of the byte, repeating or combining bits
diff --git a/lib/dma.h b/lib/dma.h
index 0a614de..d3c4e18 100644
--- a/lib/dma.h
+++ b/lib/dma.h
@@ -388,12 +388,13 @@ dma_controller_dirty_page_get(dma_controller_t *dma, vfu_dma_addr_t addr,
                               char *bitmap);
 
 void
-dirty_page_get_simple(dma_memory_region_t *region, char *bitmap,
-                      size_t bitmap_size);
+dirty_page_get_same_pgsize(dma_memory_region_t *region, char *bitmap,
+                           size_t bitmap_size);
 void
-dirty_page_get_complex(dma_memory_region_t *region, char *bitmap,
-                      size_t bitmap_size, size_t converted_bitmap_size,
-                      size_t pgsize, size_t converted_pgsize);
+dirty_page_get_different_pgsize(dma_memory_region_t *region, char *bitmap,
+                                size_t bitmap_size,
+                                size_t converted_bitmap_size, size_t pgsize,
+                                size_t converted_pgsize);
 
 bool
 dma_sg_is_mappable(const dma_controller_t *dma, const dma_sg_t *sg);
diff --git a/lib/migration.c b/lib/migration.c
index 9cadaf3..20d542c 100644
--- a/lib/migration.c
+++ b/lib/migration.c
@@ -47,14 +47,18 @@
  * The indices of each state are those in the vfio_user_device_mig_state enum.
  */
 static const char transitions[VFIO_USER_DEVICE_NUM_STATES] = {
-    0b00000000, // ERROR        -> {}
-    0b00011100, // STOP         -> {RUNNING, STOP_COPY, RESUMING}
-    0b01000010, // RUNNING      -> {STOP, PRE_COPY}
-    0b00000010, // STOP_COPY    -> {STOP}
-    0b00000010, // RESUMING     -> {STOP}
-    0b00000000, // RUNNING_P2P  -> {}
-    0b00001100, // PRE_COPY     -> {RUNNING, STOP_COPY}
-    0b00000000  // PRE_COPY_P2P -> {}
+    [VFIO_USER_DEVICE_STATE_ERROR] = 0,
+    [VFIO_USER_DEVICE_STATE_STOP] = (1 << VFIO_USER_DEVICE_STATE_RUNNING) |
+                                    (1 << VFIO_USER_DEVICE_STATE_STOP_COPY) |
+                                    (1 << VFIO_USER_DEVICE_STATE_RESUMING),
+    [VFIO_USER_DEVICE_STATE_RUNNING] = (1 << VFIO_USER_DEVICE_STATE_STOP) |
+                                       (1 << VFIO_USER_DEVICE_STATE_PRE_COPY),
+    [VFIO_USER_DEVICE_STATE_STOP_COPY] = 1 << VFIO_USER_DEVICE_STATE_STOP,
+    [VFIO_USER_DEVICE_STATE_RESUMING] = 1 << VFIO_USER_DEVICE_STATE_STOP,
+    [VFIO_USER_DEVICE_STATE_RUNNING_P2P] = 0,
+    [VFIO_USER_DEVICE_STATE_PRE_COPY] = (1 << VFIO_USER_DEVICE_STATE_RUNNING) |
+                                        (1 << VFIO_USER_DEVICE_STATE_STOP_COPY),
+    [VFIO_USER_DEVICE_STATE_PRE_COPY_P2P] = 0
 };
 
 /*
@@ -378,6 +382,9 @@ handle_mig_data_write(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg)
 
     if (ret < 0) {
         return ret;
+    } else if (ret != req->size) {
+        vfu_log(vfu_ctx, LOG_ERR, "migration data partial write");
+        return ERROR_INT(EINVAL);
     }
 
     return 0;