From d6cbb73f4e1943e34f67711a0431d8a3788bc889 Mon Sep 17 00:00:00 2001
From: Bert Belder <bertbelder@gmail.com>
Date: Sun, 3 Sep 2017 20:51:17 +0200
Subject: [PATCH] nt: separate out ntdll.dll dynamic import logic

---
 src/epoll.c | 31 ++++++++++++-------------------
 src/nt.c    | 27 +++++++++++++++++++++++++++
 src/nt.h    | 31 +++++++++++++++++++++----------
 3 files changed, 60 insertions(+), 29 deletions(-)
 create mode 100644 src/nt.c

diff --git a/src/epoll.c b/src/epoll.c
index 9505615..80496ac 100644
--- a/src/epoll.c
+++ b/src/epoll.c
@@ -6,6 +6,7 @@
 #include "afd.h"
 #include "epoll.h"
 #include "error.h"
+#include "nt.h"
 #include "tree.h"
 #include "win.h"
 
@@ -50,7 +51,6 @@ static int epoll__afd_poll(SOCKET socket,
                            OVERLAPPED* overlapped);
 
 static int epoll__initialized = 0;
-static PNTDEVICEIOCONTROLFILE pNtDeviceIoControlFile;
 
 /* State associated with a epoll handle. */
 struct epoll_port_data_s {
@@ -570,7 +570,6 @@ int epoll_close(epoll_t port_handle) {
 }
 
 int epoll__initialize() {
-  HMODULE ntdll;
   int r;
   WSADATA wsa_data;
 
@@ -578,13 +577,7 @@ int epoll__initialize() {
   if (r != 0)
     return -1;
 
-  ntdll = LoadLibraryW(L"ntdll.dll");
-  if (ntdll == NULL)
-    return -1;
-
-  pNtDeviceIoControlFile =
-      (PNTDEVICEIOCONTROLFILE) GetProcAddress(ntdll, "NtDeviceIoControlFile");
-  if (pNtDeviceIoControlFile == NULL)
+  if (nt_initialize() < 0)
     return -1;
 
   return 0;
@@ -747,16 +740,16 @@ int epoll__afd_poll(SOCKET socket,
   }
 
   iosb_ptr->Status = STATUS_PENDING;
-  status = pNtDeviceIoControlFile((HANDLE) socket,
-                                  event,
-                                  NULL,
-                                  apc_context,
-                                  iosb_ptr,
-                                  IOCTL_AFD_POLL,
-                                  info,
-                                  sizeof *info,
-                                  info,
-                                  sizeof *info);
+  status = NtDeviceIoControlFile((HANDLE) socket,
+                                 event,
+                                 NULL,
+                                 apc_context,
+                                 iosb_ptr,
+                                 IOCTL_AFD_POLL,
+                                 info,
+                                 sizeof *info,
+                                 info,
+                                 sizeof *info);
 
   if (overlapped == NULL) {
     /* If this is a blocking operation, wait for the event to become
diff --git a/src/nt.c b/src/nt.c
new file mode 100644
index 0000000..b9133e0
--- /dev/null
+++ b/src/nt.c
@@ -0,0 +1,27 @@
+#include <stdlib.h>
+
+#include "nt.h"
+#include "win.h"
+
+#define X(return_type, declarators, name, ...) \
+  return_type(declarators* name)(__VA_ARGS__) = NULL;
+NTDLL_IMPORT_LIST(X)
+#undef X
+
+int nt_initialize() {
+  HMODULE ntdll;
+
+  ntdll = LoadLibraryW(L"ntdll.dll");
+  if (ntdll == NULL)
+    return -1;
+
+#define X(return_type, declarators, name, ...)                               \
+  name =                                                                     \
+      (return_type(declarators*)(__VA_ARGS__)) GetProcAddress(ntdll, #name); \
+  if (name == NULL)                                                          \
+    return -1;
+  NTDLL_IMPORT_LIST(X)
+#undef X
+
+  return 0;
+}
diff --git a/src/nt.h b/src/nt.h
index 245d0af..177b40d 100644
--- a/src/nt.h
+++ b/src/nt.h
@@ -4,6 +4,8 @@
 #include "ntstatus.h"
 #include "win.h"
 
+int nt_initialize();
+
 typedef struct _IO_STATUS_BLOCK {
   union {
     NTSTATUS Status;
@@ -16,15 +18,24 @@ typedef VOID(NTAPI* PIO_APC_ROUTINE)(PVOID ApcContext,
                                      PIO_STATUS_BLOCK IoStatusBlock,
                                      ULONG Reserved);
 
-typedef NTSTATUS(NTAPI* PNTDEVICEIOCONTROLFILE)(HANDLE FileHandle,
-                                                HANDLE Event,
-                                                PIO_APC_ROUTINE ApcRoutine,
-                                                PVOID ApcContext,
-                                                PIO_STATUS_BLOCK IoStatusBlock,
-                                                ULONG IoControlCode,
-                                                PVOID InputBuffer,
-                                                ULONG InputBufferLength,
-                                                PVOID OutputBuffer,
-                                                ULONG OutputBufferLength);
+#define NTDLL_IMPORT_LIST(X)        \
+  X(NTSTATUS,                       \
+    NTAPI,                          \
+    NtDeviceIoControlFile,          \
+    HANDLE FileHandle,              \
+    HANDLE Event,                   \
+    PIO_APC_ROUTINE ApcRoutine,     \
+    PVOID ApcContext,               \
+    PIO_STATUS_BLOCK IoStatusBlock, \
+    ULONG IoControlCode,            \
+    PVOID InputBuffer,              \
+    ULONG InputBufferLength,        \
+    PVOID OutputBuffer,             \
+    ULONG OutputBufferLength)
+
+#define X(return_type, declarators, name, ...) \
+  extern return_type(declarators* name)(__VA_ARGS__);
+NTDLL_IMPORT_LIST(X)
+#undef X
 
 #endif /* EPOLL_NT_H_ */