///////////////////////////////////////////////////////////////////////////////
//
//	(C) Copyright 1999 - 2000 Mark Roddy
//	All Rights Reserved
//
//	Hollis Technology Solutions
//	94 Dow Road
//	Hollis, NH 03049
//	info@hollistech.com
//
//	Synopsis: 
// 
//
//	Version Information:
//
//	$Header: /iphook/sys/driver/iphook.cpp 4     1/27/00 10:35p Markr $ 
//
///////////////////////////////////////////////////////////////////////////////
#include "iphookKrnl.h"

HTS_DEBUG_THIS_FILE


//
// local functions
//
NTSTATUS
	IpfStartHooking(PVOID context);

NTSTATUS 
	IpfStopHooking(PVOID context);

IP_HOOK_GLOBAL_DATA ipGlobal;


CPP_DRIVER_ENTRY(PDRIVER_OBJECT DriverObject,
                 PUNICODE_STRING RegistryPath)
{
	NTSTATUS Status = STATUS_SUCCESS;

	//
	// setup our standard entry points
	//
	DriverObject->MajorFunction[IRP_MJ_CREATE]         = IpHookCreate;
    DriverObject->MajorFunction[IRP_MJ_CLOSE]          = IpHookClose;
    DriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL] = IpHookDeviceControl;
	DriverObject->DriverUnload = IpHookUnload;
	//
	// before we get too deeply involved in this mess, lets see if the
	// ipfilter driver is actually installed.
	//
	Status = ipGlobal.initDevice();

	if (!NT_SUCCESS(Status)) {

		HtsDebugPrint(HTS_DEBUG_HIGH, "initDevice failed error %x\n", Status);
		return Status;
	}

	Status = ipGlobal.createThread();

	if (!NT_SUCCESS(Status)) {

		HtsDebugPrint(HTS_DEBUG_HIGH, "createThread failed error %x\n", Status);
		return Status;

	}  
	//
	// create a device object that applications can open to control the ip hook driver.
	//
	UNICODE_STRING deviceName;

	RtlInitUnicodeString(&deviceName, IPHOOK_DEV_NAME);

	PDEVICE_OBJECT ourDevice = NULL;

	Status = IoCreateDevice(
				DriverObject,
				0,
				&deviceName,
				IPHOOK_DEVICE_TYPE,
				0,
				FALSE, // well perhaps this should be true.
				&ourDevice);  

	if (!NT_SUCCESS(Status)) {

		//
		// oh well, what can we do?
		//
		HtsDebugPrint(HTS_DEBUG_HIGH, "IoCreateDevice failed status %x\n", Status);

		return Status;
	}

	UNICODE_STRING SymbolicLinkName;

	RtlInitUnicodeString(&SymbolicLinkName, IPHOOK_USER_DEV_NAME);

	//
	// create a symbolic link that exposes our device object to user space.
	//
	Status = IoCreateSymbolicLink(
				&SymbolicLinkName,
				&deviceName);

	if (!NT_SUCCESS(Status)) {
		//
		// we have to delete the device object here
		//
		IoDeleteDevice(ourDevice);

		HtsDebugPrint(HTS_DEBUG_HIGH, "IoCreateSymbolicLink failed status %x\n", Status);
	}

	return Status;
}

void
	IpHookUnload(PDRIVER_OBJECT DriverObject)
{
	HtsDebugPrint(HTS_DEBUG_LOW, "IpHookUnload\n");

	//
	HtsAssert(!ipGlobal.owned()); 
	//
	// delete our symbolic link
	//
	UNICODE_STRING SymbolicLinkName;

	RtlInitUnicodeString(&SymbolicLinkName, IPHOOK_USER_DEV_NAME);


	(void) IoDeleteSymbolicLink( &SymbolicLinkName );

	//
	// delete our device object(s)
	//
	while (DriverObject->DeviceObject) {
		
		IoDeleteDevice(DriverObject->DeviceObject);
	}

}

