conn_state.c

开源的防火墙代码

	KeAcquireSpinLock(&g_ot_hash_guard, &irql);	// lock obj_tbl first (avoid deadlocks!)
	KeAcquireSpinLockAtDpcLevel(&g_conn_guard);	// lock our hash
	
	ote_conn = ot_find_fileobj(connobj, NULL);		// obj_tbl already locked
	if (ote_conn == NULL)
		goto done;
	
	ce = ote_conn->conn_entry;
	ote_conn->conn_entry = NULL;
	
	if (is_disconnect && ote_conn->log_disconnect)
		log_disconnect(ote_conn);
	
	if (ce != NULL)
		del_tcp_conn_obj(ce, TRUE);
	
	done:
	KeReleaseSpinLockFromDpcLevel(&g_conn_guard);	// unlock our hash
	KeReleaseSpinLock(&g_ot_hash_guard, irql);	// unlock obj_tbl
}

void	del_tcp_conn_obj(struct conn_entry *ce, BOOLEAN no_guard)
{
	KIRQL irql;
	
	KdPrint(("[tdi_fw] del_tcp_conn_obj: CLOSED %x:%u <-> %x:%u (state=%d)\n",
		ce->laddr, ntohs(ce->lport), ce->raddr, ntohs(ce->rport), ce->state));
	
	if (!no_guard)
		KeAcquireSpinLock(&g_conn_guard, &irql);	// lock our hash
	
	// set state to TCP_STATE_CLOSED and add to special list to queue it for deleting
	ce->state = TCP_STATE_CLOSED;
	ce->next_to_del = g_conn_to_del;
	g_conn_to_del = ce;
	
	KeQueryTickCount(&ce->ticks);
	
	KeSetEvent(&g_conn_new_to_del, IO_NO_INCREMENT, FALSE);
	
	ce->connobj = NULL;		// no connection object related for now!!!
	
	KdPrint(("[tdi_fw] del_tcp_conn_obj: state table entry scheduled for deletion!\n"));
	
	if (!no_guard)
		KeReleaseSpinLock(&g_conn_guard, irql);	// unlock our hash
}

NTSTATUS	set_tcp_conn_state(PFILE_OBJECT connobj, int state)
{
	KIRQL irql;
	struct ot_entry *ote_conn;
	NTSTATUS status;
	
	KeAcquireSpinLock(&g_ot_hash_guard, &irql);	// lock obj_tbl first (avoid deadlocks!)
	KeAcquireSpinLockAtDpcLevel(&g_conn_guard);	// lock our hash
	
	ote_conn = ot_find_fileobj(connobj, NULL);		// obj_tbl already locked
	if (ote_conn == NULL)
		{
		KdPrint(("[tdi_fw] set_tcp_conn_state: ot_find_fileobj(0x%x)!\n", connobj));
		status = STATUS_OBJECT_NAME_NOT_FOUND;
		goto done;
		}
	
	if (ote_conn->conn_entry != NULL)
		{
		ote_conn->conn_entry->state = state;
		
		KdPrint(("[tdi_fw] set_tcp_conn_state: got CHANGE CONNECT STATE %x:%u <-> %x:%u (state=%d)\n",
			ote_conn->conn_entry->laddr, ntohs(ote_conn->conn_entry->lport),
			ote_conn->conn_entry->raddr, ntohs(ote_conn->conn_entry->rport),
			ote_conn->conn_entry->state));
		
		status = STATUS_SUCCESS;
		}
	else	{
		KdPrint(("[tdi_fw] set_tcp_conn_state: no conn_entry!\n"));
		status = STATUS_INVALID_PARAMETER;
		}
	
	done:
	KeReleaseSpinLockFromDpcLevel(&g_conn_guard);		// unlock our hash
	KeReleaseSpinLock(&g_ot_hash_guard, irql);		// unlock obj_tbl
	
	return status;
}

