[NTOSKRNL] Implement ObpSetCurrentProcessDeviceMap

[reactos.git] / ntoskrnl / ob / devicemap.c

diff --git a/ntoskrnl/ob/devicemap.c b/ntoskrnl/ob/devicemap.c
index 4dcaff1..8a47144 100644 (file)
--- a/ntoskrnl/ob/devicemap.c
+++ b/ntoskrnl/ob/devicemap.c
@@ -5,6 +5,7 @@
  * PURPOSE:         Device map implementation
  * PROGRAMMERS:     Eric Kohl (eric.kohl@reactos.org)
  *                  Alex Ionescu (alex.ionescu@reactos.org)
+ *                  Pierre Schweitzer (pierre@reactos.org)
  */
 
 /* INCLUDES ***************************************************************/
@@ -143,12 +144,211 @@ ObSetDeviceMap(IN PEPROCESS Process,
 }
 
 
+NTSTATUS
+NTAPI
+ObpSetCurrentProcessDeviceMap(VOID)
+{
+    PTOKEN Token;
+    LUID LogonId;
+    NTSTATUS Status;
+    PEPROCESS CurrentProcess;
+    LUID SystemLuid = SYSTEM_LUID;
+    PDEVICE_MAP DeviceMap, OldDeviceMap;
+
+    /* Get our impersonation token */
+    CurrentProcess = PsGetCurrentProcess();
+    Token = PsReferencePrimaryToken(CurrentProcess);
+    if (Token == NULL)
+    {
+        return STATUS_NO_TOKEN;
+    }
+
+    /* Query the Logon ID */
+    Status = SeQueryAuthenticationIdToken(Token, &LogonId);
+    if (!NT_SUCCESS(Status))
+    {
+        goto done;
+    }
+
+    /* If that's system, then use system device map */
+    if (RtlEqualLuid(&LogonId, &SystemLuid))
+    {
+        DeviceMap = ObSystemDeviceMap;
+    }
+    /* Otherwise ask Se for the device map */
+    else
+    {
+        Status = SeGetLogonIdDeviceMap(&LogonId, &DeviceMap);
+        if (!NT_SUCCESS(Status))
+        {
+            /* Normalize failure status */
+            Status = STATUS_OBJECT_PATH_INVALID;
+            goto done;
+        }
+    }
+
+    /* Fail if no device map */
+    if (DeviceMap == NULL)
+    {
+        Status = STATUS_OBJECT_PATH_INVALID;
+        goto done;
+    }
+
+    /* Acquire the device map lock */
+    KeAcquireGuardedMutex(&ObpDeviceMapLock);
+
+    /* Save old device map attached to the process */
+    OldDeviceMap = CurrentProcess->DeviceMap;
+
+    /* Set new device map & reference it */
+    ++DeviceMap->ReferenceCount;
+    CurrentProcess->DeviceMap = DeviceMap;
+
+    /* Release the device map lock */
+    KeReleaseGuardedMutex(&ObpDeviceMapLock);
+
+    /* If we had a device map, dereference it */
+    if (OldDeviceMap != NULL)
+    {
+        ObfDereferenceDeviceMap(OldDeviceMap);
+    }
+
+done:
+    /* We're done with the token! */
+    ObDereferenceObject(Token);
+
+    return Status;
+}
+
+
 PDEVICE_MAP
 NTAPI
 ObpReferenceDeviceMap(VOID)
 {
-    UNIMPLEMENTED;
-    return NULL;
+    LUID LogonId;
+    NTSTATUS Status;
+    PTOKEN Token = NULL;
+    PDEVICE_MAP DeviceMap;
+    PETHREAD CurrentThread;
+    BOOLEAN LookingForSystem;
+    LUID SystemLuid = SYSTEM_LUID;
+    BOOLEAN CopyOnOpen, EffectiveOnly;
+    SECURITY_IMPERSONATION_LEVEL ImpersonationLevel;
+
+    LookingForSystem = FALSE;
+
+    /* If LUID mapping is enable, try to get appropriate device map */
+    if (ObpLUIDDeviceMapsEnabled != 0)
+    {
+        /* In case of impersonation, we've got a bit of work to do */
+        CurrentThread = PsGetCurrentThread();
+        if (CurrentThread->ActiveImpersonationInfo)
+        {
+            /* Get impersonation token */
+            Token = PsReferenceImpersonationToken(CurrentThread,
+                                                  &CopyOnOpen,
+                                                  &EffectiveOnly,
+                                                  &ImpersonationLevel);
+            /* Get logon LUID */
+            if (Token != NULL)
+            {
+                Status = SeQueryAuthenticationIdToken(Token, &LogonId);
+            }
+            else
+            {
+                /* Force failure */
+                Status = STATUS_NO_TOKEN;
+            }
+
+            /* If we got logon LUID */
+            if (NT_SUCCESS(Status))
+            {
+                /*
+                 * Check it's not system, system is easy to handle,
+                 * we just need to return ObSystemDeviceMap
+                 */
+                if (!RtlEqualLuid(&LogonId, &SystemLuid))
+                {
+                    /* Ask Se for the device  map */
+                    Status = SeGetLogonIdDeviceMap(&LogonId, &DeviceMap);
+                    if (NT_SUCCESS(Status))
+                    {
+                        /* Acquire the device map lock */
+                        KeAcquireGuardedMutex(&ObpDeviceMapLock);
+
+                        /* Reference the device map if any */
+                        if (DeviceMap != NULL)
+                        {
+                            ++DeviceMap->ReferenceCount;
+                        }
+
+                        /* Release the device map lock */
+                        KeReleaseGuardedMutex(&ObpDeviceMapLock);
+
+                        /* If we got the device map, we're done! */
+                        if (DeviceMap != NULL)
+                        {
+                            ObDereferenceObject(Token);
+
+                            return DeviceMap;
+                        }
+                    }
+                }
+                else
+                {
+                    LookingForSystem = TRUE;
+                }
+            }
+        }
+
+        /*
+         * Fall back case of the LUID mapping, make sure there's a
+         * a device map attached to the current process
+         */
+        if (PsGetCurrentProcess()->DeviceMap == NULL &&
+            !NT_SUCCESS(ObpSetCurrentProcessDeviceMap()))
+        {
+            /* We may have failed after we got impersonation token */
+            if (Token != NULL)
+            {
+                ObDereferenceObject(Token);
+            }
+
+            return NULL;
+        }
+    }
+
+    /* Acquire the device map lock */
+    KeAcquireGuardedMutex(&ObpDeviceMapLock);
+
+    /* If we're looking for system map, use it */
+    if (LookingForSystem)
+    {
+        DeviceMap = ObSystemDeviceMap;
+    }
+    /* Otherwise, use current process device map */
+    else
+    {
+        DeviceMap = PsGetCurrentProcess()->DeviceMap;
+    }
+
+    /* If we got one, reference it */
+    if (DeviceMap != NULL)
+    {
+        ++DeviceMap->ReferenceCount;
+    }
+
+    /* Release the device map lock */
+    KeReleaseGuardedMutex(&ObpDeviceMapLock);
+
+    /* We may have impersonation token (if we failed in impersonation branch) */
+    if (Token != NULL)
+    {
+        ObDereferenceObject(Token);
+    }
+
+    /* Return the potentially found device map */
+    return DeviceMap;
 }