Added wrapper generator for functions that require an import library.

 * generate wrapper functions that can be retrieved via `MmGetSystemRoutineAddress`
 * for now, only ZwTraceControl and ZwTraceEvent from ntdll are used

Signed-off-by: Toni Uhlig <matzeton@googlemail.com>

4 files changed, 163 insertions, 0 deletions

diff --git a/CRT/gen_wrapper.sh b/CRT/gen_wrapper.sh
new file mode 100755
index 0000000..565859b
--- /dev/null
+++ b/CRT/gen_wrapper.sh
@@ -0,0 +1,99 @@
+#!/usr/bin/env sh
+
+MYDIR="$(dirname ${0})"
+FN_FILE="${1:-${MYDIR}/ntdll_zw_functions.txt}"
+
+TYPEDEFS=""
+STATICS=""
+CURLINE=0
+while read -r line; do
+    CURLINE=$(expr ${CURLINE} + 1)
+    VALID=1
+
+    rtype=$(printf '%s\n' "${line}" | grep -oE '(NTSTATUS NTAPI|VOID NTAPI)')
+    if [ -z "${rtype}" ]; then
+        printf '%s\n' "Line ${CURLINE}: Missing return value of either type 'NTSTATUS NTAPI' or 'VOID NTAPI'." >&2
+        VALID=0
+    fi
+
+    fnname=$(printf '%s\n' "${line}" | grep -oE 'Zw[^ (]*')
+    if [ -z "${fnname}" ]; then
+        printf '%s\n' "Line ${CURLINE}: Missing function name." >&2
+        VALID=0
+    fi
+
+    fnsig=$(printf '%s\n' "${line}" | grep -oE '\([^;]*')
+    if [ -z "${fnsig}" ]; then
+        printf '%s\n' "Line ${CURLINE}: Missing function signature." >&2
+        VALID=0
+    fi
+
+    param_names=$(printf '%s\n' "${fnsig}" | tr -d '()' | sed 's/\([^,]*\)/\1\n/g' | grep -oE '[^ ]*$')
+    if [ -z "${param_names}" ]; then
+        printf '%s\n' "Line ${CURLINE}: Could not parse function parameters." >&2
+        VALID=0
+    fi
+    params=""
+    for param in ${param_names}; do
+        if [ -z "${param}" ]; then
+            printf '%s\n' "Line ${CURLINE}: Invalid parameter found. Please re-check regex'es used." >&2
+            VALID=0
+        fi
+        params="${params}${param}, "
+    done
+    params=$(printf '%s\n' "${params}" | sed 's/^\(.*\), $/\1/g')
+    if [ -z "${params}" ]; then
+        printf '%s\n' "Line ${CURLINE}: Parameters empty. Please re-check regex'es used." >&2
+        VALID=0
+    fi
+
+    if [ ${VALID} -eq 1 ]; then
+        TYPEDEFS="${TYPEDEFS}\ntypedef ${rtype} (*${fnname}_t) ${fnsig};"
+        STATICS="${STATICS}\nstatic ${fnname}_t _${fnname} = NULL;"
+        INITS=$(cat <<EOF
+${INITS}
+    RtlInitUnicodeString(&fnName, L"${fnname}");
+    _${fnname} = MmGetSystemRoutineAddress(&fnName);
+    if (_${fnname} == NULL)
+    {
+        DbgPrint("%s\\\n", "System routine ${fnname} not found.");
+        retval++;
+    }
+EOF
+        )
+        WRAPPERS=$(cat <<EOF
+${WRAPPERS}
+
+${rtype} ${fnname} ${fnsig}
+{
+    if (_${fnname} == NULL)
+        return STATUS_PROCEDURE_NOT_FOUND;
+
+    return _${fnname} (${params});
+}
+EOF
+        )
+    fi
+done < "${FN_FILE}"
+
+cat <<EOF
+/* This file was auto generated by $(basename ${0}) */
+#include <ntddk.h>
+
+EOF
+echo "${TYPEDEFS}"
+echo "${STATICS}"
+cat <<EOF
+
+int __cdecl $(basename -a -s '.txt' ${FN_FILE}) (void)
+{
+    int retval = 0;
+    UNICODE_STRING fnName;
+EOF
+echo "${INITS}"
+cat <<EOF
+
+    return retval;
+}
+EOF
+echo "${WRAPPERS}"
diff --git a/CRT/kcrt.c b/CRT/kcrt.c
index 8e96ca6..fd57b29 100644
--- a/CRT/kcrt.c
+++ b/CRT/kcrt.c
@@ -12,6 +12,7 @@ extern void (*__CTOR_LIST__)();
 extern void (*__DTOR_LIST__)();
 extern NTSTATUS __cdecl DriverEntry(struct _DRIVER_OBJECT * DriverObject, PUNICODE_STRING RegistryPath);
 extern void __cdecl DriverUnload(struct _DRIVER_OBJECT * DriverObject);