NTSTATUS	set_tcp_conn_local(PFILE_OBJECT connobj, TA_ADDRESS *local)
{
	KIRQL irql;
	struct ot_entry *ote_conn;
	NTSTATUS status;
	
	KeAcquireSpinLock(&g_ot_hash_guard, &irql);	// lock obj_tbl first (avoid deadlocks!)
	KeAcquireSpinLockAtDpcLevel(&g_conn_guard);	// lock our hash
	
	ote_conn = ot_find_fileobj(connobj, NULL);		// obj_tbl already locked
	if (ote_conn == NULL)
		{
		KdPrint(("[tdi_fw] set_tcp_conn_local: ot_find_fileobj(0x%x)!\n", connobj));
		status = STATUS_OBJECT_NAME_NOT_FOUND;
		goto done;
		}
	
	if (ote_conn->conn_entry != NULL)
		{
		struct conn_entry *ce = ote_conn->conn_entry;
		ULONG hash;
		
		// remove our conn_entry from list and recalculate hash
		
		if (ce->prev != NULL)
			ce->prev->next = ce->next;
		else	{
			hash = CALC_CONN_HASH(ce->laddr, ce->lport, ce->raddr, ce->rport);
			g_conn[hash] = ce->next;
			}
		
		if (ce->next != NULL)
			ce->next->prev = ce->prev;
		
		ce->laddr = ((TDI_ADDRESS_IP *)(local->Address))->in_addr;
		ce->lport = ((TDI_ADDRESS_IP *)(local->Address))->sin_port;
		
		hash = CALC_CONN_HASH(ce->laddr, ce->lport, ce->raddr, ce->rport);
		
		KdPrint(("[tdi_fw] set_tcp_conn_local: got CHANGE CONNECT LOCAL %x:%u <-> %x:%u (state=%d)\n",
			ce->laddr, ntohs(ce->lport), ce->raddr, ntohs(ce->rport), ce->state));
		
		// and now add our conn_entry under new hash value
		
		ce->next = g_conn[hash];
		if (g_conn[hash] != NULL)
			g_conn[hash]->prev = ce;
		g_conn[hash] = ce;
		
		ce->prev = NULL;
		
		status = STATUS_SUCCESS;
		
		}
	else	{
		KdPrint(("[tdi_fw] set_tcp_conn_local: no conn_entry!\n"));
		status = STATUS_INVALID_PARAMETER;
		}
	
	done:
	KeReleaseSpinLockFromDpcLevel(&g_conn_guard);		// unlock our hash
	KeReleaseSpinLock(&g_ot_hash_guard, irql);		// unlock obj_tbl
	
	return status;
}

int	get_tcp_conn_state(ULONG laddr, USHORT lport, ULONG raddr, USHORT rport)
{
	ULONG hash = CALC_CONN_HASH(laddr, lport, raddr, rport);
	KIRQL irql;
	struct conn_entry *ce;
	int result;
	
	KeAcquireSpinLock(&g_conn_guard, &irql);
	
	result = TCP_STATE_NONE;
	
	for (ce = g_conn[hash]; ce != NULL; ce = ce->next)
		if (ce->laddr == laddr && ce->lport == lport &&
		ce->raddr == raddr && ce->rport == rport)
			{
			
			result = ce->state;
			break;
			}
	
	KeReleaseSpinLock(&g_conn_guard, irql);
	
	return result;
}

int	get_tcp_conn_state_by_obj(PFILE_OBJECT connobj)
{
	KIRQL irql;
	struct ot_entry *ote_conn;
	NTSTATUS status;
	int result = TCP_STATE_NONE;
	
	KeAcquireSpinLock(&g_ot_hash_guard, &irql);	// lock obj_tbl first (avoid deadlocks!)
	KeAcquireSpinLockAtDpcLevel(&g_conn_guard);	// lock our hash
	
	ote_conn = ot_find_fileobj(connobj, NULL);		// obj_tbl already locked
	if (ote_conn == NULL)
		{
		KdPrint(("[tdi_fw] set_tcp_conn_state: ot_find_fileobj(0x%x)!\n", connobj));
		goto done;
		}
	
	if (ote_conn->conn_entry != NULL)
		result = ote_conn->conn_entry->state;
	
	done:
	KeReleaseSpinLockFromDpcLevel(&g_conn_guard);	// unlock our hash
	KeReleaseSpinLock(&g_ot_hash_guard, irql);	// unlock obj_tbl
	
	return result;
}

NTSTATUS	enum_tcp_conn(struct tcp_conn_nfo *buf, ULONG *buf_len, ULONG buf_size)
{
	NTSTATUS status = STATUS_SUCCESS;
	KIRQL irql;
	ULONG hash;
	struct conn_entry *ce;
	
	*buf_len = 0;
	
	if (buf_size < sizeof(struct tcp_conn_nfo))
		return STATUS_INVALID_PARAMETER;
	
	KeAcquireSpinLock(&g_ot_hash_guard, &irql);	// lock obj_tbl first (avoid deadlocks!)
	KeAcquireSpinLockAtDpcLevel(&g_conn_guard);	// lock our hash
	
	for (hash = 0; hash < CONN_HASH_SIZE; hash++)
		{
		for (ce = g_conn[hash]; ce != NULL; ce = ce->next)
			{
			struct ot_entry *ote;
			
			if (ce->state == TCP_STATE_CLOSED)
				continue;		// don't log "closed"
			
			buf->laddr = ce->laddr;
			buf->lport = ce->lport;
			buf->raddr = ce->raddr;
			buf->rport = ce->rport;
			buf->state = ce->state;
			
			// try to get pid
			ote = ot_find_fileobj(ce->connobj, NULL);		// obj_tbl already locked
			if (ote != NULL)
				{
				buf->pid = ote->pid;
				buf->bytes_in = ote->bytes_in;
				buf->bytes_out = ote->bytes_out;
				}
			else	{
				buf->pid = 0;
				buf->bytes_in = 0;
				buf->bytes_out = 0;
				}
			
			*buf_len += sizeof(struct tcp_conn_nfo);
			buf++;
			
			if (*buf_len + sizeof(struct tcp_conn_nfo) > buf_size)
				{
				status = STATUS_BUFFER_TOO_SMALL;
				break;
				}
			}
		
		if (status != STATUS_SUCCESS)
			break;
		}
	
	KeReleaseSpinLockFromDpcLevel(&g_conn_guard);	// unlock our_hash
	KeReleaseSpinLock(&g_ot_hash_guard, irql);	// unlock obj_tbl
	
	return status;
}

