TDIFW를 Vista 이후에서 동작시키기

Kernel Inside/TDI 2012. 4. 23. 15:08

TDI는 NDIS와 달리 제대로 돌아가는 샘플코드가 많지 않습니다. 그래서... 공부하기도 어렵죠.

그나마 제일 쓸만한 샘플코드 역할을 하는 것이 바로 오픈소스 방화벽인 TDIFW인데요... 이걸 Vista 이후에서 돌리면... 제대로 돌아가지 않습니다. 이 문제 관련해서 몇군데 수정을 해서 (일단은...) 제대로 동작하게 만들었는데요... 제가 찾아낸 문제있는 부분과 수정한 결과를 공유합니다.

제가 테스트한 바로는 아래와 같이 두군데만 수정하면 잘 동작하는 것 같습니다. (상업적인 환경에서 돌려본 건 아닙니다. 물론 책임도 안지구요. ^^)

=================================================================================================================

문제 1 : Address Object를 쿼리하는 부분.

문제의 코드는 다음과 같습니다.

/* this completion routine queries address and port from address object */
NTSTATUS
tdi_create_addrobj_complete
        (IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp, IN PVOID Context)
{
	NTSTATUS status;
	PIO_STACK_LOCATION irps = IoGetCurrentIrpStackLocation(Irp);
	PIRP query_irp = (PIRP)Context;
	PDEVICE_OBJECT devobj;
	TDI_CREATE_ADDROBJ2_CTX *ctx = NULL;
	PMDL mdl = NULL;

// 중간 생략
	devobj = get_original_devobj(DeviceObject, NULL);	// use original devobj!
	if (devobj == NULL) {
		KdPrint(("[tdi_fw] tdi_create_addrobj_complete: get_original_devobj!\n"));

status = STATUS_INVALID_PARAMETER;
		goto done;
	}

TdiBuildQueryInformation(query_irp, devobj, irps->FileObject,
		tdi_create_addrobj_complete2, ctx,
		TDI_QUERY_ADDRESS_INFO, mdl);

status = IoCallDriver(devobj, query_irp);  // 문제의 부분
	query_irp = NULL;
	mdl = NULL;
	ctx = NULL;

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

status = STATUS_SUCCESS;

done:
	// cleanup
	if (mdl != NULL)
		IoFreeMdl(mdl);
	
	if (ctx != NULL) {
		if (ctx->tai != NULL)
			free(ctx->tai);
		free(ctx);
	}
	
	if (query_irp != NULL)
		IoCompleteRequest(query_irp, IO_NO_INCREMENT);

Irp->IoStatus.Status = status;
	
	if (status != STATUS_SUCCESS) {
		// tdi_create failed - remove fileobj from hash
		ot_del_fileobj(irps->FileObject, NULL);
	}

return tdi_generic_complete(DeviceObject, Irp, Context);
}

위의 코드 중 IoCallDriver 하는 부분을 보면 결과가 STATUS_PENDING인 경우에 대한 처리가 되어 있지 않음을 알 수 있습니다. XP에서는 이부분이 동기적으로 Return되었기 때문에 별 문제가 없었지만 Vista 이후에서는 비동기적으로 처리되기 때문에 IoCallDriver가 실패한 것으로 처리되어 done: 부분에서 Hash Table에 들어간 Address Object를 삭제시켜버리는 로직이 실행되어 버립니다.

보통 이럴 땐 완료루틴이 호출될 때까지 기다려주는 작업을 해 주죠. (IRQL 체크가 필요하겠지만... 테스트 결과로는 항상 PASSIVE로 들어오는 것 같습니다.)

제가 수정한 코드는 다음과 같습니다.

