From 1074a9aaff0b8d03782132cbe08f035afa85de9b Mon Sep 17 00:00:00 2001
From: Pierre Schweitzer <pierre@reactos.org>
Date: Sat, 1 Jun 2019 15:18:52 +0200
Subject: [PATCH] [NTOSKRNL] Implement support for
 PROCESS_DEVICEMAP_INFORMATION_EX in NtQueryInformationProcess

---
 ntoskrnl/include/internal/ob.h |  3 +-
 ntoskrnl/ob/devicemap.c        |  3 +-
 ntoskrnl/ps/query.c            | 53 +++++++++++++++++++++++++++-------
 3 files changed, 47 insertions(+), 12 deletions(-)

diff --git a/ntoskrnl/include/internal/ob.h b/ntoskrnl/include/internal/ob.h
index 132b47f5a6b..910858ac5de 100644
--- a/ntoskrnl/include/internal/ob.h
+++ b/ntoskrnl/include/internal/ob.h
@@ -296,7 +296,8 @@ VOID
 NTAPI
 ObQueryDeviceMapInformation(
     IN PEPROCESS Process,
-    OUT PPROCESS_DEVICEMAP_INFORMATION DeviceMapInfo
+    OUT PPROCESS_DEVICEMAP_INFORMATION DeviceMapInfo,
+    IN ULONG Flags
 );
 
 //
diff --git a/ntoskrnl/ob/devicemap.c b/ntoskrnl/ob/devicemap.c
index a4cd5afbfae..d1b7368bb8c 100644
--- a/ntoskrnl/ob/devicemap.c
+++ b/ntoskrnl/ob/devicemap.c
@@ -229,7 +229,8 @@ ObInheritDeviceMap(IN PEPROCESS Parent,
 VOID
 NTAPI
 ObQueryDeviceMapInformation(IN PEPROCESS Process,
-                            IN PPROCESS_DEVICEMAP_INFORMATION DeviceMapInfo)
+                            IN PPROCESS_DEVICEMAP_INFORMATION DeviceMapInfo,
+                            IN ULONG Flags)
 {
     PDEVICE_MAP DeviceMap;
 
diff --git a/ntoskrnl/ps/query.c b/ntoskrnl/ps/query.c
index 12036fb01a1..29c8171a474 100644
--- a/ntoskrnl/ps/query.c
+++ b/ntoskrnl/ps/query.c
@@ -78,7 +78,7 @@ NtQueryInformationProcess(IN HANDLE ProcessHandle,
     PVM_COUNTERS VmCounters = (PVM_COUNTERS)ProcessInformation;
     PIO_COUNTERS IoCounters = (PIO_COUNTERS)ProcessInformation;
     PQUOTA_LIMITS QuotaLimits = (PQUOTA_LIMITS)ProcessInformation;
-    PROCESS_DEVICEMAP_INFORMATION DeviceMap;
+    PROCESS_DEVICEMAP_INFORMATION_EX DeviceMap;
     PUNICODE_STRING ImageName;
     ULONG Cookie, ExecuteOptions = 0;
     ULONG_PTR Wow64 = 0;
@@ -564,22 +564,55 @@ NtQueryInformationProcess(IN HANDLE ProcessHandle,
         /* DOS Device Map */
         case ProcessDeviceMap:
 
-            if (ProcessInformationLength != RTL_FIELD_SIZE(PROCESS_DEVICEMAP_INFORMATION, Query))
+            if (ProcessInformationLength < sizeof(PROCESS_DEVICEMAP_INFORMATION))
             {
-                if (ProcessInformationLength == sizeof(PROCESS_DEVICEMAP_INFORMATION_EX))
+                Status = STATUS_INFO_LENGTH_MISMATCH;
+                break;
+            }
+
+            if (ProcessInformationLength == sizeof(PROCESS_DEVICEMAP_INFORMATION_EX))
+            {
+                /* Protect read in SEH */
+                _SEH2_TRY
                 {
-                    DPRINT1("PROCESS_DEVICEMAP_INFORMATION_EX not supported!\n");
-                    Status = STATUS_NOT_IMPLEMENTED;
+                    PPROCESS_DEVICEMAP_INFORMATION_EX DeviceMapEx = ProcessInformation;
+
+                    DeviceMap.Flags = DeviceMapEx->Flags;
                 }
-                else
+                _SEH2_EXCEPT(EXCEPTION_EXECUTE_HANDLER)
+                {
+                    /* Get the exception code */
+                    Status = _SEH2_GetExceptionCode();
+                }
+                _SEH2_END;
+
+                if (!NT_SUCCESS(Status))
+                {
+                    break;
+                }
+
+                /* Only one flag is supported and it needs LUID mappings */
+                if ((DeviceMap.Flags & ~PROCESS_LUID_DOSDEVICES_ONLY) != 0 ||
+                    !ObIsLUIDDeviceMapsEnabled())
+                {
+                    Status = STATUS_INVALID_PARAMETER;
+                    break;
+                }
+            }
+            else
+            {
+                if (ProcessInformationLength != sizeof(PROCESS_DEVICEMAP_INFORMATION))
                 {
                     Status = STATUS_INFO_LENGTH_MISMATCH;
+                    break;
                 }
-                break;
+
+                /* No flags for standard call */
+                DeviceMap.Flags = 0;
             }
 
             /* Set the return length */
-            Length = sizeof(PROCESS_DEVICEMAP_INFORMATION);
+            Length = ProcessInformationLength;
 
             /* Reference the process */
             Status = ObReferenceObjectByHandle(ProcessHandle,
@@ -591,12 +624,12 @@ NtQueryInformationProcess(IN HANDLE ProcessHandle,
             if (!NT_SUCCESS(Status)) break;
 
             /* Query the device map information */
-            ObQueryDeviceMapInformation(Process, &DeviceMap);
+            ObQueryDeviceMapInformation(Process, (PPROCESS_DEVICEMAP_INFORMATION)&DeviceMap, DeviceMap.Flags);
 
             /* Enter SEH for writing back data */
             _SEH2_TRY
             {
-                *(PPROCESS_DEVICEMAP_INFORMATION)ProcessInformation = DeviceMap;
+                RtlCopyMemory(ProcessInformation, &DeviceMap, sizeof(PROCESS_DEVICEMAP_INFORMATION));
             }
             _SEH2_EXCEPT(EXCEPTION_EXECUTE_HANDLER)
             {
-- 
2.17.1