//
// Create
//
//	If we are going to support a read data interface then
//	we need to do something here to track operations.
//
//
NTSTATUS 
	IpHookCreate(PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
	HtsDebugPrint(HTS_DEBUG_LOW, "IpHookCreate\n");

	return STATUS_SUCCESS;
}

//
// Close
//
//	If we are going to support a read data interface then
//	we need to do something here to track operations.
//
//	On last close terminate hooking if active.
//
NTSTATUS 
	IpHookClose(PDEVICE_OBJECT DeviceObject, PIRP Irp)
{

	PVOID context = IoGetCurrentIrpStackLocation(Irp)->FileObject;

	HtsDebugPrint(HTS_DEBUG_LOW, "IpHookClose context %x\n", context);

	//
	// ok lets see if the caller owned the interface and do an implicit release if he did.
	//
	if (ipGlobal.IsOwner(context)) {
		//
		// bingo!
		//
		(void) IpfStopHooking(context);
	}


	return STATUS_SUCCESS;
}

//
// DeviceControl
//
//	Provide an interface to start Ip Hooking, stop Ip Hooking.
//
NTSTATUS 
	IpHookDeviceControl(PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
	HtsDebugPrint(HTS_DEBUG_LOW, "IpHookDeviceControl\n");
	//
	// get our stack location so we can decode this fellow
	//
	PIO_STACK_LOCATION IoStack = IoGetCurrentIrpStackLocation(Irp);

	//
	// get the control code.
	//
	ULONG controlCode = IoStack->Parameters.DeviceIoControl.IoControlCode;

	NTSTATUS Status = STATUS_SUCCESS;

	switch (controlCode) {

	case START_IP_HOOK:

		Status = IpfStartHooking(IoStack->FileObject);

		break;


	case STOP_IP_HOOK:

		Status = IpfStopHooking(IoStack->FileObject);

		break;

	case HOOK_THIS:

		//
		// see if we are hooking,
		// a stronger test would be to see
		// if this caller owns the interface
		// as in 
		//		ipGlobal.IsOwner(IoStack->FileObject);
		//
		if ( !ipGlobal.owned() ) {
			//
			// bag it
			//
			Status = STATUS_UNSUCCESSFUL;
			break;
		}

		{
			//
			// make sure the buffer makes sense, we will not
			// revalidate later.
			//
			PIPHOOK_BUFFER buffer = (PIPHOOK_BUFFER) MmGetSystemAddressForMdl(Irp->MdlAddress);

			if (!validIpHookBuffer(buffer)) {

				HtsDebugPrint(HTS_DEBUG_LOW, "Invalid Buffer Tag\n");

				return HtsIrpReturn(Irp, STATUS_DATA_ERROR);

			}

			ULONG size = sizeof(IPHOOK_BUFFER) + 
				((buffer->entries - 1) * sizeof(IPHOOK_DATA));

			if (size > IoStack->Parameters.DeviceIoControl.OutputBufferLength) {

				HtsDebugPrint(HTS_DEBUG_LOW, "Invalid buffer length\n");

				return HtsIrpReturn(Irp, STATUS_INFO_LENGTH_MISMATCH);

			}
		}

		IoMarkIrpPending(Irp);

		ipGlobal.queueRequest(Irp);

		return STATUS_PENDING;
		
	default:

		HtsDebugPrint(HTS_DEBUG_HIGH, "Unexpected control code %x\n", controlCode);
		//
		// we don't understand this so get rid of it.
		//		
		Status = STATUS_INVALID_DEVICE_REQUEST;
		break;
	}

	return HtsIrpReturn(Irp, Status);

}


