[NTOSKRNL] Implement SepCleanupLUIDDeviceMapDirectory

This will clean up all the links (drive letters) created
by an user on session deletion once LUID device maps are
in use

ntoskrnl/se/srm.c

@@ -413,8 +413,238 @@ NTSTATUS
 SepCleanupLUIDDeviceMapDirectory(
     PLUID LogonLuid)
 {
-    UNIMPLEMENTED;
-    return STATUS_NOT_IMPLEMENTED;
+    BOOLEAN UseCurrentProc;
+    KAPC_STATE ApcState;
+    WCHAR Buffer[63];
+    UNICODE_STRING DirectoryName;
+    OBJECT_ATTRIBUTES ObjectAttributes;
+    NTSTATUS Status;
+    HANDLE DirectoryHandle, LinkHandle;
+    PHANDLE LinksBuffer;
+    POBJECT_DIRECTORY_INFORMATION DirectoryInfo;
+    ULONG LinksCount, LinksSize, DirInfoLength, ReturnLength, Context, CurrentLinks, i;
+    BOOLEAN RestartScan;
+
+    PAGED_CODE();
+
+    /* We need a logon LUID */
+    if (LogonLuid == NULL)
+    {
+        return STATUS_INVALID_PARAMETER;
+    }
+
+    /* Use current process */
+    UseCurrentProc = ObReferenceObjectSafe(PsGetCurrentProcess());
+    if (UseCurrentProc)
+    {
+        ObDereferenceObject(PsGetCurrentProcess());
+    }
+    /* Unless it's gone, then use system process */
+    else
+    {
+        KeStackAttachProcess(&PsInitialSystemProcess->Pcb, &ApcState);
+    }
+
+    /* Initialize our directory name */
+    _snwprintf(Buffer,
+               sizeof(Buffer) / sizeof(WCHAR),
+               L"\\Sessions\\0\\DosDevices\\%08x-%08x",
+               LogonLuid->HighPart,
+               LogonLuid->LowPart);
+    RtlInitUnicodeString(&DirectoryName, Buffer);
+
+    /* And open it */
+    InitializeObjectAttributes(&ObjectAttributes,
+                               &DirectoryName,
+                               OBJ_KERNEL_HANDLE | OBJ_CASE_INSENSITIVE,
+                               NULL,
+                               NULL);
+    Status = ZwOpenDirectoryObject(&DirectoryHandle,
+                                   DIRECTORY_QUERY,
+                                   &ObjectAttributes);
+    if (!NT_SUCCESS(Status))
+    {
+        if (!UseCurrentProc)
+        {
+            KeUnstackDetachProcess(&ApcState);
+        }
+
+        return Status;
+    }
+
+    /* Some initialization needed for browsing all our links... */
+    Context = 0;
+    DirectoryInfo = NULL;
+    DirInfoLength = 0;
+    /* In our buffer, we'll store at max 100 HANDLE */
+    LinksCount = 100;
+    CurrentLinks = 0;
+    /* Which gives a certain size */
+    LinksSize = LinksCount * sizeof(HANDLE);
+
+    /*
+     * This label is hit if we need to store more than a hundred
+     * of links. In that case, we jump here after having cleaned
+     * and deleted previous buffer.
+     * All handles have been already closed
+     */
+AllocateLinksAgain:
+    LinksBuffer = ExAllocatePoolWithTag(PagedPool, LinksSize, 'aHeS');
+    if (LinksBuffer == NULL)
+    {
+        /*
+         * Failure path: no need to clear handles:
+         * already closed and the buffer is already gone
+         */
+        ZwClose(DirectoryHandle);
+
+        /*
+         * On the first round, DirectoryInfo is NULL,
+         * if we grow LinksBuffer, it has been allocated
+         */
+        if (DirectoryInfo != NULL)
+        {
+            ExFreePoolWithTag(DirectoryInfo, 0);
+        }
+
+        if (!UseCurrentProc)
+        {
+            KeUnstackDetachProcess(&ApcState);
+        }
+
+        return STATUS_NO_MEMORY;
+    }
+
+    /*
+     * We always restart scan, but on the first loop
+     * if we couldn't fit everything in our buffer,
+     * then, we continue scan.
+     * But we restart if link buffer was too small
+     */
+    for (RestartScan = TRUE; ; RestartScan = FALSE)
+    {
+        /*
+         * Loop until our buffer is big enough to store
+         * one entry
+         */
+        while (TRUE)
+        {
+            Status = ZwQueryDirectoryObject(DirectoryHandle,
+                                            DirectoryInfo,
+                                            DirInfoLength,
+                                            TRUE,
+                                            RestartScan,
+                                            &Context,
+                                            &ReturnLength);
+            /* Only handle buffer growth in that loop */
+            if (Status != STATUS_BUFFER_TOO_SMALL)
+            {
+                break;
+            }
+
+            /* Get output length as new length */
+            DirInfoLength = ReturnLength;
+            /* Delete old buffer if any */
+            if (DirectoryInfo != NULL)
+            {
+                ExFreePoolWithTag(DirectoryInfo, 'bDeS');
+            }
+
+            /* And reallocate a bigger one */
+            DirectoryInfo = ExAllocatePoolWithTag(PagedPool, DirInfoLength, 'bDeS');
+            /* Fail if we cannot allocate */
+            if (DirectoryInfo == NULL)
+            {
+                Status = STATUS_INSUFFICIENT_RESOURCES;
+                break;
+            }
+        }
+
+        /* If querying the entry failed, quit */
+        if (!NT_SUCCESS(Status))
+        {
+            break;
+        }
+
+        /* We only look for symbolic links, the rest, we ignore */
+        if (wcscmp(DirectoryInfo->TypeName.Buffer, L"SymbolicLink"))
+        {
+            continue;
+        }
+
+        /* If our link buffer is out of space, reallocate */
+        if (CurrentLinks >= LinksCount)
+        {
+            /* First, close the links */
+            for (i = 0; i < CurrentLinks; ++i)
+            {
+                ZwClose(LinksBuffer[i]);
+            }
+
+            /* Allow 20 more HANDLEs */
+            LinksCount += 20;
+            CurrentLinks = 0;
+            ExFreePoolWithTag(LinksBuffer, 'aHeS');
+            LinksSize = LinksCount * sizeof(HANDLE);
+
+            /* And reloop again */
+            goto AllocateLinksAgain;
+        }
+
+        /* Open the found link */
+        InitializeObjectAttributes(&ObjectAttributes,
+                                   &DirectoryInfo->Name,
+                                   OBJ_KERNEL_HANDLE | OBJ_CASE_INSENSITIVE,
+                                   DirectoryHandle,
+                                   NULL);
+        if (NT_SUCCESS(ZwOpenSymbolicLinkObject(&LinkHandle,
+                                                SYMBOLIC_LINK_ALL_ACCESS,
+                                                &ObjectAttributes)))
+        {
+            /* If we cannot make it temporary, just close the link handle */
+            if (!NT_SUCCESS(ZwMakeTemporaryObject(LinkHandle)))
+            {
+                ZwClose(LinkHandle);
+            }
+            /* Otherwise, store it to defer deletion */
+            else
+            {
+                LinksBuffer[CurrentLinks] = LinkHandle;
+                ++CurrentLinks;
+            }
+        }
+    }
+
+    /* No more entries means we handled all links, that's not a failure */
+    if (Status == STATUS_NO_MORE_ENTRIES)
+    {
+        Status = STATUS_SUCCESS;
+    }
+
+    /* Close all the links we stored, this will like cause their deletion */
+    for (i = 0; i < CurrentLinks; ++i)
+    {
+        ZwClose(LinksBuffer[i]);
+    }
+    /* And free our links buffer */
+    ExFreePoolWithTag(LinksBuffer, 'aHeS');
+
+    /* Free our directory info buffer - it might be NULL if we failed realloc */
+    if (DirectoryInfo != NULL)
+    {
+        ExFreePoolWithTag(DirectoryInfo, 'bDeS');
+    }
+
+    /* Close our session directory */
+    ZwClose(DirectoryHandle);
+
+    /* And detach from system */
+    if (!UseCurrentProc)
+    {
+        KeUnstackDetachProcess(&ApcState);
+    }
+
+    return Status;
 }