//
// Implements an extensible system call monitoring interface
//
// skape
// mmiller@hick.org
// 5/2007
//
#include <windows.h>
#include <stdlib.h>
#include <stdio.h>
#include <wchar.h>
#include "..\common\winstrace.h"
#include "tables.h"

//
// Maximum of 65536 buckets in the thread handle hash
//
#define THREAD_HANDLE_HASH_BUCKETS 0x10000

//
// Initialize the doubly linked list
//
VOID
FORCEINLINE
InitializeListHead(
    IN PLIST_ENTRY ListHead
    )
{
    ListHead->Flink = ListHead->Blink = ListHead;
}

//
// Insert an entry to the head of the list
//
VOID
FORCEINLINE
InsertHeadList(
    IN PLIST_ENTRY ListHead,
    IN PLIST_ENTRY Entry
    )
{
    PLIST_ENTRY Flink;

    Flink = ListHead->Flink;
    Entry->Flink = Flink;
    Entry->Blink = ListHead;
    Flink->Blink = Entry;
    ListHead->Flink = Entry;
}

//
// Remove an entry from the list
//
BOOLEAN
FORCEINLINE
RemoveEntryList(
    IN PLIST_ENTRY Entry
    )
{
    PLIST_ENTRY Blink;
    PLIST_ENTRY Flink;

    Flink = Entry->Flink;
    Blink = Entry->Blink;
    Blink->Flink = Flink;
    Flink->Blink = Blink;
    return (BOOLEAN)(Flink == Blink);
}

typedef VOID (*SYSTEM_CALL_NOTIFY_ROUTINE)(
		__in PSYSTEM_CALL_EVENT Event);
typedef VOID (*SYSTEM_CALL_RETURN_NOTIFY_ROUTINE)(
		__in PSYSTEM_CALL_RETURN_EVENT Event);

//
// Plugin structure
//
typedef struct _TRACE_PLUGIN
{
	//
	// List entry
	//
	LIST_ENTRY                 Entry;

	//
	// Plugin's user context
	//
	PVOID                      Context;

	//
	// The path to load the file from
	//
	PWSTR                      FilePath;

	//
	// The plugin module
	//
	HMODULE                    Module;

	//
	// Called when a system call occurs
	//
	SYSTEM_CALL_NOTIFY_ROUTINE OnSystemCall;
	//
	// Called when a system call returns
	//
	SYSTEM_CALL_RETURN_NOTIFY_ROUTINE OnSystemCallReturn;

} TRACE_PLUGIN, *PTRACE_PLUGIN;

//
// Cache entry that stores target thread information
//
typedef struct _THREAD_HANDLE_CACHE_ENTRY
{
	//
	// The list entry
	//
	LIST_ENTRY        Entry;

	//
	// The thread's unique identifier
	//
	ULONG             ThreadId;

	//
	// The thread's handle
	//
	HANDLE            ThreadHandle;

	//
	// The last system call to occur on this thread
	//
	SYSTEM_CALL_EVENT LastSystemCall;

	//
	// The name of the last system call to occur on this thread
	//
	WCHAR             LastSystemCallName[128];

	//
	// Thread-specific context to store arguments in
	//
	ULONG_PTR         StackContext[MAX_STACK_ELEMENTS];

} THREAD_HANDLE_CACHE_ENTRY, *PTHREAD_HANDLE_CACHE_ENTRY;

////
//
// Globals
//
////

//
// The handle to the process that is being traced
//
static HANDLE ProcessHandle = NULL;
//
// Flag that controls whether the polling loop terminates or not
//
static BOOL TerminateSystemCallLoop = FALSE;
//
// The list of plugins
//
static LIST_ENTRY PluginList;
//
// Cache hash for thread handles
//
static LIST_ENTRY ThreadHandleHash[THREAD_HANDLE_HASH_BUCKETS];

////
//
// Internal functions
//
////

