From 5562eb7ca9971ca9ad895519e2ee834e721545d8 Mon Sep 17 00:00:00 2001
From: William Henderson <william.henderson@nutanix.com>
Date: Mon, 4 Sep 2023 10:21:54 +0000
Subject: fix: migration quiesce and short write

Signed-off-by: William Henderson <william.henderson@nutanix.com>
---
 lib/libvfio-user.c | 17 ++++++++++++++++-
 lib/migration.c    | 14 ++++++++++++++
 lib/migration.h    |  3 +--
 3 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/lib/libvfio-user.c b/lib/libvfio-user.c
index 01c520e..0cf4b13 100644
--- a/lib/libvfio-user.c
+++ b/lib/libvfio-user.c
@@ -1522,6 +1522,7 @@ static bool
 command_needs_quiesce(vfu_ctx_t *vfu_ctx, const vfu_msg_t *msg)
 {
     struct vfio_user_region_access *reg;
+    struct vfio_user_device_feature *feature;
 
     if (vfu_ctx->quiesce == NULL) {
         return false;
@@ -1540,7 +1541,6 @@ command_needs_quiesce(vfu_ctx_t *vfu_ctx, const vfu_msg_t *msg)
             /*
              * bad request, it will be eventually failed by
              * handle_region_access
-             *
              */
             return false;
         }
@@ -1549,8 +1549,23 @@ command_needs_quiesce(vfu_ctx_t *vfu_ctx, const vfu_msg_t *msg)
             return true;
         }
         break;
+
+    case VFIO_USER_DEVICE_FEATURE:
+        if (msg->in.iov.iov_len < sizeof(*feature)) {
+            /*
+             * bad request, it will be eventually failed by
+             * handle_region_access
+             */
+            return false;
+        }
+        feature = msg->in.iov.iov_base;
+        if (migration_feature_needs_quiesce(feature)) {
+            return true;
+        }
+        break;
     }
 
+
     return false;
 }
 
diff --git a/lib/migration.c b/lib/migration.c
index cc6d271..6148e9e 100644
--- a/lib/migration.c
+++ b/lib/migration.c
@@ -385,6 +385,12 @@ handle_mig_data_write(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg)
         return ERROR_INT(EINVAL);
     }
 
+    if (msg->in.iov.iov_len < sizeof(struct vfio_user_mig_data) + req->size) {
+        vfu_log(vfu_ctx, LOG_ERR, "short write (%d < %ld)",
+                req->argsz, sizeof(struct vfio_user_mig_data) + req->size);
+        return ERROR_INT(EINVAL);
+    }
+
     ssize_t ret = migr->callbacks.write_data(vfu_ctx, &req->data, req->size);
 
     if (ret < 0) {
@@ -434,4 +440,12 @@ migration_set_pgsize(struct migration *migr, size_t pgsize)
     return 0;
 }
 
+bool
+migration_feature_needs_quiesce(struct vfio_user_device_feature *feature)
+{
+    return ((feature->flags &
+        (VFIO_DEVICE_FEATURE_SET | VFIO_DEVICE_FEATURE_MIG_DEVICE_STATE)) != 0)
+        && !(feature->flags & VFIO_DEVICE_FEATURE_PROBE);
+}
+
 /* ex: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab: */
diff --git a/lib/migration.h b/lib/migration.h
index 0d0dd76..928a7e5 100644
--- a/lib/migration.h
+++ b/lib/migration.h
@@ -85,8 +85,7 @@ MOCK_DECLARE(ssize_t, handle_device_state, vfu_ctx_t *vfu_ctx,
              struct migration *migr, uint32_t device_state, bool notify);
 
 bool
-access_migration_needs_quiesce(const vfu_ctx_t *vfu_ctx, size_t region_index,
-                              uint64_t offset);
+migration_feature_needs_quiesce(struct vfio_user_device_feature *feature);
 
 #endif /* LIB_VFIO_USER_MIGRATION_H */
 
-- 
cgit v1.1