/* Copyright (c) 2002-2005 Vladislav Goncharov.
*
* Redistribution and use in source forms, with and without modification,
* are permitted provided that this entire comment appears intact.
*
* Redistribution in binary form may occur without any restrictions.
*
* This software is provided ``AS IS'' without any warranties of any kind.
*/

// -*- mode: C++; tab-width: 4; indent-tabs-mode: nil -*- (for GNU Emacs)
//
// $Id: tdi_fw.c,v 1.13 2003/09/04 15:20:09 dev Exp $

/*
* TDI-based open source personal firewall (TdiFw)
*/

#include <ntddk.h>
#include <tdikrnl.h>
#include "sock.h"

#include "conn_state.h"
#include "dispatch.h"
#include "filter.h"
#include "memtrack.h"
#include "obj_tbl.h"
#include "tdi_fw.h"

#define IOCTL_TRANSFER_TYPE(ioctl)	((ioctl) & 3)

/* context for tdi_skip_complete */
typedef struct
{
	PIO_COMPLETION_ROUTINE	old_cr;			/* old (original) completion routine */
	PVOID					old_context;	/* old (original) parameter for old_cr */
	PIO_COMPLETION_ROUTINE	new_cr;			/* new (replaced) completion routine */
	PVOID					new_context;	/* new (replaced) parameter for new_cr */
	PFILE_OBJECT			fileobj;		/* FileObject from IO_STACK_LOCATION */
	PDEVICE_OBJECT			new_devobj;		/* filter device object */
	UCHAR					old_control;	/* old (original) irps->Control */
} TDI_SKIP_CTX;

/* prototypes */

static NTSTATUS	DeviceDispatch(IN PDEVICE_OBJECT DeviceObject, IN PIRP irp);
static VOID		OnUnload(IN PDRIVER_OBJECT DriverObject);

#ifndef USE_TDI_HOOKING
static NTSTATUS	c_n_a_device(PDRIVER_OBJECT DriverObject, PDEVICE_OBJECT *fltobj,
			     PDEVICE_OBJECT *oldobj, wchar_t *devname);
static void		d_n_d_device(PDRIVER_OBJECT DriverObject, PDEVICE_OBJECT oldobj,
			     PDEVICE_OBJECT fltobj);
#else
static NTSTATUS	hook_tcpip(DRIVER_OBJECT *old_DriverObject, BOOLEAN b_hook);
static NTSTATUS	get_device_object(wchar_t *name, PDEVICE_OBJECT *devobj);
#endif

static NTSTATUS	tdi_skip_complete(IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp, IN PVOID Context);

/* device objects for: */
PDEVICE_OBJECT
g_tcpfltobj = NULL,		// \Device\Tcp
	      g_udpfltobj = NULL,		// \Device\Udp
	      g_ipfltobj = NULL,		// \Device\RawIp
	      g_devcontrol = NULL,	// control device (exclusive access only!)
	      g_devnfo = NULL;		// information device

BOOLEAN g_got_log = FALSE;	// got log app

#ifndef USE_TDI_HOOKING
// original device objects
PDEVICE_OBJECT g_tcpoldobj, g_udpoldobj, g_ipoldobj;
#else
// original driver object
DRIVER_OBJECT g_old_DriverObject;
BOOLEAN g_hooked = FALSE;
#endif

/* for IOCTL_TDI_QUERY_DIRECT_SEND_HANDLER */
typedef NTSTATUS  TCPSendData_t(IN PIRP Irp, IN PIO_STACK_LOCATION IrpSp);
static TCPSendData_t *g_TCPSendData = NULL;
static TCPSendData_t new_TCPSendData;

/* global traffic stats */
unsigned __int64 g_traffic[TRAFFIC_MAX];
KSPIN_LOCK g_traffic_guard;