//
// Acquires a handle to a thread, potentially re-using an existing handle from
// the thread handle cache
//
static PTHREAD_HANDLE_CACHE_ENTRY AcquireThreadHandle(
		__in ULONG ThreadId)
{
	PTHREAD_HANDLE_CACHE_ENTRY CacheEntry;
	PLIST_ENTRY                BucketEntry = &ThreadHandleHash[(ThreadId >> 2) % THREAD_HANDLE_HASH_BUCKETS];
	PLIST_ENTRY                CurrentEntry;
	HANDLE                     ThreadHandle = NULL;

	//
	// Look for a cached handle to this thread
	//
	for (CurrentEntry  = BucketEntry->Flink;
	     CurrentEntry != BucketEntry;
	     CurrentEntry  = CurrentEntry->Flink)
	{
		CacheEntry = (PTHREAD_HANDLE_CACHE_ENTRY)CurrentEntry;

		if (CacheEntry->ThreadId == ThreadId)
		{
			ThreadHandle = CacheEntry->ThreadHandle;
			break;
		}
	}

	//
	// If we failed to find a thread handle in the cache, then acquire one and
	// store it in the cache
	//
	if (!ThreadHandle)
	{
		ThreadHandle = OpenThread(
			THREAD_SET_CONTEXT | THREAD_GET_CONTEXT,
			FALSE,
			ThreadId);

		if (ThreadHandle)
		{
			CacheEntry = (PTHREAD_HANDLE_CACHE_ENTRY)malloc(sizeof(THREAD_HANDLE_CACHE_ENTRY));

			if (CacheEntry)
			{
				CacheEntry->ThreadId     = ThreadId;
				CacheEntry->ThreadHandle = ThreadHandle;

				InsertHeadList(
						BucketEntry,
						&CacheEntry->Entry);
			}
			//
			// Cache allocation failed, close the handle and pretend like we
			// weren't able to acquire one
			//
			else
			{
				CloseHandle(ThreadHandle);

				ThreadHandle = NULL;
			}
		}
	}

	return CacheEntry;
}

//
// Remove the entry associated with the supplied thread identifier
//
static VOID RemoveThreadHandleFromCache(
		__in ULONG ThreadId)
{
	PLIST_ENTRY BucketEntry = &ThreadHandleHash[(ThreadId >> 2) % THREAD_HANDLE_HASH_BUCKETS];
	PLIST_ENTRY CurrentEntry;
	HANDLE      ThreadHandle = NULL;

	for (CurrentEntry  = BucketEntry->Flink;
	     CurrentEntry != BucketEntry;
	     CurrentEntry  = CurrentEntry->Flink)
	{
		PTHREAD_HANDLE_CACHE_ENTRY CacheEntry = (PTHREAD_HANDLE_CACHE_ENTRY)CurrentEntry;

		if (CacheEntry->ThreadId == ThreadId)
		{
			CloseHandle(
					CacheEntry->ThreadHandle);

			RemoveEntryList(
					CurrentEntry);

			free(CacheEntry);
			break;
		}
	}
}

//
// Calculates the hash code of the system call name
//
static ULONG GetSystemCallHash(
		__in PWSTR SystemCallName)
{
	PWSTR CurrentCharacter = SystemCallName;
	ULONG HashCode = 0;

	while (*CurrentCharacter)	
	{
		HashCode += *CurrentCharacter++;
		HashCode  = (HashCode << 8) | ((HashCode >> 24) & 0xff);
	}

	return HashCode;
}

