api: use NT api directly for enumerating kernel modules

Signed-off-by: Jason A. Donenfeld <Jason@zx2c4.com>

5 files changed, 68 insertions, 27 deletions

diff --git a/api/adapter.c b/api/adapter.c
index 16bfc4d..c0df8e9 100644
--- a/api/adapter.c
+++ b/api/adapter.c
@@ -1012,47 +1012,45 @@ out:
 static DWORDLONG
 RunningWintunVersion(void)
 {
-    DWORD RequiredSize = 0, CurrentSize = 0;
-    VOID **Drivers = NULL;
     DWORDLONG Version = 0;
+    PRTL_PROCESS_MODULES Modules;
+    ULONG BufferSize = 2048;
     for (;;)
     {
-        if (!EnumDeviceDrivers(Drivers, CurrentSize, &RequiredSize))
-        {
-            LOG(WINTUN_LOG_ERR, L"Failed to enumerate drivers");
-            return Version;
-        }
-        if (CurrentSize == RequiredSize)
-            break;
-        if (Drivers)
-            HeapFree(ModuleHeap, 0, Drivers);
-        Drivers = HeapAlloc(ModuleHeap, 0, RequiredSize);
-        if (!Drivers)
+        Modules = HeapAlloc(ModuleHeap, 0, BufferSize);
+        if (!Modules)
         {
             LOG(WINTUN_LOG_ERR, L"Out of memory");
             return Version;
         }
-        CurrentSize = RequiredSize;
+        NTSTATUS Status = NtQuerySystemInformation(SystemModuleInformation, Modules, BufferSize, &BufferSize);
+        if (NT_SUCCESS(Status))
+            break;
+        HeapFree(ModuleHeap, 0, Modules);
+        if (Status == STATUS_INFO_LENGTH_MISMATCH)
+            continue;
+        LOG(WINTUN_LOG_ERR, L"Failed to enumerate drivers");
+        return Version;
     }
-    WCHAR MaybeWintun[11];
-    for (DWORD i = CurrentSize / sizeof(Drivers[0]); i-- > 0;)
+    for (ULONG i = 0; i < Modules->NumberOfModules; ++i)
     {
-        if (GetDeviceDriverBaseNameW(Drivers[i], MaybeWintun, _countof(MaybeWintun)) == 10 &&
-            !_wcsicmp(MaybeWintun, L"wintun.sys"))
+        if (!_stricmp((const char *)&Modules->Modules[i].FullPathName[Modules->Modules[i].OffsetToFileName], "wintun.sys"))
         {
-            WCHAR WintunPath[MAX_PATH + 2];
-            DWORD Len = GetDeviceDriverFileNameW(Drivers[i], WintunPath, _countof(WintunPath));
-            if (!Len || Len == _countof(WintunPath) - 1)
+            size_t Size = strlen((const char *)Modules->Modules[i].FullPathName) + 1;
+            WCHAR *FilePathName = HeapAlloc(ModuleHeap, 0, Size * 2);
+            if (!FilePathName)
             {
-                LOG(WINTUN_LOG_ERR, L"Failed to locate driver path");
+                LOG(WINTUN_LOG_ERR, L"Out of memory");
                 goto out;
             }
-            Version = VersionOfFile(WintunPath);
+            mbstowcs_s(&Size, FilePathName, Size, (const char *)Modules->Modules[i].FullPathName, _TRUNCATE);
+            Version = VersionOfFile(FilePathName);
+            HeapFree(ModuleHeap, 0, FilePathName);
             goto out;
         }
     }
 out:
-    HeapFree(ModuleHeap, 0, Drivers);
+    HeapFree(ModuleHeap, 0, Modules);
     return Version;
 }
 
diff --git a/api/api.vcxproj b/api/api.vcxproj
index 7fed028..42dc2c6 100644
--- a/api/api.vcxproj
+++ b/api/api.vcxproj
@@ -199,6 +199,7 @@
     <ClInclude Include="logger.h" />
     <ClInclude Include="namespace.h" />
     <ClInclude Include="nci.h" />
+    <ClInclude Include="ntldr.h" />
     <ClInclude Include="pch.h" />
     <ClInclude Include="registry.h" />
     <ClInclude Include="resource.h" />
diff --git a/api/api.vcxproj.filters b/api/api.vcxproj.filters
index c585527..c65b099 100644
--- a/api/api.vcxproj.filters
+++ b/api/api.vcxproj.filters
@@ -52,6 +52,9 @@
     <ClInclude Include="wintun.h">
       <Filter>Header Files</Filter>
     </ClInclude>
+    <ClInclude Include="ntldr.h">
+      <Filter>Header Files</Filter>
+    </ClInclude>
     <ClInclude Include="elevate.h">
       <Filter>Header Files</Filter>
     </ClInclude>
diff --git a/api/ntldr.h b/api/ntldr.h
new file mode 100644
index 0000000..dca0e78
--- /dev/null
+++ b/api/ntldr.h
@@ -0,0 +1,33 @@
+/* SPDX-License-Identifier: GPL-2.0
+ *
+ * Copyright (C) 2018-2020 WireGuard LLC. All Rights Reserved.
+ */
+
+#pragma once
+
+#include <Windows.h>
+
+enum
+{
+    SystemModuleInformation = 11
+};
+
+typedef struct _RTL_PROCESS_MODULE_INFORMATION
+{
+    HANDLE Section;
+    PVOID MappedBase;
+    PVOID ImageBase;
+    ULONG ImageSize;
+    ULONG Flags;
+    USHORT LoadOrderIndex;
+    USHORT InitOrderIndex;
+    USHORT LoadCount;
+    USHORT OffsetToFileName;
+    UCHAR FullPathName[256];
+} RTL_PROCESS_MODULE_INFORMATION, *PRTL_PROCESS_MODULE_INFORMATION;
+
+typedef struct _RTL_PROCESS_MODULES
+{
+    ULONG NumberOfModules;
+    RTL_PROCESS_MODULE_INFORMATION Modules[1];
+} RTL_PROCESS_MODULES, *PRTL_PROCESS_MODULES;
diff --git a/api/pch.h b/api/pch.h
index 45bf3b1..4defd82 100644
--- a/api/pch.h
+++ b/api/pch.h
@@ -11,6 +11,7 @@
 #include "logger.h"
 #include "namespace.h"
 #include "nci.h"
+#include "ntldr.h"
 #include "registry.h"
 #include "resource.h"
 #include "wintun.h"
@@ -19,12 +20,13 @@
 #pragma warning(disable: 4201) /* nonstandard extension used: nameless struct/union */
 #include <bcrypt.h>
 #include <cfgmgr32.h>
+#include <delayimp.h>
 #include <devguid.h>
+#include <IPExport.h>
 #include <iphlpapi.h>
 #include <locale.h>
 #include <ndisguid.h>
 #include <newdev.h>
-#include <NTSecAPI.h>
 #include <objbase.h>
 #include <Psapi.h>
 #include <sddl.h>
@@ -32,6 +34,10 @@
 #include <Shlwapi.h>
 #include <string.h>
 #include <TlHelp32.h>
-#include <delayimp.h>
 #include <wchar.h>
-#pragma warning(pop)
+#include <Windows.h>
+#include <winternl.h>
+#define _NTDEF_ //TODO: find a better way to include both ntsecapi.h and winternl.h or include ntdef.h for real somehow
+#include <NTSecAPI.h>
+#define STATUS_INFO_LENGTH_MISMATCH ((NTSTATUS)0xC0000004L) //TODO: #include <ntstatus.h> instead of this
+#pragma warning(pop)
\ No newline at end of file