fix: more implementation changes for comments

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

2 files changed, 22 insertions, 4 deletions

diff --git a/lib/libvfio-user.c b/lib/libvfio-user.c
index 145146d..6927942 100644
--- a/lib/libvfio-user.c
+++ b/lib/libvfio-user.c
@@ -1065,16 +1065,29 @@ handle_device_feature(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg)
     } else if (req->flags & VFIO_DEVICE_FEATURE_GET) {
         // all supported outgoing data is currently the same size as 
         //   vfio_user_device_feature_migration
-        msg->out.iov.iov_len = sizeof(struct vfio_user_device_feature_migration);
+        msg->out.iov.iov_len = sizeof(struct vfio_user_device_feature)
+            + sizeof(struct vfio_user_device_feature_migration);
         msg->out.iov.iov_base = calloc(1, msg->out.iov.iov_len);
 
         if (msg->out.iov.iov_base == NULL) {
             return -1;
         }
 
+        memcpy(msg->out.iov.iov_base, msg->in.iov.iov_base,
+               sizeof(struct vfio_user_device_feature));
+
         ret = migration_feature_get(vfu_ctx,
                                     req->flags & VFIO_DEVICE_FEATURE_MASK,
-                                    msg->out.iov.iov_base);
+                                    msg->out.iov.iov_base +
+                                    sizeof(struct vfio_user_device_feature));
+
+        struct vfio_user_device_feature *res = msg->out.iov.iov_base;
+
+        if (ret < 0) {
+            msg->out.iov.iov_len = sizeof(struct vfio_user_device_feature);
+        } else {
+            res->argsz = sizeof(struct vfio_user_device_feature) + ret;
+        }
     } else if (req->flags & VFIO_DEVICE_FEATURE_SET) {
         msg->out.iov.iov_base = malloc(msg->in.iov.iov_len);
         msg->out.iov.iov_len = msg->in.iov.iov_len;
diff --git a/lib/migration.c b/lib/migration.c
index 69f3f35..80af2e0 100644
--- a/lib/migration.c
+++ b/lib/migration.c
@@ -264,8 +264,13 @@ handle_mig_data_read(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg)
 
     ssize_t ret = migr->callbacks.read_data(vfu_ctx, &res->data, req->size);
 
-    res->size = ret;
-    res->argsz = ret + msg->in.iov.iov_len;
+    if (ret < 0) {
+        res->size = 0;
+        res->argsz = sizeof(struct vfio_user_mig_data);
+    } else {
+        res->size = ret;
+        res->argsz = sizeof(struct vfio_user_mig_data) + ret;
+    }
 
     return ret;
 }