// cl /Gz /D_X86_ /Od /TP /IG:\ddk\inc\ddk /IG:\ddk\inc\api myfault.c /link /driver:wdm /out:myfault.sys /entry:DriverEntry /subsystem:native /nodefaultlib g:\ddk\lib\wxp\i386\ntoskrnl.lib
#include <ntddk.h>
#include <string.h>
#include <stdio.h>
#include "ioctlcmd.h"

template<class Func>
void GetFunc(UNICODE_STRING& funcName, Func& fn)
{
	fn = static_cast<Func>(MmGetSystemRoutineAddress(&funcName));
}

VOID
InitHeadless( 
    PCHAR pMessageOut 
    )
{
	typedef BOOLEAN (NTAPI*pfnTermInit)(int enable);
    pfnTermInit termInit = NULL;
	UNICODE_STRING name = RTL_CONSTANT_STRING(L"HeadlessEnableTerminal");
	GetFunc(name, termInit);
	if(termInit)
	{
		BOOLEAN res = termInit(TRUE);
		sprintf(pMessageOut, "HeadlessEnableTerminal returned %d", res);
	}
	else
	{
		strcpy(pMessageOut, "HeadlessEnableTerminal not fouund");
	}
}

VOID ProcessDumpCommand(PCHAR pMessageOut)
{
	typedef void (NTAPI*pfnDump)();
	pfnDump dumpCommand = NULL;
	UNICODE_STRING name = RTL_CONSTANT_STRING(L"HeadlessProcessDumpCommand");
	GetFunc(name, dumpCommand);
	if(dumpCommand)
	{
		dumpCommand();
		strcpy(pMessageOut, "HeadlessProcessDumpCommand was called");
	}
	else
	{
		strcpy(pMessageOut, "HeadlessProcessDumpCommand not fouund");
	}
}

VOID QueryInfo(PCHAR pMessageOut)
{
	typedef void (NTAPI*pfnQueryInfo)(BOOLEAN* pEnabled, ULONG* pOtherData);
	pfnQueryInfo queryInfo = NULL;
	UNICODE_STRING name = RTL_CONSTANT_STRING(L"HeadlessQueryInformation");
	GetFunc(name, queryInfo);
	if(queryInfo)
	{
		BOOLEAN enabled = FALSE;
		ULONG otherData = 0;
		queryInfo(&enabled, &otherData);
		sprintf(pMessageOut, "HeadlessQueryInformation returned enabled=%d, otherData = %lu", enabled, otherData);
	}
	else
	{
		strcpy(pMessageOut, "HeadlessQueryInformation not fouund");
	}
}

VOID PutString(PCHAR pMessageOut)
{
	typedef void (NTAPI*pfnPutString)(const char* pMessage);
	pfnPutString putString = NULL;
	UNICODE_STRING name = RTL_CONSTANT_STRING(L"HeadlessPutString");
	GetFunc(name, putString);
	if(putString)
	{
		putString("Yay! I called HeadlessPutString, let's see if this does anything");
		strcpy(pMessageOut, "HeadlessPutString called");
	}
	else
	{
		strcpy(pMessageOut, "HeadlessPutString not found");
	}
}

//----------------------------------------------------------------------
//
// MyfaultDeviceControl
//
//----------------------------------------------------------------------
NTSTATUS  
MyfaultDeviceControl( 
    IN PFILE_OBJECT FileObject, 
    IN BOOLEAN Wait,
    IN PVOID InputBuffer, 
    IN ULONG InputBufferLength, 
    OUT char* OutputBuffer, 
    IN ULONG OutputBufferLength, 
    IN ULONG IoControlCode, 
    OUT PIO_STATUS_BLOCK IoStatus, 
    IN PDEVICE_OBJECT DeviceObject 
    ) 
{
    IoStatus->Status = STATUS_SUCCESS;
	IoStatus->Information = 0;
	*OutputBuffer = 0;
	char* pAfterHeader = OutputBuffer + sprintf(OutputBuffer, "HeadlessDevIoctl called with code %x\n", IoControlCode);
	switch ( IoControlCode ) 
	{
		case IOCTL_INIT_HEADLESS:
		{
			InitHeadless(pAfterHeader);
		}
		break;
		case IOCTL_HEADLESS_STRING:
		{
			PutString(pAfterHeader);
		}
		break;
		case IOCTL_HEADLESS_QUERY_INFO:
		{
			QueryInfo(pAfterHeader);
		}
		break;
		case IOCTL_HEADLESS_DUMP_COMMAND:
		{
			ProcessDumpCommand(pAfterHeader);
		}
		break;
		default:
		{
			IoStatus->Status = STATUS_ACCOUNT_DISABLED;
			break;
		}
	}
	IoStatus->Information = strlen(OutputBuffer) + 1;
	return IoStatus->Status;
}


