From 4da92f3203117a6689644637f8a1568eb7abff71 Mon Sep 17 00:00:00 2001
From: Zebediah Figura <z.figura12@gmail.com>
Date: Mon, 15 Jun 2020 18:19:55 -0500
Subject: [PATCH] ntoskrnl.exe/tests: Test completion callback parameters.

Signed-off-by: Zebediah Figura <z.figura12@gmail.com>
Signed-off-by: Alexandre Julliard <julliard@winehq.org>
---
 dlls/ntoskrnl.exe/tests/driver.c | 47 ++++++++++++++++++++++++++++++++
 dlls/ntoskrnl.exe/tests/driver.h |  1 +
 include/ddk/wdm.h                |  8 ++++++
 3 files changed, 56 insertions(+)

diff --git a/dlls/ntoskrnl.exe/tests/driver.c b/dlls/ntoskrnl.exe/tests/driver.c
index fa8d7b20e3e..f4b373264fc 100644
--- a/dlls/ntoskrnl.exe/tests/driver.c
+++ b/dlls/ntoskrnl.exe/tests/driver.c
@@ -1570,6 +1570,33 @@ static void test_stack_limits(void)
     ok(low < (ULONG_PTR)&low && (ULONG_PTR)&low < high, "stack variable is not in stack limits\n");
 }
 
+static unsigned int got_completion;
+
+static NTSTATUS WINAPI completion_cb(DEVICE_OBJECT *device, IRP *irp, void *context)
+{
+    todo_wine ok(device == context, "Got device %p; expected %p.\n", device, context);
+    ++got_completion;
+    return STATUS_SUCCESS;
+}
+
+static void test_completion(void)
+{
+    IO_STATUS_BLOCK io;
+    NTSTATUS ret;
+    KEVENT event;
+    IRP *irp;
+
+    KeInitializeEvent(&event, NotificationEvent, FALSE);
+
+    irp = IoBuildDeviceIoControlRequest(IOCTL_WINETEST_COMPLETION, upper_device,
+            NULL, 0, NULL, 0, FALSE, &event, &io);
+
+    IoSetCompletionRoutine(irp, completion_cb, NULL, TRUE, TRUE, TRUE);
+    ret = IoCallDriver(upper_device, irp);
+    ok(ret == STATUS_SUCCESS, "IoCallDriver returned %#x\n", ret);
+    ok(got_completion == 2, "got %u calls to completion routine\n");
+}
+
 static void test_IoAttachDeviceToDeviceStack(void)
 {
     DEVICE_OBJECT *dev1, *dev2, *dev3, *ret;
@@ -1714,6 +1741,7 @@ static void WINAPI main_test_task(DEVICE_OBJECT *device, void *context)
     test_call_driver(device);
     test_cancel_irp(device);
     test_stack_limits();
+    test_completion();
 
     /* print process report */
     if (winetest_debug)
@@ -2174,6 +2202,23 @@ static NTSTATUS test_mismatched_status_ioctl(IRP *irp, IO_STACK_LOCATION *stack,
     return STATUS_SUCCESS;
 }
 
+static NTSTATUS test_completion_ioctl(DEVICE_OBJECT *device, IRP *irp)
+{
+    if (device == upper_device)
+    {
+        IoCopyCurrentIrpStackLocationToNext(irp);
+        IoSetCompletionRoutine(irp, completion_cb, upper_device, TRUE, TRUE, TRUE);
+        return IoCallDriver(lower_device, irp);
+    }
+    else
+    {
+        ok(device == lower_device, "Got wrong device.\n");
+        irp->IoStatus.Status = STATUS_SUCCESS;
+        IoCompleteRequest(irp, IO_NO_INCREMENT);
+        return STATUS_SUCCESS;
+    }
+}
+
 static NTSTATUS WINAPI driver_Create(DEVICE_OBJECT *device, IRP *irp)
 {
     IO_STACK_LOCATION *irpsp = IoGetCurrentIrpStackLocation( irp );
@@ -2237,6 +2282,8 @@ static NTSTATUS WINAPI driver_IoControl(DEVICE_OBJECT *device, IRP *irp)
             break;
         case IOCTL_WINETEST_MISMATCHED_STATUS:
             return test_mismatched_status_ioctl(irp, stack, &irp->IoStatus.Information);
+        case IOCTL_WINETEST_COMPLETION:
+            return test_completion_ioctl(device, irp);
         default:
             break;
     }
diff --git a/dlls/ntoskrnl.exe/tests/driver.h b/dlls/ntoskrnl.exe/tests/driver.h
index 6012efc41c8..4b70046b8e2 100644
--- a/dlls/ntoskrnl.exe/tests/driver.h
+++ b/dlls/ntoskrnl.exe/tests/driver.h
@@ -34,6 +34,7 @@
 #define IOCTL_WINETEST_GET_FSCONTEXT    CTL_CODE(FILE_DEVICE_UNKNOWN, 0x809, METHOD_BUFFERED, FILE_ANY_ACCESS)
 #define IOCTL_WINETEST_RETURN_STATUS    CTL_CODE(FILE_DEVICE_UNKNOWN, 0x80a, METHOD_BUFFERED, FILE_ANY_ACCESS)
 #define IOCTL_WINETEST_MISMATCHED_STATUS CTL_CODE(FILE_DEVICE_UNKNOWN, 0x80b, METHOD_NEITHER, FILE_ANY_ACCESS)
+#define IOCTL_WINETEST_COMPLETION       CTL_CODE(FILE_DEVICE_UNKNOWN, 0x80c, METHOD_NEITHER, FILE_ANY_ACCESS)
 
 static const char teststr[] = "Wine is not an emulator";
 
diff --git a/include/ddk/wdm.h b/include/ddk/wdm.h
index 85d637639f5..23ac630f410 100644
--- a/include/ddk/wdm.h
+++ b/include/ddk/wdm.h
@@ -1599,6 +1599,14 @@ static inline void IoMarkIrpPending(IRP *irp)
     IoGetCurrentIrpStackLocation(irp)->Control |= SL_PENDING_RETURNED;
 }
 
+static inline void IoCopyCurrentIrpStackLocationToNext(IRP *irp)
+{
+    IO_STACK_LOCATION *current = IoGetCurrentIrpStackLocation(irp);
+    IO_STACK_LOCATION *next = IoGetNextIrpStackLocation(irp);
+    memcpy(next, current, FIELD_OFFSET(IO_STACK_LOCATION, CompletionRoutine));
+    next->Control = 0;
+}
+
 #define KernelMode 0
 #define UserMode   1