void	conn_thread(PVOID param)
{
	LARGE_INTEGER timeout;
	
	timeout.QuadPart = 0;
	
	do	{
		LARGE_INTEGER ticks;
		struct conn_entry *ce, *prev_ce, *ce2;
		KIRQL irql;
		
		KeWaitForSingleObject(&g_conn_new_to_del, Executive, KernelMode, FALSE,
				      (timeout.QuadPart != 0) ? &timeout : NULL);
		
		KdPrint(("[tdi_fw] conn_thread: wake up!\n"));
		
		KeQueryTickCount(&ticks);
		
		// delete all entries older than MAX_CLOSED_TIME sec
		// or all entries if (g_conn_thread == NULL)
		
		// and also find min timeout to next entry
		
		KeAcquireSpinLock(&g_conn_guard, &irql);
		
		prev_ce = NULL;
		timeout.QuadPart = 0;
		for (ce = g_conn_to_del; ce != NULL;)
			{
			__int64 delta = (ticks.QuadPart - ce->ticks.QuadPart) * KeQueryTimeIncrement();
			
			KdPrint(("[tdi_fw] conn_thread: delta %d msec\n", delta / 10000));
			
			if (delta >= MAX_CLOSED_TIME * 1000 * 10000 ||
				g_conn_thread == NULL)
				{
				
				KdPrint(("[tdi_fw] conn_thread: remove it!\n"));
				
				// remove this entry!!!
				
				if (prev_ce == NULL)
					g_conn_to_del = ce->next_to_del;
				else
					prev_ce->next_to_del = ce->next_to_del;
				
				// delete ce from our hash
				
				if (ce->prev != NULL)
					ce->prev->next = ce->next;
				else	{
					ULONG hash = CALC_CONN_HASH(ce->laddr, ce->lport, ce->raddr, ce->rport);
					g_conn[hash] = ce->next;
					}
				
				if (ce->next != NULL)
					ce->next->prev = ce->prev;
				
				ce2 = ce->next_to_del;
				free(ce);
				
				ce = ce2;
				
				}
			else	{
				delta = (MAX_CLOSED_TIME * 1000 * 10000 - delta);
				
				if (delta > -timeout.QuadPart || timeout.QuadPart == 0)
					{
					timeout.QuadPart = -delta;	// how much to wait?
					
					KdPrint(("[tdi_fw] conn_thread: we must wait %d ms!\n", timeout.QuadPart / 10000));
					}
				
				prev_ce = ce;
				ce = ce->next_to_del;
				}
			}
		
		KeReleaseSpinLock(&g_conn_guard, irql);
		
		} while(g_conn_thread != NULL);
}

void	log_disconnect(struct ot_entry *ote_conn)
{
	TA_ADDRESS *local_addr, *remote_addr;
	struct flt_request request;
	
	local_addr = (TA_ADDRESS *)(ote_conn->local_addr);
	remote_addr = (TA_ADDRESS *)(ote_conn->remote_addr);
	
	KdPrint(("[tdi_flt] del_tcp_conn: %x:%u -> %x:%u\n",
		ntohl(((TDI_ADDRESS_IP *)(local_addr->Address))->in_addr),
	ntohs(((TDI_ADDRESS_IP *)(local_addr->Address))->sin_port),
	ntohl(((TDI_ADDRESS_IP *)(remote_addr->Address))->in_addr),
	ntohs(((TDI_ADDRESS_IP *)(remote_addr->Address))->sin_port)));
	
	memset(&request, 0, sizeof(request));
	
	request.struct_size = sizeof(request);
	
	request.result = FILTER_DISCONNECT;
	request.proto = ote_conn->ipproto;
	request.direction = DIRECTION_ANY;
	
	request.pid = (ULONG)-1;		// don't use pid because on close there's no info in database
	
	// get user SID & attributes!
	if ((request.sid_a = copy_sid_a(ote_conn->sid_a, ote_conn->sid_a_size)) != NULL)
		request.sid_a_size = ote_conn->sid_a_size;
	
	memcpy(&request.addr.from, &local_addr->AddressType, sizeof(struct sockaddr));
	memcpy(&request.addr.to, &remote_addr->AddressType, sizeof(struct sockaddr));
	request.addr.len = sizeof(struct sockaddr_in);
	
	request.log_bytes_in = ote_conn->bytes_in;
	request.log_bytes_out = ote_conn->bytes_out;
	
	log_request(&request);
}