/* initialization */
NTSTATUS
DriverEntry(IN PDRIVER_OBJECT theDriverObject,
	    IN PUNICODE_STRING theRegistryPath)
{
	NTSTATUS status = STATUS_SUCCESS;
	int i;
	UNICODE_STRING name, linkname;
	
	memtrack_init();
	KeInitializeSpinLock(&g_traffic_guard);
	
#ifdef USE_TDI_HOOKING
	KdPrint(("[tdi_fw] WARNING! Using unstable working mode: TDI hooking!\n"));
#endif
	
	status = ot_init();
	if (status != STATUS_SUCCESS)
		{
		KdPrint(("[tdi_fw] DriverEntry: ot_init: 0x%x\n", status));
		goto done;
		}
	
	status = filter_init();
	if (status != STATUS_SUCCESS)
		{
		KdPrint(("[tdi_fw] DriverEntry: filter_init: 0x%x\n", status));
		goto done;
		}
	
	status = conn_state_init();
	if (status != STATUS_SUCCESS)
		{
		KdPrint(("[tdi_fw] DriverEntry: conn_state_init: 0x%x\n", status));
		goto done;
		}
	
	for (i = 0; i < IRP_MJ_MAXIMUM_FUNCTION; i++)
		theDriverObject->MajorFunction[i] = DeviceDispatch;
	
#if DBG
	// register UnLoad procedure
	theDriverObject->DriverUnload = OnUnload;
#endif
	
	/* create control device and symbolic link */
	
	RtlInitUnicodeString(&name, L"\\Device\\tdifw");
	
	status = IoCreateDevice(theDriverObject,
				0,
				&name,
				0,
				0,
				TRUE,		// exclusive!
				&g_devcontrol);
	if (status != STATUS_SUCCESS)
		{
		KdPrint(("[tdi_fw] DriverEntry: IoCreateDevice(control): 0x%x!\n", status));
		goto done;
		}
	
	RtlInitUnicodeString(&linkname, L"\\??\\tdifw");
	
	status = IoCreateSymbolicLink(&linkname, &name);
	if (status != STATUS_SUCCESS)
		{
		KdPrint(("[tdi_fw] DriverEntry: IoCreateSymbolicLink: 0x%x!\n", status));
		goto done;
		}
	
	RtlInitUnicodeString(&name, L"\\Device\\tdifw_nfo");
	
	status = IoCreateDevice(theDriverObject,
				0,
				&name,
				0,
				0,
				FALSE,		// not exclusive!
				&g_devnfo);
	if (status != STATUS_SUCCESS)
		{
		KdPrint(("[tdi_fw] DriverEntry: IoCreateDevice(nfo): 0x%x!\n", status));
		goto done;
		}
	
	RtlInitUnicodeString(&linkname, L"\\??\\tdifw_nfo");
	
	status = IoCreateSymbolicLink(&linkname, &name);
	if (status != STATUS_SUCCESS)
		{
		KdPrint(("[tdi_fw] DriverEntry: IoCreateSymbolicLink: 0x%x!\n", status));
		goto done;
		}
	
#ifndef USE_TDI_HOOKING
	
	status = c_n_a_device(theDriverObject, &g_tcpfltobj, &g_tcpoldobj, L"\\Device\\Tcp");
	if (status != STATUS_SUCCESS)
		{
		KdPrint(("[tdi_fw] DriverEntry: c_n_a_device: 0x%x\n", status));
		goto done;
		}
	
	status = c_n_a_device(theDriverObject, &g_udpfltobj, &g_udpoldobj, L"\\Device\\Udp");
	if (status != STATUS_SUCCESS)
		{
		KdPrint(("[tdi_fw] DriverEntry: c_n_a_device: 0x%x\n", status));
		goto done;
		}
	
	status = c_n_a_device(theDriverObject, &g_ipfltobj, &g_ipoldobj, L"\\Device\\RawIp");
	if (status != STATUS_SUCCESS)
		{
		KdPrint(("[tdi_fw] DriverEntry: c_n_a_device: 0x%x\n", status));
		goto done;
		}
	
#else	/* USE_TDI_HOOKING */
	
	/* get device objects for tcp/udp/ip */
	
	status = get_device_object(L"\\Device\\Tcp", &g_tcpfltobj);
	if (status != STATUS_SUCCESS)
		{
		KdPrint(("[tdi_fw] DriverEntry: get_device_object(tcp): 0x%x\n", status));
		goto done;
		}
	
	status = get_device_object(L"\\Device\\Udp", &g_udpfltobj);
	if (status != STATUS_SUCCESS)
		{
		KdPrint(("[tdi_fw] DriverEntry: get_device_object(udp): 0x%x\n", status));
		goto done;
		}
	
	status = get_device_object(L"\\Device\\RawIp", &g_ipfltobj);
	if (status != STATUS_SUCCESS)
		{
		KdPrint(("[tdi_fw] DriverEntry: get_device_object(ip): 0x%x\n", status));
		goto done;
		}
	
	/* hook tcpip */
	
	status = hook_tcpip(&g_old_DriverObject, TRUE);
	if (status != STATUS_SUCCESS)
		{
		KdPrint(("[tdi_fw] DriverEntry: hook_driver: 0x%x\n", status));
		goto done;
		}
	g_hooked = TRUE;
	
#endif	/* USE_TDI_HOOKING */
	
	status = STATUS_SUCCESS;
	
	done:
	if (status != STATUS_SUCCESS)
		{
		// cleanup
		OnUnload(theDriverObject);
		}
	
	return status;
}