+extern int __cdecl ntdll_zw_functions(void);
 
 DRIVER_INITIALIZE __cdecl _CRT_DriverEntry;
 DRIVER_UNLOAD __cdecl _CRT_DriverUnload;
@@ -274,6 +275,13 @@ NTSTATUS __cdecl _CRT_DriverEntry(struct _DRIVER_OBJECT * DriverObject, PUNICODE
 
     KCRT_OnDriverEntry();
 
+    int zw_retval = ntdll_zw_functions();
+    if (zw_retval != 0)
+    {
+        DbgPrint("ERROR: Missing %d required system routines.\n", zw_retval);
+        return STATUS_NOT_IMPLEMENTED;
+    }
+
     retval = DriverEntry(DriverObject, RegistryPath);
 
     /* support for service stopping and CRT de-init */
@@ -291,3 +299,8 @@ void __cdecl _disable(void)
 {
     __asm__ __volatile__("cli");
 }
+
+void * __cdecl _AddressOfReturnAddress(void)
+{
+    return __builtin_extract_return_addr(__builtin_return_address(0));
+}
diff --git a/CRT/ntdll_zw_functions.c b/CRT/ntdll_zw_functions.c
new file mode 100644
index 0000000..d2df299
--- /dev/null
+++ b/CRT/ntdll_zw_functions.c
@@ -0,0 +1,49 @@
+/* This file was auto generated by gen_wrapper.sh */
+#include <ntddk.h>
+
+
+typedef NTSTATUS NTAPI (*ZwTraceControl_t) (_In_ ULONG FunctionCode, PVOID InBuffer, _In_ ULONG InBufferLen, PVOID OutBuffer, _In_ ULONG OutBufferLen, _Out_ PULONG ReturnLength);
+typedef NTSTATUS NTAPI (*ZwTraceEvent_t) (_In_ HANDLE TraceHandle, _In_ ULONG Flags, _In_ ULONG FieldSize, _In_ PVOID Fields);
+
+static ZwTraceControl_t _ZwTraceControl = NULL;
+static ZwTraceEvent_t _ZwTraceEvent = NULL;
+
+int __cdecl ntdll_zw_functions (void)
+{
+    int retval = 0;
+    UNICODE_STRING fnName;
+
+    RtlInitUnicodeString(&fnName, L"ZwTraceControl");
+    _ZwTraceControl = MmGetSystemRoutineAddress(&fnName);
+    if (_ZwTraceControl == NULL)
+    {
+        DbgPrint("%s\n", "System routine ZwTraceControl not found.");
+        retval++;
+    }
+    RtlInitUnicodeString(&fnName, L"ZwTraceEvent");
+    _ZwTraceEvent = MmGetSystemRoutineAddress(&fnName);
+    if (_ZwTraceEvent == NULL)
+    {
+        DbgPrint("%s\n", "System routine ZwTraceEvent not found.");
+        retval++;
+    }
+
+    return retval;
+}
+
+
+NTSTATUS NTAPI ZwTraceControl (_In_ ULONG FunctionCode, PVOID InBuffer, _In_ ULONG InBufferLen, PVOID OutBuffer, _In_ ULONG OutBufferLen, _Out_ PULONG ReturnLength)
+{
+    if (_ZwTraceControl == NULL)
+        return STATUS_PROCEDURE_NOT_FOUND;
+
+    return _ZwTraceControl (FunctionCode, InBuffer, InBufferLen, OutBuffer, OutBufferLen, ReturnLength);
+}
+
+NTSTATUS NTAPI ZwTraceEvent (_In_ HANDLE TraceHandle, _In_ ULONG Flags, _In_ ULONG FieldSize, _In_ PVOID Fields)
+{
+    if (_ZwTraceEvent == NULL)
+        return STATUS_PROCEDURE_NOT_FOUND;
+
+    return _ZwTraceEvent (TraceHandle, Flags, FieldSize, Fields);
+}
diff --git a/CRT/ntdll_zw_functions.txt b/CRT/ntdll_zw_functions.txt
new file mode 100644
index 0000000..2581487
--- /dev/null
+++ b/CRT/ntdll_zw_functions.txt
@@ -0,0 +1,2 @@
+NTSYSCALLAPI NTSTATUS NTAPI ZwTraceControl (_In_ ULONG FunctionCode, PVOID InBuffer, _In_ ULONG InBufferLen, PVOID OutBuffer, _In_ ULONG OutBufferLen, _Out_ PULONG ReturnLength)
+NTSYSCALLAPI NTSTATUS NTAPI ZwTraceEvent (_In_ HANDLE TraceHandle, _In_ ULONG Flags, _In_ ULONG FieldSize, _In_ PVOID Fields)