//
// Notifies all registered plugins of a system call
//
static VOID NotifyPluginsOfSystemCall(
		__in ULONG ThreadId,
		__in PTHREAD_HANDLE_CACHE_ENTRY CacheEntry,
		__in PCONTEXT ThreadContext)
{
	PSYSTEM_CALL_DESCRIPTOR Descriptor;
	PSYSTEM_CALL_EVENT      Event = &CacheEntry->LastSystemCall;
	PLIST_ENTRY             CurrentEntry;
	ULONG                   NumberOfArguments;
	BOOL                    Suppress = FALSE;

	//
	// Read the stack context at the time that the system call is issued
	//
	if (!ReadProcessMemory(
			ProcessHandle,
			(PVOID)ThreadContext->Esp,
			CacheEntry->StackContext,
			sizeof(CacheEntry->StackContext),
			NULL))
		return;

	//
	// Initialize the event's fields
	//
	Event->SystemCallNumber = ThreadContext->Eax;
	Event->ThreadId         = ThreadId;
	Event->ThreadHandle     = CacheEntry->ThreadHandle;
	Event->ThreadContext    = ThreadContext;
	Event->ProcessHandle    = ProcessHandle;
	Event->Arguments        = &CacheEntry->StackContext[2];

	//
	// Calculate the number of arguments that were passed in, we assume that our
	// address space is the same as the target
	//
	if ((*(PULONG)CacheEntry->StackContext[0] & 0xff) == 0xc3)
		Event->NumberOfArguments = 0;
	else
		Event->NumberOfArguments = (*(PULONG)CacheEntry->StackContext[0] >> 10) & 0x0fff;

	if (Event->NumberOfArguments > MAX_STACK_ELEMENTS - 2)
		Event->NumberOfArguments = MAX_STACK_ELEMENTS - 2;

	//
	// Figure out what system call it is
	//
	Descriptor = GetSystemCallDescriptor(
			Event->SystemCallNumber,
			&Event->Suppressed);

	//
	// If we shouldn't display this call, then don't do it.
	//
	if (Event->Suppressed)
		return;

	//
	// Get the system call name
	//
	if ((!Descriptor) || 
	    (!Descriptor->Name))
	{
		swprintf(
				CacheEntry->LastSystemCallName, 
				L"syscall_%x",
				Event->SystemCallNumber);

		Event->SystemCallName = CacheEntry->LastSystemCallName;
		Event->SystemCallHash = 0;
	}
	else
	{
		Event->SystemCallName = Descriptor->Name;
		Event->SystemCallHash = GetSystemCallHash(Descriptor->Name);
	}

	//
	// Now, notify each of the plugins of this event
	//
	for (CurrentEntry  = PluginList.Flink;
	     CurrentEntry != &PluginList;
	     CurrentEntry  = CurrentEntry->Flink)
	{
		PTRACE_PLUGIN Plugin = (PTRACE_PLUGIN)CurrentEntry;

		Event->PluginContext = Plugin->Context;

		if (Plugin->OnSystemCall)
			Plugin->OnSystemCall(Event);
	}
}

//
// Notify plugins of the system call return event
//
static VOID NotifyPluginsOfSystemCallReturn(
		__in PSYSTEM_CALL_RETURN_EVENT Event)
{
	PLIST_ENTRY CurrentEntry;

	//
	// Don't display the return value if this was suppressed
	//
	if (Event->SystemCall->Suppressed)
		return;

	//
	// Now, notify each of the plugins of this event
	//
	for (CurrentEntry  = PluginList.Flink;
	     CurrentEntry != &PluginList;
	     CurrentEntry  = CurrentEntry->Flink)
	{
		PTRACE_PLUGIN Plugin = (PTRACE_PLUGIN)CurrentEntry;

		if (Plugin->OnSystemCallReturn)
			Plugin->OnSystemCallReturn(Event);
	}
}

//
// Handles a system call intercept
//
static DWORD HandleSystemCall(
		__in LPDEBUG_EVENT DebugEvent,
		__in PEXCEPTION_RECORD ExceptionRecord)
{
	PTHREAD_HANDLE_CACHE_ENTRY CacheEntry;
	CONTEXT                    ThreadContext;
	DWORD                      ContinueStatus = DBG_EXCEPTION_NOT_HANDLED;

	do
	{
		//
		// Get a handle to thread
		//
		CacheEntry = AcquireThreadHandle(
				DebugEvent->dwThreadId);

		if (!CacheEntry)
		{
			wprintf(L"AcquireThreadHandle failed, %lu.\n",
					GetLastError());
			break;
		}

		//
		// Get the thread's context
		//
		ThreadContext.ContextFlags = CONTEXT_FULL;

		if (!GetThreadContext(
				CacheEntry->ThreadHandle,
				&ThreadContext))
		{
			wprintf(L"GetThreadContext failed, %lu.\n",
					GetLastError());
			break;
		}

		//
		// Set EDX to ESP and progress EIP to the sysenter instruction
		//
		ThreadContext.Edx = ThreadContext.Esp;
		ThreadContext.Eip = ThreadContext.Eip + 1;

		//
		// Display the system call information
		//
		NotifyPluginsOfSystemCall(
				DebugEvent->dwThreadId,
				CacheEntry,
				&ThreadContext);

		//
		// Restore the thread context
		//
		if (!SetThreadContext(
				CacheEntry->ThreadHandle,
				&ThreadContext))
		{
			wprintf(L"SetThreadContext failed, %lu.\n",
					GetLastError());
			break;
		}

		//
		// Success, let's allow the thread to continue executing
		//
		ContinueStatus = DBG_CONTINUE;

	} while (0);

	return ContinueStatus;
}