//
// the hooker herself!
//
PF_FORWARD_ACTION 
	IpHookFilter(IN unsigned char *PacketHeader,
				 IN unsigned char *Packet, 
				 IN unsigned int PacketLength, 
				 IN unsigned int RecvInterfaceIndex, 
				 IN unsigned int SendInterfaceIndex, 
				 IN IPAddr RecvLinkNextHop, 
				 IN IPAddr SendLinkNextHop)
{
	PF_FORWARD_ACTION action = PF_PASS;

	__try {
		//
		// is this interface single threaded or multi threaded?
		//
		ULONG sequence = ipGlobal.getSequence();

		HtsDebugPrint(HTS_DEBUG_LOW, "IpHookFilter PacketHeader %x Packet %x PacketLength %x\n",
					  PacketHeader, Packet, PacketLength);

		UCHAR header0 = *PacketHeader; //see what this is		
		//
		// if iph_verlen & 0xf0 != 0x40 we have 
		// something other than an IPV4 packet.
		//
		// if iph_verlen & 0xf0 != 0x60 then not
		// only is this not an IPV4 packet, but it 
		// isn't an IPV6 packet either.
		//

		if ((header0 & 0xf0) == 0x40) {

			PIPHOOK_DATA iphookdata = ipGlobal.getBuffer();
			//
			// ok its IPv4 or reasonably close!
			//
			if (iphookdata) {

				ULARGE_INTEGER ticks;
				
				KeQueryTickCount((PLARGE_INTEGER)&ticks);

				iphookdata->tag = IPHOOK_DATA_TAG;
  
				iphookdata->sequence = sequence; // ? see above ?
				iphookdata->header = *(IPHeader *) PacketHeader;
				iphookdata->timestamp = ticks.QuadPart;
				iphookdata->dataLength = PacketLength;
				if (RecvInterfaceIndex != INVALID_PF_IF_INDEX ) {

					iphookdata->direction = 0;
					iphookdata->nextHop = RecvLinkNextHop;
					iphookdata->ifIndex = RecvInterfaceIndex;

				} else {

					iphookdata->direction = 1;
					iphookdata->nextHop = SendLinkNextHop;
					iphookdata->ifIndex = SendInterfaceIndex;

				}

			} else {

				HtsDebugPrint(HTS_DEBUG_LOW, "IpHookFilter: no buffer\n");
			}
		}
	}

	__except(HTS_EXCEPTION_FILTER ) {

		//
		// cross your fingers and keep going
		//
		HtsDebugPrint(HTS_DEBUG_HIGH, "IpHookFilter continuing from exception");

	}

	return action;
}


NTSTATUS
	IpfStartHooking(PVOID context)
{
	HtsDebugPrint(HTS_DEBUG_LOW, "IpfStartHooking context %x owner %x\n",
				  PsGetCurrentThread(),
				  ipGlobal.owner());

	if (FALSE == ipGlobal.setOwner(context)) {
		//
		// some other thread owns the interface,
		// bail out
		//
		HtsDebugPrint(HTS_DEBUG_LOW, "In Use\n");

		return STATUS_CONNECTION_IN_USE; 

	} 
	//
	// we own the interface! now try to start it up
	//
	PF_SET_EXTENSION_HOOK_INFO extensionHook = { IpHookFilter, };

	KEVENT event;

	KeInitializeEvent(&event, NotificationEvent, FALSE);

	IO_STATUS_BLOCK IoStatus;

	PIRP Irp = 
		IoBuildDeviceIoControlRequest(IOCTL_PF_SET_EXTENSION_POINTER,
									  ipGlobal.getIpDevice(),
									  (PVOID) &extensionHook,
									  sizeof(PF_SET_EXTENSION_HOOK_INFO),
									  NULL,
									  0,
									  FALSE,
									  &event,
									  &IoStatus);

	if (!Irp) {
		//
		// we need to free up the interface!
		//
		HtsDebugPrint(HTS_DEBUG_HIGH, "IoBuildDeviceIoControlRequest no IRP!\n");

		ipGlobal.freeOwner(context);

		return STATUS_INSUFFICIENT_RESOURCES;

	}

	//
	// ok now send our IRP to the filter driver
	//
	NTSTATUS Status = IoCallDriver(ipGlobal.getIpDevice(), Irp);

	if (Status == STATUS_PENDING) {

		NTSTATUS WaitStatus = KeWaitForSingleObject(&event, UserRequest, KernelMode, FALSE, NULL);

		if (WaitStatus 	!= STATUS_SUCCESS ) {
			//
			// bummer!
			//
			HtsDebugPrint(HTS_DEBUG_HIGH, "KeWaitForSingleObject failed status %x\n", WaitStatus);
			//
			// we need to free up the interface!
			//		
			ipGlobal.freeOwner(context);

			return WaitStatus;
		}
	}

	Status = IoStatus.Status;

	if (!NT_SUCCESS(Status)) {
		//
		// we need to free up the interface!
		//		
		ipGlobal.freeOwner(context);
	}

	HtsDebugPrint(HTS_DEBUG_LOW, "IpfStartHooking returns %x\n", Status);
	
	return Status;

}


