obj_tbl.c

非常使用的基于TDI驱动开发的应用程序过滤 防火墙的例子

// -*- mode: C++; tab-width: 4; indent-tabs-mode: nil -*- (for GNU Emacs)
//
// $Id: obj_tbl.c,v 1.3 2002/12/03 12:14:27 dev Exp $

/*
 * Working with connection objects, address objects and links between them
 */

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

#include "memtrack.h"
#include "obj_tbl.h"
#include "tdi_fw.h"

#define HASH_SIZE	0x1000
#define CALC_HASH(fileobj)  (((ULONG)(fileobj) >> 5) % HASH_SIZE)

static struct ot_entry **g_ot_hash;
KSPIN_LOCK g_ot_hash_guard;

struct ctx_entry {
	struct ctx_entry *next;
	PFILE_OBJECT addrobj;
	CONNECTION_CONTEXT conn_ctx;
	PFILE_OBJECT connobj;
};

static struct ctx_entry **g_cte_hash;
KSPIN_LOCK g_cte_hash_guard;

static NTSTATUS		do_set_old_event_handler(struct ot_entry *ote, int event_type);

//----------------------------------------------------------------------------

NTSTATUS
ot_init(void)
{
	g_ot_hash = (struct ot_entry **)malloc_np(sizeof(*g_ot_hash) * HASH_SIZE);
	if (g_ot_hash == NULL) {
		KdPrint(("[tdi_fw] ot_init: malloc_np\n"));
		return STATUS_INSUFFICIENT_RESOURCES;
	}
	memset(g_ot_hash, 0, sizeof(*g_ot_hash) * HASH_SIZE);

	KeInitializeSpinLock(&g_ot_hash_guard);

	g_cte_hash = (struct ctx_entry **)malloc_np(sizeof(*g_cte_hash) * HASH_SIZE);
	if (g_cte_hash == NULL) {
		KdPrint(("[tdi_fw] ot_init: malloc_np\n"));
		free(g_ot_hash);
		return STATUS_INSUFFICIENT_RESOURCES;
	}
	memset(g_cte_hash, 0, sizeof(*g_cte_hash) * HASH_SIZE);

	KeInitializeSpinLock(&g_cte_hash_guard);

	return STATUS_SUCCESS;
}

void
ot_free(void)
{
	KIRQL irql;
	int i;

	if (g_ot_hash != NULL) {

		KeAcquireSpinLock(&g_ot_hash_guard, &irql);
		
		for (i = 0; i < HASH_SIZE; i++) {
			struct ot_entry *ote = g_ot_hash[i];
			while (ote) {
				struct ot_entry *ote2 = ote->next;
				int j;

				KdPrint(("[tdi_fw] ot_free: Warning! fileobj 0x%x type %d exists!\n",
					ote->fileobj, ote->type));
				
				if (ote->signature != 'OTE ') {
					KdPrint(("[tdi_fw] ot_free: Warning! fileobj 0x%x invalid signature 0x%x!\n",
						ote->fileobj, ote->signature));
				}

				for (j = 0; j < MAX_EVENT; j++)
					if (ote->ctx[j].old_handler != NULL)
						do_set_old_event_handler(ote, j);

				free(ote);
				ote = ote2;
			}
		}
		free(g_ot_hash);
		g_ot_hash = NULL;

		KeReleaseSpinLock(&g_ot_hash_guard, irql);
	}

	if (g_cte_hash != NULL) {
		KeAcquireSpinLock(&g_cte_hash_guard, &irql);
		
		for (i = 0; i < HASH_SIZE; i++) {
			struct ctx_entry *cte = g_cte_hash[i];
			while (cte) {
				struct ctx_entry *cte2 = cte->next;
				free(cte);
				cte = cte2;
			}
		}
		free(g_cte_hash);
		g_cte_hash = NULL;

		KeReleaseSpinLock(&g_cte_hash_guard, irql);
	}
}

//----------------------------------------------------------------------------