//
// Handles the system call return intercept
//
static DWORD HandleSystemCallReturn(
		__in LPDEBUG_EVENT DebugEvent,
		__in PEXCEPTION_RECORD ExceptionRecord)
{
	PTHREAD_HANDLE_CACHE_ENTRY CacheEntry;
	SYSTEM_CALL_RETURN_EVENT   Event;
	CONTEXT                    ThreadContext;
	DWORD                      ContinueStatus = DBG_EXCEPTION_NOT_HANDLED;

	do
	{
		//
		// Get a handle to thread
		//
		CacheEntry = AcquireThreadHandle(
				DebugEvent->dwThreadId);

		if (!CacheEntry)
		{
			wprintf(L"AcquireThreadHandle failed, %lu.\n",
					GetLastError());
			break;
		}

		//
		// Get the thread's context
		//
		ThreadContext.ContextFlags = CONTEXT_FULL;

		if (!GetThreadContext(
				CacheEntry->ThreadHandle,
				&ThreadContext))
		{
			wprintf(L"GetThreadContext failed, %lu.\n",
					GetLastError());
			break;
		}

		Event.Status     = (LONG)ThreadContext.Eax;
		Event.SystemCall = &CacheEntry->LastSystemCall;

		//
		// Display the system call information
		//
		NotifyPluginsOfSystemCallReturn(
				&Event);

		//
		// Success, let's allow the thread to continue executing
		//
		ContinueStatus = DBG_CONTINUE;

	} while (0);

	return ContinueStatus;
}

//
// Handles the exception by checking to see if it occurred at an expected
// location
//
static DWORD HandleException(
		__in LPDEBUG_EVENT DebugEvent)
{
	PEXCEPTION_RECORD ExceptionRecord  = &DebugEvent->u.Exception.ExceptionRecord;
	ULONG_PTR         ExceptionAddress = (ULONG_PTR)ExceptionRecord->ExceptionAddress;
	DWORD             ContinueStatus   = DBG_EXCEPTION_NOT_HANDLED;

	if (ExceptionRecord->ExceptionCode == EXCEPTION_BREAKPOINT)
		ContinueStatus = ActiveSystemCallTables->FilterBreakpoint(
				ExceptionAddress,
				DebugEvent,
				ExceptionRecord);
	else
	{
		if (DebugEvent->u.Exception.dwFirstChance)
			ContinueStatus = DBG_CONTINUE;
		else
		{
			wprintf(L"Unhandled exception %p @ %p\n",
					ExceptionRecord->ExceptionCode,
					ExceptionAddress);

			TerminateSystemCallLoop = TRUE;
		}
	}

	return ContinueStatus;
}

