www.pudn.com > tdi_fw.zip > obj_tbl.c

// -*- 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  
#include  
#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; 
}