/* deinitialization */
VOID	OnUnload(IN PDRIVER_OBJECT DriverObject)
{
#ifndef USE_TDI_HOOKING
	d_n_d_device(DriverObject, g_tcpoldobj, g_tcpfltobj);
	d_n_d_device(DriverObject, g_udpoldobj, g_udpfltobj);
	d_n_d_device(DriverObject, g_ipoldobj, g_ipfltobj);
#else
	if (g_hooked)
		hook_tcpip(&g_old_DriverObject, FALSE);
#endif
	
	// delete control device and symbolic link
	if (g_devcontrol != NULL)
		{
		UNICODE_STRING linkname;
		
		RtlInitUnicodeString(&linkname, L"\\??\\tdifw");
		IoDeleteSymbolicLink(&linkname);
		
		IoDeleteDevice(g_devcontrol);
		}
	
	// delete info device and symbolic link
	if (g_devnfo != NULL)
		{
		UNICODE_STRING linkname;
		
		RtlInitUnicodeString(&linkname, L"\\??\\tdifw_nfo");
		IoDeleteSymbolicLink(&linkname);
		
		IoDeleteDevice(g_devnfo);
		}
	
	filter_free();
	ot_free();
	conn_state_free();		// call after ot_free()
	
	memtrack_free();
}

#ifndef USE_TDI_HOOKING

/* create & attach device */
NTSTATUS	c_n_a_device(PDRIVER_OBJECT DriverObject, PDEVICE_OBJECT *fltobj, PDEVICE_OBJECT *oldobj,
		     wchar_t *devname)
{
	NTSTATUS status;
	UNICODE_STRING str;
	
	/* create filter device */
	
	status = IoCreateDevice(DriverObject,
				0,
				NULL,
				FILE_DEVICE_UNKNOWN,
				0,
				TRUE,
				fltobj);
	if (status != STATUS_SUCCESS)
		{
		KdPrint(("[tdi_fw] c_n_a_device: IoCreateDevice(%S): 0x%x\n", devname, status));
		return status;
		}
	
	(*fltobj)->Flags |= DO_DIRECT_IO;
	
	RtlInitUnicodeString(&str, devname);
	
	status = IoAttachDevice(*fltobj, &str, oldobj);
	if (status != STATUS_SUCCESS)
		{
		KdPrint(("[tdi_fw] DriverEntry: IoAttachDevice(%S): 0x%x\n", devname, status));
		return status;
		}
	
	KdPrint(("[tdi_fw] DriverEntry: %S fileobj: 0x%x\n", devname, *fltobj));
	
	return STATUS_SUCCESS;
}

/* detach & delete device */
void	d_n_d_device(PDRIVER_OBJECT DriverObject, PDEVICE_OBJECT oldobj, PDEVICE_OBJECT fltobj)
{
	/*
	* Detaching of a filter driver at runtime is a high-risk deal
	*/
	
#if 1
	// for extremal guys only!
	if (oldobj != NULL && fltobj != NULL)
		{
		int i;
		for (i = 0; i < IRP_MJ_MAXIMUM_FUNCTION; i++)
			DriverObject->MajorFunction[i] = g_tcpoldobj->DriverObject->MajorFunction[i];
		}
#endif
	
	if (oldobj != NULL)
		IoDetachDevice(oldobj);
	
	if (fltobj != NULL)
		IoDeleteDevice(fltobj);
}

#else	/* USE_TDI_HOOKING */

/* hook/unhook driver */
NTSTATUS	hook_tcpip(DRIVER_OBJECT *old_DriverObject, BOOLEAN b_hook)
{
	UNICODE_STRING drv_name;
	NTSTATUS status;
	PDRIVER_OBJECT new_DriverObject;
	int i;
	
	RtlInitUnicodeString(&drv_name, L"\\Driver\\Tcpip");
	
	status = ObReferenceObjectByName(&drv_name, OBJ_CASE_INSENSITIVE, NULL, 0,
					 IoDriverObjectType, KernelMode, NULL, &new_DriverObject);
	if (status != STATUS_SUCCESS)
		{
		KdPrint(("[tdi_fw] hook_driver: ObReferenceObjectByName\n"));
		return status;
		}
	
	for (i = 0; i < IRP_MJ_MAXIMUM_FUNCTION; i++)
		{
		if (b_hook)
			{
			old_DriverObject->MajorFunction[i] = new_DriverObject->MajorFunction[i];
			new_DriverObject->MajorFunction[i] = DeviceDispatch;
			}
		else
			new_DriverObject->MajorFunction[i] = old_DriverObject->MajorFunction[i];
		}
	
	return STATUS_SUCCESS;
}