//
// Monitor for debug events
//
static VOID WatchDebugEvents()
{
	DEBUG_EVENT DebugEvent;
	DWORD       ContinueStatus;
	BOOL        ProcessTerminated = FALSE;

	//
	// Enable system call traps
	//
	if ((!ActiveSystemCallTables->HookNativeSystemCalls(TRUE)) ||
	    (!ActiveSystemCallTables->HookNativeSystemCallReturns(TRUE)))
	{
		wprintf(L"Failed to hook system call routines, %lu.\n",
				GetLastError());
		return;
	}

	//
	// Keep looping until a ctrl C occurs
	//
	while (!TerminateSystemCallLoop)
	{
		//
		// Wait for our next debug event
		//
		if (!WaitForDebugEvent(
				&DebugEvent,
				1000))
		{
			if (GetLastError() == ERROR_SEM_TIMEOUT)
				continue;

			wprintf(L"WaitForDebugEvent failed, %lu.\n",
					GetLastError());
			break;
		}

		switch (DebugEvent.dwDebugEventCode)
		{
			case EXCEPTION_DEBUG_EVENT:
				ContinueStatus = HandleException(
						&DebugEvent);
				break;
			case EXIT_THREAD_DEBUG_EVENT:
				RemoveThreadHandleFromCache(
						DebugEvent.dwThreadId);
				break;
			case EXIT_PROCESS_DEBUG_EVENT:
				TerminateSystemCallLoop = 
				ProcessTerminated       = TRUE;
				break;
			default:
				ContinueStatus = DBG_EXCEPTION_NOT_HANDLED;
				break;
		}

		//
		// Proceed by either handling or not handling this exception
		//
		ContinueDebugEvent(
				DebugEvent.dwProcessId,
				DebugEvent.dwThreadId,
				ContinueStatus);
	}

	//
	// Enable system call traps
	//
	if ((!ProcessTerminated) && 
	    ((!ActiveSystemCallTables->HookNativeSystemCalls(FALSE)) ||
	     (!ActiveSystemCallTables->HookNativeSystemCallReturns(FALSE))))
	{
		wprintf(L"Failed to unhook system call routines, %lu.\n",
				GetLastError());
		return;
	}
}

//
// Handle the control C event
//
static BOOL WINAPI ConsoleCtrlHandler(
		__in DWORD ControlType)
{
	if (ControlType == CTRL_C_EVENT)
	{
		TerminateSystemCallLoop = TRUE;
		return TRUE;
	}
	else
		return FALSE;
}

//
// Initialize global state
//
static VOID InitializeGlobals()
{
	ULONG Index;

	InitializeListHead(&PluginList);

	//
	// Initialize the list entries in the thread handle hash
	//
	for (Index = 0;
	     Index < THREAD_HANDLE_HASH_BUCKETS;
	     Index++)
		InitializeListHead(ThreadHandleHash + Index);

	//
	// We force USER32 to be loaded so that we can access it whenever win32k
	// system calls are issued.
	//
	LoadLibrary("user32");

	//
	// Get the system call tables to use for this version of Windows
	//
	DetermineActiveSystemCallTables();
}