/* this completion routine queries address and port from address object */
NTSTATUS
tdi_create_addrobj_complete
       (IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp, IN PVOID Context)
{
	NTSTATUS status;
	PIO_STACK_LOCATION irps = IoGetCurrentIrpStackLocation(Irp);
	PIRP query_irp = (PIRP)Context;
	PDEVICE_OBJECT devobj;
	TDI_CREATE_ADDROBJ2_CTX *ctx = NULL;
	PMDL mdl = NULL;
	KEVENT		CompEvent;

KeInitializeEvent(&CompEvent, NotificationEvent, FALSE);

KdPrint(("[tdi_fw] tdi_create_addrobj_complete: devobj 0x%x; addrobj 0x%x\n",
		DeviceObject, irps->FileObject));

// 중간생략

ctx->pEvent = &CompEvent;

mdl = IoAllocateMdl(ctx->tai, TDI_ADDRESS_INFO_MAX, FALSE, FALSE, NULL);
	if (mdl == NULL) {
		KdPrint(("[tdi_fw] tdi_create_addrobj_complete: IoAllocateMdl!\n"));
		
		status = STATUS_INSUFFICIENT_RESOURCES;
		goto done;
	}
	MmBuildMdlForNonPagedPool(mdl);

devobj = get_original_devobj(DeviceObject, NULL);	// use original devobj!
	if (devobj == NULL) {
		KdPrint(("[tdi_fw] tdi_create_addrobj_complete: get_original_devobj!\n"));

status = STATUS_INVALID_PARAMETER;
		goto done;
	}

TdiBuildQueryInformation(query_irp, devobj, irps->FileObject,
		tdi_create_addrobj_complete2, ctx,
		TDI_QUERY_ADDRESS_INFO, mdl);

status = IoCallDriver(devobj, query_irp);

if (status == STATUS_PENDING)
	{
		KeWaitForSingleObject(&CompEvent, Executive, KernelMode, FALSE, NULL);
	} else if (status != STATUS_SUCCESS) {
		KdPrint(("[tdi_fw] tdi_create_addrobj_complete: IoCallDriver: 0x%x\n", status));
		goto done;
	}

query_irp = NULL;
	mdl = NULL;
	ctx = NULL;

status = STATUS_SUCCESS;

Irp->IoStatus.Status = status;
	
	if (status != STATUS_SUCCESS) {
		// tdi_create failed - remove fileobj from hash
		ot_del_fileobj(irps->FileObject, NULL);
	}
	
	return status;
	//return tdi_generic_complete(DeviceObject, Irp, Context);
}

/* this completion routine gets address and port from reply to 
       TDI_QUERY_ADDRESS_INFO */
NTSTATUS
tdi_create_addrobj_complete2
          (IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp, IN PVOID Context)
{
	NTSTATUS status;
	TDI_CREATE_ADDROBJ2_CTX *ctx = (TDI_CREATE_ADDROBJ2_CTX *)Context;
	TA_ADDRESS *addr = ctx->tai->Address.Address;
	struct ot_entry *ote_addr;
	KIRQL irql;
	int ipproto;

KdPrint(("[tdi_fw] tdi_create_addrobj_complete2: : %x:%addressu\n", 
		 ntohl(((TDI_ADDRESS_IP *)(addr->Address))->in_addr),
		 ntohs(((TDI_ADDRESS_IP *)(addr->Address))->sin_port)));

// 중간 생략

status = STATUS_SUCCESS;
done:
	if (ote_addr != NULL)
		KeReleaseSpinLock(&g_ot_hash_guard, irql);

if (Irp->PendingReturned == TRUE)
	{
 		if (ctx->pEvent != NULL)
 		{
 			KeSetEvent(ctx->pEvent, IO_NO_INCREMENT, FALSE);
 		}
	}

// cleanup MDL to avoid unlocking pages from NonPaged pool
	if (Irp->MdlAddress != NULL) {
		IoFreeMdl(Irp->MdlAddress);
		Irp->MdlAddress = NULL;
	}

free(ctx->tai);
	free(ctx);

// success anyway
	return STATUS_MORE_PROCESSING_REQUIRED;
}

대충 보시면... 수정 내용은 다음과 같습니다.

