unit test exec_command and friends w.r.t. migration device state (#346)

Signed-off-by: Thanos Makatos <thanos.makatos@nutanix.com>
Reviewed-by: Swapnil Ingle <swapnil.ingle@nutanix.com>

4 files changed, 196 insertions, 3 deletions

diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt
index 3e3b0cc..21116eb 100644
--- a/test/CMakeLists.txt
+++ b/test/CMakeLists.txt
@@ -53,6 +53,10 @@ target_link_libraries(unit-tests PUBLIC "-Wl,--wrap=free")
 target_link_libraries(unit-tests PUBLIC "-Wl,--wrap=process_request")
 target_link_libraries(unit-tests PUBLIC "-Wl,--wrap=bind")
 target_link_libraries(unit-tests PUBLIC "-Wl,--wrap=listen")
+target_link_libraries(unit-tests PUBLIC "-Wl,--wrap=device_is_stopped_and_copying")
+target_link_libraries(unit-tests PUBLIC "-Wl,--wrap=cmd_allowed_when_stopped_and_copying")
+target_link_libraries(unit-tests PUBLIC "-Wl,--wrap=device_is_stopped")
+target_link_libraries(unit-tests PUBLIC "-Wl,--wrap=should_exec_command")
 target_link_libraries(unit-tests PUBLIC "-Wl,--wrap=handle_dirty_pages")
 
 enable_testing()
diff --git a/test/mocks.c b/test/mocks.c
index 80545e6..363e9ca 100644
--- a/test/mocks.c
+++ b/test/mocks.c
@@ -86,12 +86,12 @@ __wrap__dma_controller_do_remove_region(dma_controller_t *dma,
 }
 
 bool
-__wrap_device_is_stopped(struct migration *migr)
+__wrap_device_is_stopped(struct migration *migration)
 {
     if (!is_patched(device_is_stopped)) {
-        return __real_device_is_stopped(migr);
+        return __real_device_is_stopped(migration);
     }
-    check_expected(migr);
+    check_expected(migration);
     return mock();
 }
 
@@ -195,6 +195,47 @@ int __wrap_listen(int sockfd __attribute__((unused)),
     return 0;
 }
 
+bool
+__wrap_device_is_stopped_and_copying(struct migration *migration)
+{
+    if (!is_patched(device_is_stopped_and_copying)) {
+        return __real_device_is_stopped_and_copying(migration);
+    }
+    check_expected(migration);
+    return mock();
+}
+
+bool
+__wrap_cmd_allowed_when_stopped_and_copying(uint16_t cmd)
+{
+    if (!is_patched(cmd_allowed_when_stopped_and_copying)) {
+        return __real_cmd_allowed_when_stopped_and_copying(cmd);
+    }
+    check_expected(cmd);
+    return mock();
+}
+
+bool
+__wrap_cmd_allowed_when_stopped(struct migration *migration)
+{
+    if (!is_patched(device_is_stopped)) {
+        return __real_device_is_stopped(migration);
+    }
+    check_expected(migration);
+    return mock();
+}
+
+bool
+__wrap_should_exec_command(vfu_ctx_t *vfu_ctx, uint16_t cmd)
+{
+    if (!is_patched(should_exec_command)) {
+        return __real_should_exec_command(vfu_ctx, cmd);
+    }
+    check_expected(vfu_ctx);
+    check_expected(cmd);
+    return mock();
+}
+
 int
 __wrap_handle_dirty_pages(vfu_ctx_t *vfu_ctx, uint32_t size,
                           struct iovec **iovecs, size_t *nr_iovecs,
@@ -226,6 +267,10 @@ static struct function funcs[] = {
     {.addr = &__wrap_process_request},
     {.addr = &__wrap_bind},
     {.addr = &__wrap_listen},
+    {.addr = &__wrap_device_is_stopped_and_copying},
+    {.addr = &__wrap_cmd_allowed_when_stopped_and_copying},
+    {.addr = &__wrap_device_is_stopped},
+    {.addr = &__wrap_should_exec_command},
     {.addr = &__wrap_handle_dirty_pages},
 };
 
diff --git a/test/mocks.h b/test/mocks.h
index 06b853c..6f5765b 100644
--- a/test/mocks.h
+++ b/test/mocks.h
@@ -37,6 +37,9 @@ void patch(void *fn);
 
 bool is_patched(void *fn);
 
+bool
+__real_cmd_allowed_when_stopped_and_copying(u_int16_t cmd);
+
 int
 handle_dirty_pages(vfu_ctx_t *vfu_ctx, uint32_t size,
                    struct iovec **iovecs, size_t *nr_iovecs,
@@ -69,5 +72,10 @@ __real_free(void *ptr);
 int
 __real_process_request(vfu_ctx_t *vfu_ctx);
 
+bool
+__real_device_is_stopped_and_copying(struct migration *migration);
+
+bool
+__real_should_exec_command(vfu_ctx_t *vfu_ctx, uint16_t cmd);
 
 /* ex: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab: */
diff --git a/test/unit-tests.c b/test/unit-tests.c
index 4ab51e4..1520f31 100644
--- a/test/unit-tests.c
+++ b/test/unit-tests.c
@@ -48,6 +48,7 @@
 #include "migration.h"
 #include "mocks.h"
 #include "tran_sock.h"
+#include "migration_priv.h"
 
 static void
 test_dma_map_without_dma(void **state __attribute__((unused)))
@@ -1392,6 +1393,136 @@ test_setup_migration_callbacks(void **state)
 }
 
 static void
+test_device_is_stopped_and_copying(UNUSED void **state)
+{
+    vfu_ctx_t vfu_ctx = { .migration = NULL };
+
+    assert_false(device_is_stopped_and_copying(vfu_ctx.migration));
+    assert_false(device_is_stopped(vfu_ctx.migration));
+
+    size_t i;
+    struct migration migration;
+    vfu_ctx.migration = &migration;
+    for (i = 0; i < ARRAY_SIZE(migr_states); i++) {
+        if (migr_states[i].name == NULL) {
+            continue;
+        }
+        migration.info.device_state = i;
+        bool r = device_is_stopped_and_copying(vfu_ctx.migration);
+        if (i == VFIO_DEVICE_STATE_SAVING) {
+            assert_true(r);
+        } else {
+            assert_false(r);
+        }
+        r = device_is_stopped(vfu_ctx.migration);
+        if (i == VFIO_DEVICE_STATE_STOP) {
+            assert_true(r);
+        } else {
+            assert_false(r);
+        }
+    }
+}
+
+static void
+test_cmd_allowed_when_stopped_and_copying(UNUSED void **state)
+{
+    size_t i;
+
+    for (i = 0; i < VFIO_USER_MAX; i++) {
+        bool r = cmd_allowed_when_stopped_and_copying(i);
+        if (i == VFIO_USER_REGION_READ || i == VFIO_USER_REGION_WRITE ||
+            i == VFIO_USER_DIRTY_PAGES) {
+            assert_true(r);
+        } else {
+            assert_false(r);
+        }
+    }
+}
+
+static void
+test_should_exec_command(UNUSED void **state)
+{
+    struct migration migration = { { 0 } };
+    vfu_ctx_t vfu_ctx = { .migration = &migration };
+
+    patch(device_is_stopped_and_copying);
+    patch(cmd_allowed_when_stopped_and_copying);
+    patch(device_is_stopped);
+
+    /* XXX stopped and copying, command allowed */
+    will_return(__wrap_device_is_stopped_and_copying, true);
+    expect_value(__wrap_device_is_stopped_and_copying, migration, &migration);
+    will_return(__wrap_cmd_allowed_when_stopped_and_copying, true);
+    expect_value(__wrap_cmd_allowed_when_stopped_and_copying, cmd, 0xbeef);
+    assert_true(should_exec_command(&vfu_ctx, 0xbeef));
+
+    /* XXX stopped and copying, command not allowed */
+    will_return(__wrap_device_is_stopped_and_copying, true);
+    expect_any(__wrap_device_is_stopped_and_copying, migration);
+    will_return(__wrap_cmd_allowed_when_stopped_and_copying, false);
+    expect_any(__wrap_cmd_allowed_when_stopped_and_copying, cmd);
+    assert_false(should_exec_command(&vfu_ctx, 0xbeef));
+
+    /* XXX stopped */
+    will_return(__wrap_device_is_stopped_and_copying, false);
+    expect_any(__wrap_device_is_stopped_and_copying, migration);
+    will_return(__wrap_device_is_stopped, true);
+    expect_value(__wrap_device_is_stopped, migration, &migration);
+    assert_false(should_exec_command(&vfu_ctx, 0xbeef));
+
+    /* XXX none of the above */
+    will_return(__wrap_device_is_stopped_and_copying, false);
+    expect_any(__wrap_device_is_stopped_and_copying, migration);
+    will_return(__wrap_device_is_stopped, false);
+    expect_any(__wrap_device_is_stopped, migration);
+    assert_true(should_exec_command(&vfu_ctx, 0xbeef));
+}
+
+int recv_body(UNUSED vfu_ctx_t *vfu_ctx, UNUSED const struct vfio_user_header *hdr,
+              UNUSED void **datap)
+{
+    /* hack to avoid having to refactor the rest of exec_command */
+    return -1;
+}
+
+static void
+test_exec_command(UNUSED void **state)
+{
+    vfu_ctx_t vfu_ctx = { 0 };
+    struct vfio_user_header hdr = {
+        .cmd = 0xbeef,
+        .flags.type = VFIO_USER_F_TYPE_COMMAND,
+        .msg_size = sizeof hdr + 1
+    };
+    size_t size = sizeof hdr;
+    int fds = 0;
+    struct iovec _iovecs = { 0 };
+    struct iovec *iovecs = NULL;
+    size_t nr_iovecs = 0;
+    bool free_iovec_data = false;
+    int r;
+
+    /* XXX should NOT execute command */
+    patch(should_exec_command);
+    will_return(__wrap_should_exec_command, false);
+    expect_value(__wrap_should_exec_command, vfu_ctx, &vfu_ctx);
+    expect_value(__wrap_should_exec_command, cmd, 0xbeef);
+    r = exec_command(&vfu_ctx, &hdr, size, &fds, 0, NULL, NULL, &_iovecs,
+                     &iovecs, &nr_iovecs, &free_iovec_data);
+    assert_int_equal(-EINVAL, r);
+
+    /* XXX should execute command */
+    struct transport_ops tran = { .recv_body = recv_body };
+    vfu_ctx.tran = &tran;
+    will_return(__wrap_should_exec_command, true);
+    expect_value(__wrap_should_exec_command, vfu_ctx, &vfu_ctx);
+    expect_value(__wrap_should_exec_command, cmd, 0xbeef);
+    r = exec_command(&vfu_ctx, &hdr, size, &fds, 0, NULL, NULL, &_iovecs,
+                     &iovecs, &nr_iovecs, &free_iovec_data);
+    assert_int_equal(-1, r);
+}
+
+static void
 test_dirty_pages_without_dma(UNUSED void **state)
 {
     vfu_ctx_t vfu_ctx = { .migration = NULL };
@@ -1410,6 +1541,7 @@ test_dirty_pages_without_dma(UNUSED void **state)
     bool free_iovec_data = false;
     int r;
 
+
     patch(handle_dirty_pages);
 
     /* XXX w/o DMA controller */
@@ -1480,6 +1612,10 @@ int main(void)
         cmocka_unit_test_setup_teardown(test_setup_migration_callbacks,
             setup_test_setup_migration_region,
             teardown_test_setup_migration_region),
+        cmocka_unit_test_setup(test_device_is_stopped_and_copying, setup),
+        cmocka_unit_test_setup(test_cmd_allowed_when_stopped_and_copying, setup),
+        cmocka_unit_test_setup(test_should_exec_command, setup),
+        cmocka_unit_test_setup(test_exec_command, setup),
         cmocka_unit_test_setup(test_dirty_pages_without_dma, setup),
     };