int _cdecl wmain(ULONG argc, WCHAR **argv)
{
	PTRACE_PLUGIN Plugin;
	PVOID         (*Initialize)();
	ULONG         ProcessId = 0;
	PWSTR         ExecutableString = NULL;
	ULONG         Index = 1;
	BOOL          Attached = FALSE;

	do
	{
		//
		// Initialize global state
		//
		InitializeGlobals();

		//
		// Usage
		//
		if (argc == 1)
		{
			wprintf(L"Usage: %s COMMAND OPTIONS ...\n", argv[0]);
			wprintf(
				L"\n"
				L"Commands:\n\n"
				L"\t/p arg\tProcess ID to attach to\n"
				L"\t/e arg\tExecutes the specified application\n"
				L"\nOptions:\n\n"
				L"\t/g\tSuppress GDI system calls\n"
				L"\t/l arg\tPlugin path to load\n"
				);
			break;
		}

		//
		// Parse command line arguments
		//
		while (Index < argc)
		{
			ULONG Increment = 1;

			if (argv[Index][0] != L'/')
			{
				Index++;
				continue;
			}

			switch (argv[Index][1])
			{
				//
				// Handles /p <pid>
				//
				case L'p':
					if (Index + 1 >= argc)
						break;

					ProcessId = _wtoi(argv[Index + 1]);
					Increment = 2;
					break;

				case L'e':
					if (Index + 1 >= argc)
						break;

					ExecutableString = argv[Index + 1];
					Increment = 2;
					break;

				//
				// Handles /g
				//
				case L'g':
					DisableGdiCalls = TRUE;
					break;
				//
				// Handles /l <file path>
				//
				case L'l':
					if (Index + 1 >= argc)
						break;

					// 
					// Allocate and initialize the plugin
					//
					Plugin = (PTRACE_PLUGIN)malloc(sizeof(TRACE_PLUGIN));

					if (!Plugin)
						break;

					Plugin->FilePath = argv[Index + 1];

					//
					// Load the plugin
					//
					Plugin->Module = LoadLibraryW(Plugin->FilePath);

					if (!Plugin->Module)
					{
						wprintf(L"Failed to load plugin from %s (%lu)\n", 
								Plugin->FilePath,
								GetLastError());

						free(Plugin);
						break;
					}

					//
					// Resolve exports
					//
					Initialize = (PVOID (*)())GetProcAddress(
							Plugin->Module,
							"_InitializePlugin@0");

					Plugin->OnSystemCall = (SYSTEM_CALL_NOTIFY_ROUTINE)GetProcAddress(
							Plugin->Module,
							"_OnSystemCall@4");
					Plugin->OnSystemCallReturn = (SYSTEM_CALL_RETURN_NOTIFY_ROUTINE)GetProcAddress(
							Plugin->Module,
							"_OnSystemCallReturn@4");

					if (Initialize)
						Plugin->Context = Initialize();

					//
					// Insert the plugin into the list
					//
					InsertHeadList(
							&PluginList,
							(PLIST_ENTRY)Plugin);
					Increment = 2;
					break;
				default: 
					break;
			}

			Index += Increment;
		}

		if ((!ProcessId) &&
		    (!ExecutableString))
		{
			wprintf(L"A process identifier or executable must be specified.\n");
			break;
		}

		//
		// If they specified a process identifier, then we'll attach.
		//
		if (ProcessId)
		{
			//
			// Acquire a handle to the process
			//
			ProcessHandle = OpenProcess(
					PROCESS_VM_OPERATION | PROCESS_VM_WRITE | PROCESS_VM_READ | PROCESS_DUP_HANDLE,
					FALSE,
					ProcessId);

			if (!ProcessHandle)
			{
				wprintf(L"OpenProcess failed, %lu.\n",
						GetLastError());
				break;
			}

			//
			// Attach to the target process
			//
			if (!DebugActiveProcess(
					ProcessId))
			{
				wprintf(L"DebugActiveProcess failed, %lu.\n",
						GetLastError());
				break;
			}

		}
		//
		// Otherwise, if they specified an executable string, then spawn the
		// process
		//
		else
		{
			PROCESS_INFORMATION Pi;
			STARTUPINFOW        Si = { 0 };

			Si.cb = sizeof(Si);

			if (!CreateProcessW(
					NULL,
					ExecutableString,
					NULL,
					NULL,
					FALSE,
					DEBUG_PROCESS | DEBUG_ONLY_THIS_PROCESS,
					NULL,
					NULL,
					&Si,
					&Pi))
			{
				wprintf(L"CreateProcess failed, %lu.\n",
						GetLastError());
				break;
			}

			if (Pi.hThread)
				CloseHandle(Pi.hThread);

			ProcessHandle = Pi.hProcess;
		}
			
		Attached = TRUE;

		//
		// Register the handler to be notified when a ctrl-c event occurs
		//
		if (!SetConsoleCtrlHandler(
				ConsoleCtrlHandler,
				TRUE))
		{
			wprintf(L"SetConsoleCtrlHandler failed, %lu.\n",
					GetLastError());
			break;
		}

		//
		// Start watching debug events for the specified process identifier
		//
		WatchDebugEvents();

		//
		// Flag that we've succeeded
		//
		SetLastError(
				ERROR_SUCCESS);

	} while (0);

	//
	// If we had attached to the process, then now we detach
	//
	if (Attached)
		DebugActiveProcessStop(
				ProcessId);
	if (ProcessHandle)
		CloseHandle(ProcessHandle);

	return GetLastError();
}