1. Dispatch 루틴(tdi_create_addrobj_complete)에서 TDI_CREATE_ADDROBJ2_CTX 구조체에 pEvent 멤버를 추가하여 IoCallDriver 전에 이벤트를 초기화하고 구조체에 Setting해줍니다.

2. IoCallDriver가 STATUS_PENDING을 리턴한 경우 KeWaitForSingle... 로 대기 처리를 해줍니다.

3. Completion 루틴 (tdi_create_addrobj_complete2) 에서는 Pending Flag를 체크하여 SetEvent해준 후 STATUS_MORE_PROCESSING_REQUIRED 를 리턴해줍니다.

4. Dispatch 루틴에서는 대기 처리를 해 주었기 때문에 IoMarkPending을 Propargate할 필요가 없습니다. (없을 것... 같은데요?? ^^;;)

=================================================================================================================

문제 2 : Event Receive Routine

문제의 코드는 다음과 같습니다.

NTSTATUS
tdi_event_receive(
    IN PVOID TdiEventContext,
    IN CONNECTION_CONTEXT ConnectionContext,
    IN ULONG ReceiveFlags,
    IN ULONG BytesIndicated,
    IN ULONG BytesAvailable,
    OUT ULONG *BytesTaken,
    IN PVOID Tsdu,
    OUT PIRP *IoRequestPacket)
{
	TDI_EVENT_CONTEXT *ctx = (TDI_EVENT_CONTEXT *)TdiEventContext;
	PFILE_OBJECT connobj = ot_find_conn_ctx(ctx->fileobj, ConnectionContext);
	NTSTATUS status;

KdPrint(("[tdi_fw] tdi_event_receive: addrobj 0x%x; connobj: 0x%x; %u/%u; flags: 0x%x\n",
		ctx->fileobj, connobj, BytesIndicated, BytesAvailable, ReceiveFlags));

// status 에 대한 체크가 누락된 부분!!
	status = ((PTDI_IND_RECEIVE)(ctx->old_handler))
		(ctx->old_context, ConnectionContext, ReceiveFlags, BytesIndicated,
		BytesAvailable, BytesTaken, Tsdu, IoRequestPacket);

KdPrint(("[tdi_fw] tdi_event_receive: status 0x%x; BytesTaken: %u; Irp: 0x%x\n",
		status, *BytesTaken, *IoRequestPacket));

if (*BytesTaken != 0) {
		struct ot_entry *ote_conn;
		KIRQL irql;

ote_conn = ot_find_fileobj(connobj, &irql);
		if (ote_conn != NULL) {
			ULONG bytes = *BytesTaken;

ote_conn->bytes_in += bytes;

// traffic stats
			KeAcquireSpinLockAtDpcLevel(&g_traffic_guard);
			
			g_traffic[TRAFFIC_TOTAL_IN] += bytes;
			
			if (ote_conn->log_disconnect)
				g_traffic[TRAFFIC_COUNTED_IN] += bytes;
			
			KeReleaseSpinLockFromDpcLevel(&g_traffic_guard);

KeReleaseSpinLock(&g_ot_hash_guard, irql);
		}
	}

if (*IoRequestPacket != NULL) {
	    // got IRP. replace completion.
	    struct tdi_client_irp_ctx *new_ctx;
		PIO_STACK_LOCATION irps = IoGetCurrentIrpStackLocation(*IoRequestPacket);

new_ctx = (struct tdi_client_irp_ctx *)malloc_np(sizeof(*new_ctx));
		if (new_ctx != NULL) {

new_ctx->connobj = connobj;

if (irps->CompletionRoutine != NULL) {
				new_ctx->completion = irps->CompletionRoutine;
				new_ctx->context = irps->Context;
				new_ctx->old_control = irps->Control;

} else {

// we don't use IoSetCompletionRoutine 
				// because it uses next not current location
				new_ctx->completion = NULL;
				new_ctx->context = NULL;

}

irps->CompletionRoutine = tdi_client_irp_complete;
			irps->Context = new_ctx;
			irps->Control = 
			SL_INVOKE_ON_SUCCESS | SL_INVOKE_ON_ERROR | SL_INVOKE_ON_CANCEL;
		}
	}

return status;
}