//----------------------------------------------------------------------
//
// MyfaultDispatch
//
// In this routine we Myfault requests to our own device. The only 
// requests we care about handling explicitely are IOCTL commands that
// we will get from the GUI. We also expect to get Create and Close 
// commands when the GUI opens and closes communications with us.
//
//----------------------------------------------------------------------
NTSTATUS 
MyfaultDispatch( 
    IN PDEVICE_OBJECT DeviceObject, 
    IN PIRP Irp 
    )
{
	PIO_STACK_LOCATION      iosp;
	PVOID                   inputBuffer;
	PVOID                   outputBuffer;
	ULONG                   inputBufferLength;
	ULONG                   outputBufferLength;
	ULONG                   ioControlCode;
	NTSTATUS                status;

    //
    // Switch on the request type
    //
	iosp = IoGetCurrentIrpStackLocation (Irp);
	switch (iosp->MajorFunction) 
	{
		case IRP_MJ_CREATE:
		case IRP_MJ_CLOSE:
		{
			KdPrint(("HeadlessTest: Got open/close\n"));
			status = STATUS_SUCCESS;
		}
		break;
		case IRP_MJ_DEVICE_CONTROL:
		{
			inputBuffer        = Irp->AssociatedIrp.SystemBuffer;
			inputBufferLength  = iosp->Parameters.DeviceIoControl.InputBufferLength;
			outputBuffer       = Irp->AssociatedIrp.SystemBuffer;
			outputBufferLength = iosp->Parameters.DeviceIoControl.OutputBufferLength;
			ioControlCode      = iosp->Parameters.DeviceIoControl.IoControlCode;

			//
			// Special case: handle the IRP hang so as not to complete the IRP
			//
			status = MyfaultDeviceControl( iosp->FileObject, TRUE,
												   inputBuffer, inputBufferLength, 
												   (char*)outputBuffer, outputBufferLength,
												   ioControlCode, &Irp->IoStatus, 
												   DeviceObject );
		}
		break;
		default:
		{
			status = STATUS_NONE_MAPPED;
			break;        
		}
	}

    //
    // Complete the request
    //
    Irp->IoStatus.Status = status;
	IoCompleteRequest( Irp, IO_NO_INCREMENT );
	return status;
}


//----------------------------------------------------------------------
//
// MyfaultUnload
//
// Our job is done - time to leave.
//
//----------------------------------------------------------------------
VOID 
MyfaultUnload( 
    IN PDRIVER_OBJECT DriverObject 
    )
{
	UNICODE_STRING          deviceLinkUnicodeString = RTL_CONSTANT_STRING(L"\\DosDevices\\Headless");
DbgPrint(("HeadlessTest: Unloading\n"));

    //
	// Delete the symbolic link for our device
    //
	IoDeleteSymbolicLink( &deviceLinkUnicodeString );
DbgPrint(("HeadlessTest: Deleted the symbolic link\n"));

    //
	// Delete the device object
    //
	IoDeleteDevice( DriverObject->DeviceObject );
DbgPrint(("HeadlessTest: Deleted the device\n"));
}

extern "C"
//----------------------------------------------------------------------
//
// DriverEntry
//
// Installable driver initialization. Here we just set ourselves up.
//
//----------------------------------------------------------------------
NTSTATUS 
DriverEntry(
    IN PDRIVER_OBJECT DriverObject, 
    IN PUNICODE_STRING RegistryPath 
    )
{
    NTSTATUS                status;
    UNICODE_STRING          deviceNameUnicodeString = RTL_CONSTANT_STRING(L"\\Device\\Headless");
    UNICODE_STRING          deviceLinkUnicodeString = RTL_CONSTANT_STRING(L"\\DosDevices\\Headless");
    PDEVICE_OBJECT          interfaceDevice = NULL;
    ULONG                   startType, demandStart;

DbgPrint(("HeadlessTest: In DriverEntry\n"));

	MmPageEntireDriver(&DriverEntry);

    //
    // Create a named device object
    //
    status = IoCreateDevice ( DriverObject,
                                0,
                                &deviceNameUnicodeString,
                                FILE_DEVICE_MYFAULT,
                                0,
                                TRUE,
                                &interfaceDevice );
    if (NT_SUCCESS(status)) {

DbgPrint(("HeadlessTest: IoCreateDevice Success\n"));
	   //
	   // Create a symbolic link that the GUI can specify to gain access
	   // to this driver/device
	   //
	   status = IoCreateSymbolicLink (&deviceLinkUnicodeString,
										&deviceNameUnicodeString );
KdPrint(("HeadlessTest: IoCreateSymbolicLink returned 0x%x\n", status));
	   //
	   // Create dispatch points for all routines that must be Myfaultd
	   //
	   DriverObject->MajorFunction[IRP_MJ_CREATE]          =
       DriverObject->MajorFunction[IRP_MJ_CLOSE]           =
       DriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL]  = MyfaultDispatch;
	   DriverObject->DriverUnload                          = MyfaultUnload;
    }

    if (!NT_SUCCESS(status)) {
      
	   //
	   // Something went wrong, so clean up 
	   //
	   if( interfaceDevice ) {

           IoDeleteDevice( interfaceDevice );
       }
    }
	else
	{
		char buffer[80];
		InitHeadless(buffer);
	}
KdPrint(("HeadlessTest: In DriverEntry - returning 0x%x\n", status));
    return status;
}