//
// NB failure modes here other than the caller isn't the owning thread are
// truly bogus. Only a cursory attempt is made to deal with them. Basically the 
// strategy is to propogate the failure back to the caller but leave the 
// filter interface in place.
//
NTSTATUS
	IpfStopHooking(PVOID context)
{
	HtsDebugPrint(HTS_DEBUG_LOW, "IpfStopHooking context %x owner %x\n",
				  context,
				  ipGlobal.owner());

	//
	// see if this is even legitimate
	//
	if (ipGlobal.IsOwner(context) == FALSE) {

		return STATUS_INVALID_OWNER;
	}
	//
	// ok so it is us, and we is the thread, so we can only be here ONCE
	// so we are already serialized, so everything is groovey. Except other
	// threads could open the hook device after we release ownership,
	// and new reads could arrive (but not on this thread.)
	//
	//
	// we own the interface! now try to start it up
	//
	PF_SET_EXTENSION_HOOK_INFO extensionHook = { (PacketFilterExtensionPtr) NULL, };

	KEVENT event;

	KeInitializeEvent(&event, NotificationEvent, FALSE);

	IO_STATUS_BLOCK IoStatus;

	PIRP Irp = 
		IoBuildDeviceIoControlRequest(IOCTL_PF_SET_EXTENSION_POINTER,
									  ipGlobal.getIpDevice(),
									  (PVOID) &extensionHook,
									  sizeof(PF_SET_EXTENSION_HOOK_INFO),
									  NULL,
									  0,
									  FALSE,
									  &event,
									  &IoStatus);

	if (!Irp) {

		HtsDebugPrint(HTS_DEBUG_HIGH,  "IoBuildDeviceIoControlRequest No IRP\n");
		
		return STATUS_INSUFFICIENT_RESOURCES;

	}

	//
	// ok now send our IRP to the filter driver
	//
	NTSTATUS Status = IoCallDriver(ipGlobal.getIpDevice(), Irp);

	if (Status == STATUS_PENDING) {

		NTSTATUS WaitStatus = KeWaitForSingleObject(&event, UserRequest, KernelMode, FALSE, NULL);


		if (WaitStatus 	!= STATUS_SUCCESS ) {

			HtsDebugPrint(HTS_DEBUG_HIGH, "KeWaitForSingleObject failed status %x\n", WaitStatus);
			//
			// bummer!
			//
			return WaitStatus;
		}
	}

	Status = IoStatus.Status;

	if (NT_SUCCESS(Status)) {
		//
		// we need to free up the interface!
		//
		ipGlobal.freeOwner(context);
	}

	HtsDebugPrint(HTS_DEBUG_LOW, "IpfStopHooking return %x\n", Status);

	return Status;

}

//
// IP_HOOK_GLOBAL_DATA and associates, LLC
//

VOID
GlobalCancel(
    IN PDEVICE_OBJECT DeviceObject,
    IN PIRP Irp
    )
{
	KIRQL irql = Irp->CancelIrql;

	IoReleaseCancelSpinLock(irql);

	PVOID context;

	ipGlobal.lock(context);
	//
	// this irp is in our queue
	//
	RemoveEntryList(&Irp->Tail.Overlay.ListEntry);

	HtsIrpReturn(Irp, STATUS_CANCELLED);

	ipGlobal.unlock(context);

}
///////////////////////////////////////////////////////////////////////////////
// 
// Change History Log
//
// $Log: /iphook/sys/driver/iphook.cpp $
// 
// 4     1/27/00 10:35p Markr
// Prepare to release!
//
///////////////////////////////////////////////////////////////////////////////