실제로 위의 코드를 실행해보면 두가지 문제가 있는데요...

1. Original Event Handler를 호출한 다음 결과 (status)값에 대한 검증이 없습니다. 이 결과가 SUCCESS가 아니라면 계속 진행하면 안되겠죠.

2. MSDN에 따르면 Original Event Handler의 마지막 OUT 인자인 *IoRequestPacket 에는 NULL 혹은 TDI Client에서 할당한 IRP 값이 들어있어야 하지만... Win7에서 돌려본 결과... '1'이라는 숫자값이나 '0xFFFFFFFF', 혹은... 아예 이상한 주소값이 들어있는 등... 문서대로 동작하지 않는 현상이 발견되었습니다. 그래서 결과는? 당근 BSOD가 뜹니다. ㅡ.ㅡ

이 부분은 어떻게 처리해야 할지 몰라서... 저는 일단 아래와 같이 주석처리를 했습니다. 단순한 방화벽 같은 프로그램이라면 이정도로도 충분하리라고 봅니다만... Incoming 패킷을 재조합해야 하거나 트래픽을 카운트해야 하는 프로그램이라면.. 이 문제를 어떻게든 처리해야 할 것으로 보여집니다.

이 문제... 근본적으로 해결하신 분이 계시면... 공유 좀 부탁드려요~~ ^^;;;

KdPrint(("[tdi_fw] tdi_event_receive: addrobj 0x%x; connobj: 0x%x; %u/%u; flags: 0x%x\n",
		ctx->fileobj, connobj, BytesIndicated, BytesAvailable, ReceiveFlags));

status = ((PTDI_IND_RECEIVE)(ctx->old_handler))
		(ctx->old_context, ConnectionContext, ReceiveFlags, BytesIndicated,
		BytesAvailable, BytesTaken, Tsdu, IoRequestPacket);
	
	if (!NT_SUCCESS(status))
	{
		return status;
	}

KdPrint(("[tdi_fw] tdi_event_receive: status 0x%x; BytesTaken: %u; Irp: 0x%x\n",
		status, *BytesTaken, *IoRequestPacket));

if (*BytesTaken != 0) {
		struct ot_entry *ote_conn;
		KIRQL irql;

ote_conn = ot_find_fileobj(connobj, &irql);
		if (ote_conn != NULL) {
			ULONG bytes = *BytesTaken;

ote_conn->bytes_in += bytes;

KeReleaseSpinLock(&g_ot_hash_guard, irql);
		}
	}
/*	
	if (*IoRequestPacket != NULL) {
	    // got IRP. replace completion.
	    struct tdi_client_irp_ctx *new_ctx;
		PIO_STACK_LOCATION irps = IoGetCurrentIrpStackLocation(*IoRequestPacket);

new_ctx = (struct tdi_client_irp_ctx *)malloc_np(sizeof(*new_ctx));
		if (new_ctx != NULL) {

new_ctx->connobj = connobj;

if (irps->CompletionRoutine != NULL) {
				new_ctx->completion = irps->CompletionRoutine;
				new_ctx->context = irps->Context;
				new_ctx->old_control = irps->Control;

} else {

// we don't use IoSetCompletionRoutine 
				// because it uses next not current location

new_ctx->completion = NULL;
				new_ctx->context = NULL;

}

irps->CompletionRoutine = tdi_client_irp_complete;
			irps->Context = new_ctx;
			irps->Control = 
			  SL_INVOKE_ON_SUCCESS | SL_INVOKE_ON_ERROR | SL_INVOKE_ON_CANCEL;
		}
	} 
	*/

return status;
}

Posted by kuaaan

달토끼 대박나라~!! ^^

TDIFW를 Vista 이후에서 동작시키기

카테고리

최근에 올라온 글

최근에 달린 댓글

링크

티스토리툴바