NTSTATUS
ot_add_fileobj(PDEVICE_OBJECT devobj, PFILE_OBJECT fileobj, int fileobj_type,
			   CONNECTION_CONTEXT conn_ctx, int no_guard)
{
	ULONG hash = CALC_HASH(fileobj);
	KIRQL irql;
	struct ot_entry *ote;
	NTSTATUS status;
	int i;

	if (fileobj == NULL)
		return STATUS_INVALID_PARAMETER_2;

	if (!no_guard)
		KeAcquireSpinLock(&g_ot_hash_guard, &irql);
	
	for (ote = g_ot_hash[hash]; ote != NULL; ote = ote->next)
		if (ote->fileobj == fileobj)
			break;

	if (ote == NULL) {
		ote = (struct ot_entry *)malloc_np(sizeof(*ote));
		if (ote == NULL) {
			KdPrint(("[tdi_fw] ot_add_fileobj: malloc_np\n"));
			status = STATUS_INSUFFICIENT_RESOURCES;
			goto done;
		}
		memset(ote, 0, sizeof(*ote));

		ote->next = g_ot_hash[hash];
		g_ot_hash[hash] = ote;

		ote->fileobj = fileobj;
		for (i = 0; i < MAX_EVENT; i++)
			ote->ctx[i].fileobj = fileobj;

	} else {
		KdPrint(("[tdi_fw] ot_add_fileobj: reuse fileobj 0x%x\n", fileobj));

		// set neccessary fields to zero
		
		ote->associated_fileobj = NULL;
		memset(ote->ctx, 0, sizeof(ote->ctx));
		memset(ote->local_addr, 0, sizeof(ote->local_addr));
		memset(ote->remote_addr, 0, sizeof(ote->remote_addr));
	}

	ote->signature = 'OTE ';
	ote->pid = (ULONG)PsGetCurrentProcessId();

	ote->devobj = devobj;

	ote->type = fileobj_type;
	ote->conn_ctx = conn_ctx;
	
	status = STATUS_SUCCESS;

done:
	if (!no_guard)
		KeReleaseSpinLock(&g_ot_hash_guard, irql);

	return status;
}

NTSTATUS
ot_del_fileobj(PFILE_OBJECT fileobj, int *fileobj_type)
{
	ULONG hash = CALC_HASH(fileobj);
	KIRQL irql;
	struct ot_entry *ote, *prev_ote;
	NTSTATUS status;

	if (fileobj == NULL)
		return STATUS_INVALID_PARAMETER_1;

	KeAcquireSpinLock(&g_ot_hash_guard, &irql);

	prev_ote = NULL;
	for (ote = g_ot_hash[hash]; ote; ote = ote->next) {
		if (ote->fileobj == fileobj)
			break;
		prev_ote = ote;
	}

	if (ote == NULL) {
		KdPrint(("[tdi_fw] ot_del_fileobj: fileobj 0x%x not found!\n", fileobj));
		status = STATUS_OBJECT_NAME_NOT_FOUND;
		goto done;
	}

	if (fileobj_type != NULL)
		*fileobj_type = ote->type;

	if (prev_ote != NULL)
		prev_ote->next = ote->next;
	else
		g_ot_hash[hash] = ote->next;

	free(ote);
	status = STATUS_SUCCESS;

done:
	KeReleaseSpinLock(&g_ot_hash_guard, irql);

	return status;
}

struct ot_entry *
ot_find_fileobj(PFILE_OBJECT fileobj, KIRQL *irql)
{
	ULONG hash = CALC_HASH(fileobj);
	struct ot_entry *ote;

	if (fileobj == NULL)
		return NULL;

	if (irql != NULL)
		KeAcquireSpinLock(&g_ot_hash_guard, irql);

	for (ote = g_ot_hash[hash]; ote != NULL; ote = ote->next)
		if (ote->fileobj == fileobj)
			break;

	if (ote == NULL) {
		KdPrint(("[tdi_fw] ot_find_fileobj: fileobj 0x%x not found!\n", fileobj));
		if (irql != NULL)
			KeReleaseSpinLock(&g_ot_hash_guard, *irql);
	}

	return ote;
}

//----------------------------------------------------------------------------

#define CALC_HASH_2(addrobj, conn_ctx)	CALC_HASH((ULONG)(addrobj) ^ (ULONG)(conn_ctx))

NTSTATUS
ot_add_conn_ctx(PFILE_OBJECT addrobj, CONNECTION_CONTEXT conn_ctx, PFILE_OBJECT connobj)
{
	ULONG hash = CALC_HASH_2(addrobj, conn_ctx);
	KIRQL irql;
	struct ctx_entry *cte;
	NTSTATUS status;

	KeAcquireSpinLock(&g_cte_hash_guard, &irql);
	
	for (cte = g_cte_hash[hash]; cte != NULL; cte = cte->next)
		if (cte->addrobj == addrobj && cte->conn_ctx == conn_ctx)
			break;

	if (cte == NULL) {
		KdPrint(("[tdi_fw] ot_add_fileobj: reuse addrobj 0x%x, conn_ctx 0x%x\n",
			addrobj, conn_ctx));

		cte = (struct ctx_entry *)malloc_np(sizeof(*cte));
		if (cte == NULL) {
			KdPrint(("[tdi_fw] ot_add_conn_ctx: malloc_np\n"));
			status = STATUS_INSUFFICIENT_RESOURCES;
			goto done;
		}
		cte->next = g_cte_hash[hash];
		g_cte_hash[hash] = cte;
	
		cte->addrobj = addrobj;
		cte->conn_ctx = conn_ctx;
	}

	cte->connobj = connobj;
	
	status = STATUS_SUCCESS;
done:

	KeReleaseSpinLock(&g_cte_hash_guard, irql);
	return status;
}