/* get device object by its name */
NTSTATUS	get_device_object(wchar_t *name, PDEVICE_OBJECT *devobj)
{
	UNICODE_STRING str;
	NTSTATUS status;
	PFILE_OBJECT fileobj;
	
	RtlInitUnicodeString(&str, name);
	
	status = IoGetDeviceObjectPointer(&str, FILE_ALL_ACCESS, &fileobj, devobj);
	if (status == STATUS_SUCCESS)
		ObDereferenceObject(fileobj);
	
	return status;
}

#endif	/* USE_TDI_HOOKING */

/* dispatch */
NTSTATUS	DeviceDispatch(IN PDEVICE_OBJECT DeviceObject, IN PIRP irp)
{
	PIO_STACK_LOCATION irps;
	NTSTATUS status;
	
	// sanity check
	if (irp == NULL)
		{
		KdPrint(("[tdi_fw] DeviceDispatch: !irp\n"));
		return STATUS_SUCCESS;
		}
	
	irps = IoGetCurrentIrpStackLocation(irp);
	
	if (DeviceObject == g_tcpfltobj || DeviceObject == g_udpfltobj ||
		DeviceObject == g_ipfltobj)
		{
		
		/*
		* This IRP is for filtered device
		*/
		
		int result;
		struct completion completion;
		
		memset(&completion, 0, sizeof(completion));
		
		// Analyze MajorFunction
		switch (irps->MajorFunction)
			{
			
			case IRP_MJ_CREATE:		/* create fileobject */
				
				result = tdi_create(irp, irps, &completion);
				
				status = tdi_dispatch_complete(DeviceObject, irp, result,
							       completion.routine, completion.context);
				
				break;
				
			case IRP_MJ_DEVICE_CONTROL:
				
				KdPrint(("[tdi_fw] DeviceDispatch: IRP_MJ_DEVICE_CONTROL, control 0x%x for 0x%08X\n",
					irps->Parameters.DeviceIoControl.IoControlCode, irps->FileObject));
				
				if (KeGetCurrentIrql() == PASSIVE_LEVEL)
					{
					/*
					* try to convert it to IRP_MJ_INTERNAL_DEVICE_CONTROL
					* (works on PASSIVE_LEVEL only!)
					*/
					status = TdiMapUserRequest(DeviceObject, irp, irps);
					
					}
				else
					status = STATUS_NOT_IMPLEMENTED; // set fake status
				
				if (status != STATUS_SUCCESS)
					{
					void *buf = (irps->Parameters.DeviceIoControl.IoControlCode == IOCTL_TDI_QUERY_DIRECT_SEND_HANDLER) ?
						    irps->Parameters.DeviceIoControl.Type3InputBuffer : NULL;
					
					// send IRP to original driver
					status = tdi_dispatch_complete(DeviceObject, irp, FILTER_ALLOW, NULL, NULL);
					
					if (buf != NULL && status == STATUS_SUCCESS)
						{
						
						g_TCPSendData = *(TCPSendData_t **)buf;
						
						KdPrint(("[tdi_fw] DeviceDispatch: IOCTL_TDI_QUERY_DIRECT_SEND_HANDLER: TCPSendData = 0x%x\n",
							g_TCPSendData));
						
						*(TCPSendData_t **)buf = new_TCPSendData;
						}
					
					break;
					}
				
				// don't break! go to internal device control!
				
			case IRP_MJ_INTERNAL_DEVICE_CONTROL:
					{
					/*
					* Analyze ioctl for TDI driver
					*/
					int i;
					
					for (i = 0; g_tdi_ioctls[i].MinorFunction != 0; i++)
						if (g_tdi_ioctls[i].MinorFunction == irps->MinorFunction)
							{
							
#if DBG
						// print description
							KdPrint(("[tdi_fw] DeviceDispatch: %s (0x%x) for 0x%x\n",
								g_tdi_ioctls[i].desc,
								irps->MinorFunction,
								irps->FileObject));
#endif
							
							if (g_tdi_ioctls[i].fn == NULL)
							{
							// send IRP to original driver
							status = tdi_dispatch_complete(DeviceObject, irp, FILTER_ALLOW,
										       NULL, NULL);
							break;
							}
						
						// call dispatch function
						
						result = g_tdi_ioctls[i].fn(irp, irps, &completion);
						
						// complete request
						status = tdi_dispatch_complete(DeviceObject, irp, result,
									       completion.routine, completion.context);