NTSTATUS
ot_del_conn_ctx(PFILE_OBJECT addrobj, CONNECTION_CONTEXT conn_ctx)
{
	ULONG hash = CALC_HASH_2(addrobj, conn_ctx);
	KIRQL irql;
	struct ctx_entry *cte, *prev_cte;
	NTSTATUS status;

	KeAcquireSpinLock(&g_cte_hash_guard, &irql);

	prev_cte = NULL;
	for (cte = g_cte_hash[hash]; cte != NULL; cte = cte->next) {
		if (cte->addrobj == addrobj && cte->conn_ctx == conn_ctx)
			break;
		prev_cte = cte;
	}

	if (cte == NULL) {
		KdPrint(("[tdi_fw] ot_del_conn_ctx: addrobj 0x%x not found!\n", addrobj));
		status = STATUS_OBJECT_NAME_NOT_FOUND;
		goto done;
	}

	if (prev_cte != NULL)
		prev_cte->next = cte->next;
	else
		g_cte_hash[hash] = cte->next;

	free(cte);

	status = STATUS_SUCCESS;
done:

	KeReleaseSpinLock(&g_cte_hash_guard, irql);
	return status;
}

PFILE_OBJECT
ot_find_conn_ctx(PFILE_OBJECT addrobj, CONNECTION_CONTEXT conn_ctx)
{
	ULONG hash = CALC_HASH_2(addrobj, conn_ctx);
	KIRQL irql;
	struct ctx_entry *cte;
	PFILE_OBJECT result = NULL;

	KeAcquireSpinLock(&g_cte_hash_guard, &irql);
	
	for (cte = g_cte_hash[hash]; cte != NULL; cte = cte->next)
		if (cte->addrobj == addrobj && cte->conn_ctx == conn_ctx) {
			result = cte->connobj;
			break;
		}

	KeReleaseSpinLock(&g_cte_hash_guard, irql);
	return result;
}

//----------------------------------------------------------------------------

NTSTATUS
do_set_old_event_handler(struct ot_entry *ote, int event_type)
{
	NTSTATUS status;
	PIRP query_irp = NULL;
	PDEVICE_OBJECT devobj;
	
#ifndef USE_TDI_HOOKING
	// get original (unhooked) device object
	if (ote->devobj == g_tcpfltobj)
		devobj = g_tcpoldobj;
	else if (ote->devobj == g_udpfltobj)
		devobj = g_udpoldobj;
	else if (ote->devobj == g_ipfltobj)
		devobj = g_ipoldobj;
	else
		return STATUS_UNSUCCESSFUL;
#else
	// original and hooked device objects are the same
	devobj = ote->devobj;
#endif

	KdPrint(("[tdi_fw] do_set_old_event_handler: devobj 0x%x, fileobj 0x%x, handler 0x%x, context 0x%x\n",
		devobj, ote->fileobj, ote->ctx[event_type].old_handler, ote->ctx[event_type].old_context));

	// FIXME!!! Calling TdiBuildInternalDeviceControlIrp at DISPATCH_LEVEL!!!
	query_irp = TdiBuildInternalDeviceControlIrp(TDI_SET_EVENT_HANDLER, devobj,
		ote->fileobj, NULL, NULL);
	if (query_irp == NULL) {
		KdPrint(("[tdi_fw] do_set_old_event_handler: TdiBuildInternalDeviceControlIrp\n"));
		status = STATUS_UNSUCCESSFUL;
		goto done;
	}

	TdiBuildSetEventHandler(query_irp, devobj, ote->fileobj, NULL, NULL, event_type,
		ote->ctx[event_type].old_handler, ote->ctx[event_type].old_context);

	status = IoCallDriver(devobj, query_irp);
	query_irp = NULL;

	if (status != STATUS_SUCCESS) {
		KdPrint(("[tdi_fw] do_set_old_event_handler: IoCallDriver: 0x%x\n", status));
		goto done;
	}

	// don't wait to complete

done:
	if (query_irp)
		IoFreeIrp(query_irp);

	return status;
}