1. 程式人生 > >TDI簡單的TCP網路通訊

TDI簡單的TCP網路通訊

#include <ntifs.h>
#include <ntddk.h>
#include <Ntstrsafe.h>
#include <tdi.h>
#include <tdikrnl.h>
#include <stdlib.h>
#include "KernelSocket.h"

#define SENDBUFFLEN			4096
#define RECVBUFFERLEN		32768

#define DELAY_ONE_MICROSECOND (-10)
#define DELAY_ONE_MILLISECOND (DELAY_ONE_MICROSECOND*1000)


// 建立並連線TDI
NTSTATUS OpenTDIConnection(char* szIpAddress, unsigned short Port);
unsigned long SendData(char* pData, unsigned long bufferLength);
unsigned long  RecvData(char* pData, unsigned long Length);

NTSTATUS CloseTDIConnection();
void FreeConnection();

VOID Sleep(unsigned long msec);

PDEVICE_OBJECT				pDeviceObject = NULL;
PFILE_OBJECT				pFileObject = NULL;
PFILE_FULL_EA_INFORMATION	pFileInfo = NULL;

VOID DriverUnload(IN PDRIVER_OBJECT DriverObject)
{

	return;
}

NTSTATUS DriverEntry(IN PDRIVER_OBJECT DriverObject, IN PUNICODE_STRING RegistryPath)
{

	DbgBreakPoint();
	DriverObject->DriverUnload = DriverUnload;

	//連線伺服器
	if (!NT_SUCCESS(OpenTDIConnection("14.17.32.211", 80)))
	{
		FreeConnection();
	}

	//傳送資料
	unsigned char pData[] = "GET / HTTP/1.1\r\nAccept: */*\r\nAccept-Language: zh-cn\r\nAccept-Encoding: gzip, deflate\r\nUser-Agent: Mozilla/4.0\r\nHost:www.qq.com\r\nConnection: Keep-Alive\r\n\r\n";
	if (NT_SUCCESS(SendData(pData, sizeof(pData))))
	{
		Sleep(2000);
		char* pRecvBuffer = ExAllocatePool(NonPagedPool, RECVBUFFERLEN);
		if (pRecvBuffer)
		{
			//接受資料
			RtlZeroMemory(pRecvBuffer, RECVBUFFERLEN);
			if (RecvData(pRecvBuffer, RECVBUFFERLEN)>0)
			{
				KdPrint(("%s\n", pRecvBuffer));
			}
			ExFreePool(pRecvBuffer);

		}
		CloseTDIConnection();
	}

	return STATUS_SUCCESS;
}


static NTSTATUS TDICompletionRoutine(IN PDEVICE_OBJECT theDeviceObject, IN PIRP theIrp, IN PVOID theContextP)
{

	if (theContextP != NULL)
		KeSetEvent((PKEVENT)theContextP, 0, FALSE);
	return STATUS_MORE_PROCESSING_REQUIRED;	//IO操作忙,是不完整的
}

// 建立並連線TDI
NTSTATUS OpenTDIConnection(char* szIpAddress, unsigned short Port)
{

	NTSTATUS					status;
	UNICODE_STRING				TdiTransportDeviceName;
	OBJECT_ATTRIBUTES           TdiAttributes;
	HANDLE						TdiAddressHandle;
	HANDLE						TdiEndpointHandle;
	IO_STATUS_BLOCK             IoStatusBlock;
	PTA_IP_ADDRESS				pAddress;
	CONNECTION_CONTEXT			connectionContext = NULL;
	ULONG						eaSize;
	PIRP						pIrp;
	PVOID						pAddressFileObject;
	KEVENT                      irpCompleteEvent;
	KEVENT						connectionEvent;
	TA_IP_ADDRESS				controllerTaIpAddress = {0};
	TDI_CONNECTION_INFORMATION  controllerConnection = {0};


	//引數效驗
	if (szIpAddress == NULL || Port <= 0)return STATUS_UNSUCCESSFUL;


	static char eaBuffer[sizeof(FILE_FULL_EA_INFORMATION) + TDI_TRANSPORT_ADDRESS_LENGTH + sizeof(TA_IP_ADDRESS)] = {0};
	PFILE_FULL_EA_INFORMATION	pEaBuffer = (PFILE_FULL_EA_INFORMATION)eaBuffer;
	pEaBuffer->NextEntryOffset = 0;
	pEaBuffer->Flags = 0;
	pEaBuffer->EaNameLength = TDI_TRANSPORT_ADDRESS_LENGTH;
	RtlCopyMemory(pEaBuffer->EaName,TdiTransportAddress,pEaBuffer->EaNameLength + 1);

	pEaBuffer->EaValueLength = sizeof(TA_IP_ADDRESS);
	pAddress = (PTA_IP_ADDRESS)(pEaBuffer->EaName + pEaBuffer->EaNameLength + 1);
	pAddress->TAAddressCount = 1;
	pAddress->Address[0].AddressLength = TDI_ADDRESS_LENGTH_IP;
	pAddress->Address[0].AddressType = TDI_ADDRESS_TYPE_IP;
	pAddress->Address[0].Address[0].sin_port = 0; // any port
	pAddress->Address[0].Address[0].in_addr = 0; // local address
	RtlZeroMemory(pAddress->Address[0].Address[0].sin_zero, 0,sizeof(pAddress->Address[0].Address[0].sin_zero));


	// 開啟裝置L"\\Device\\Tcp"
	RtlInitUnicodeString(&TdiTransportDeviceName, L"\\Device\\Tcp");
	InitializeObjectAttributes(&TdiAttributes,&TdiTransportDeviceName,OBJ_CASE_INSENSITIVE | OBJ_KERNEL_HANDLE,0,0);
	status = ZwCreateFile(&TdiAddressHandle,GENERIC_READ | GENERIC_WRITE | SYNCHRONIZE,&TdiAttributes,&IoStatusBlock,0,FILE_ATTRIBUTE_NORMAL,FILE_SHARE_READ,FILE_OPEN,0,pEaBuffer,sizeof(eaBuffer));
	if (!NT_SUCCESS(status))
	{
		return STATUS_UNSUCCESSFUL;
	}

	status = ObReferenceObjectByHandle(TdiAddressHandle,FILE_ANY_ACCESS,0,KernelMode,(PVOID *)&pAddressFileObject,NULL);
	if (!NT_SUCCESS(status))
	{
		return STATUS_UNSUCCESSFUL;
	}

	eaSize = FIELD_OFFSET(FILE_FULL_EA_INFORMATION, EaName) +TDI_CONNECTION_CONTEXT_LENGTH + 1 +sizeof(CONNECTION_CONTEXT);
	pFileInfo = (PFILE_FULL_EA_INFORMATION)ExAllocatePool(NonPagedPool, eaSize);
	if (pFileInfo == NULL)
	{
		return STATUS_INSUFFICIENT_RESOURCES;
	}
	RtlZeroMemory(pFileInfo, 0, eaSize);
	pFileInfo->NextEntryOffset = 0;
	pFileInfo->Flags = 0;
	pFileInfo->EaNameLength = TDI_CONNECTION_CONTEXT_LENGTH;
	RtlCopyMemory(pFileInfo->EaName,TdiConnectionContext,pFileInfo->EaNameLength + 1);


	pFileInfo->EaValueLength = sizeof(CONNECTION_CONTEXT);
	*(CONNECTION_CONTEXT*)(pFileInfo->EaName + (pFileInfo->EaNameLength + 1)) =(CONNECTION_CONTEXT)connectionContext;

	status = ZwCreateFile(&TdiEndpointHandle,GENERIC_READ | GENERIC_WRITE | SYNCHRONIZE,&TdiAttributes,&IoStatusBlock,0,FILE_ATTRIBUTE_NORMAL,FILE_SHARE_READ,FILE_OPEN,0,pFileInfo,sizeof(eaBuffer));
	if (!NT_SUCCESS(status))
	{
		return STATUS_UNSUCCESSFUL;
	}

	status = ObReferenceObjectByHandle(TdiEndpointHandle,FILE_ANY_ACCESS,0,KernelMode,(PVOID *)&pFileObject,NULL);
	if (!NT_SUCCESS(status))
	{
		return STATUS_UNSUCCESSFUL;
	}

	//獲取原始裝置物件
	pDeviceObject = IoGetRelatedDeviceObject(pAddressFileObject);
	if (pDeviceObject==NULL)
	{
		return STATUS_UNSUCCESSFUL;
	}

	KeInitializeEvent(&irpCompleteEvent, NotificationEvent, FALSE);
	//分配一個客戶一個IRP開始內部裝置的控制要求
	pIrp = TdiBuildInternalDeviceControlIrp(TDI_ASSOCIATE_ADDRESS,pDeviceObject,pFileObject,&irpCompleteEvent,&IoStatusBlock);
	if (pIrp == NULL)
	{
		return STATUS_INSUFFICIENT_RESOURCES;
	}
	//本地節點的客戶端已經打開了一個地址
	TdiBuildAssociateAddress(pIrp,pDeviceObject,pFileObject,NULL,NULL,TdiAddressHandle);
	IoSetCompletionRoutine(pIrp, TDICompletionRoutine, &irpCompleteEvent, TRUE, TRUE, TRUE);
	// Send the packet
	status = IoCallDriver(pDeviceObject, pIrp);
	if (status == STATUS_PENDING)
	{
		KeWaitForSingleObject(&irpCompleteEvent, Executive, KernelMode, FALSE, 0);
	}
	if ((status != STATUS_SUCCESS) &&(status != STATUS_PENDING))
	{
		return STATUS_UNSUCCESSFUL;
	}


	//建立一個tdi_connect要求底層傳輸
	KeInitializeEvent(&connectionEvent, NotificationEvent, FALSE);
	pIrp = TdiBuildInternalDeviceControlIrp(TDI_CONNECT,pDeviceObject,pFileObject,&connectionEvent,&IoStatusBlock);
	if (pIrp == NULL)
	{
		return STATUS_INSUFFICIENT_RESOURCES;
	}
	controllerTaIpAddress.TAAddressCount = 1;
	controllerTaIpAddress.Address[0].AddressLength = TDI_ADDRESS_LENGTH_IP;
	controllerTaIpAddress.Address[0].AddressType = TDI_ADDRESS_TYPE_IP;
	controllerTaIpAddress.Address[0].Address[0].sin_port = htons(Port);
	controllerTaIpAddress.Address[0].Address[0].in_addr = inet_addr(szIpAddress);
	controllerConnection.RemoteAddressLength = sizeof(controllerTaIpAddress);
	controllerConnection.RemoteAddress = &controllerTaIpAddress;
	TdiBuildConnect(pIrp,pDeviceObject,pFileObject,NULL,NULL,NULL,&controllerConnection,0);
	IoSetCompletionRoutine(pIrp, TDICompletionRoutine, &connectionEvent, TRUE, TRUE, TRUE);
	status = IoCallDriver(pDeviceObject, pIrp);
	if (status == STATUS_PENDING)
	{
		KeWaitForSingleObject(&connectionEvent, Executive, KernelMode, FALSE, 0);
	}

	//設定變數
	status = pIrp->IoStatus.Status;
	return status;
}

unsigned long SendData(char* pData, unsigned long bufferLength)
{

	NTSTATUS		status = STATUS_UNSUCCESSFUL;
	KEVENT			SendEvent;
	PIRP			pIrp=NULL;
	IO_STATUS_BLOCK	IoStatusBlock;
	unsigned long uSendSize = 0;


	//引數效驗
	if (pData == NULL || bufferLength <= 0)return status;
	char* pSendBuffer = ExAllocatePool(NonPagedPool, SENDBUFFLEN);
	if (pSendBuffer == NULL)return NULL;
	PMDL pSendMdl = IoAllocateMdl(pSendBuffer, SENDBUFFLEN, FALSE, FALSE, NULL);
	if (pSendMdl == NULL)return NULL;
	__try
	{
		MmProbeAndLockPages(pSendMdl, KernelMode, IoModifyAccess);
	}
	__except (EXCEPTION_EXECUTE_HANDLER)
	{
		return NULL;
	}

	do 
	{
		RtlZeroMemory(pSendBuffer, SENDBUFFLEN);
		RtlCopyMemory(pSendBuffer, pData, bufferLength);

		// build an IO Request Packet
		KeInitializeEvent(&SendEvent, NotificationEvent, FALSE);
		pIrp = TdiBuildInternalDeviceControlIrp(TDI_SEND, pDeviceObject, pFileObject, &SendEvent, &IoStatusBlock);
		if (pIrp == NULL)break;

		TdiBuildSend(pIrp, pDeviceObject, pFileObject, NULL, NULL, pSendMdl, 0, bufferLength);
		IoSetCompletionRoutine(pIrp, TDICompletionRoutine, &SendEvent, TRUE, TRUE, TRUE);
		status = IoCallDriver(pDeviceObject, pIrp);
		if (status == STATUS_PENDING)
			KeWaitForSingleObject(&SendEvent, Executive, KernelMode, FALSE, 0);

		uSendSize = pIrp->IoStatus.Information;
	} while (FALSE);


	ExFreePool(pSendBuffer);
	IoFreeMdl(pSendMdl);
	return uSendSize;

}

unsigned long  RecvData(char* pData, unsigned long Length)
{

	NTSTATUS		status;
	KEVENT			RecvEvent;
	PIRP			pIrp=NULL;
	IO_STATUS_BLOCK	IoStatusBlock;
	unsigned long uRecvSize = 0;


	//引數效驗
	if (MmIsAddressValid(pData) == FALSE || Length <= 0)return STATUS_UNSUCCESSFUL;
	PMDL pReceiveMdl = IoAllocateMdl(pData, RECVBUFFERLEN, FALSE, FALSE, NULL);
	__try
	{
		MmProbeAndLockPages(pReceiveMdl, KernelMode, IoModifyAccess);
	}
	__except (EXCEPTION_EXECUTE_HANDLER)
	{
		return STATUS_UNSUCCESSFUL;
	}



	do 
	{

		//分配recv
		KeInitializeEvent(&RecvEvent, NotificationEvent, FALSE);
		pIrp = TdiBuildInternalDeviceControlIrp(TDI_RECEIVE, pDeviceObject, pFileObject, &RecvEvent, &IoStatusBlock);
		if (pIrp == NULL)break;

		TdiBuildReceive(pIrp, pDeviceObject, pFileObject, NULL, NULL, pReceiveMdl, TDI_RECEIVE_NORMAL, Length);
		IoSetCompletionRoutine(pIrp, TDICompletionRoutine, &RecvEvent, TRUE, TRUE, TRUE);

		status = IoCallDriver(pDeviceObject, pIrp);
		if (status == STATUS_PENDING)
			KeWaitForSingleObject(&RecvEvent, Executive, KernelMode, FALSE, 0);
		if ((status != STATUS_SUCCESS) && (status != STATUS_PENDING))break;


		uRecvSize = pIrp->IoStatus.Information;
	} while (FALSE);



	IoFreeMdl(pReceiveMdl);
	return uRecvSize;
}

VOID Sleep(unsigned long msec)
{
	LARGE_INTEGER my_interval;
	my_interval.QuadPart = DELAY_ONE_MILLISECOND;
	my_interval.QuadPart *= msec;
	KeDelayExecutionThread(KernelMode, 0, &my_interval);
	return;
}

NTSTATUS CloseTDIConnection()
{
	
	NTSTATUS		status;
	KEVENT			DisconEvent;
	PIRP			pIrp;
	IO_STATUS_BLOCK	IoStatusBlock;
	KEVENT CompleteEvent;
	TDI_CONNECTION_INFORMATION RequestDisconnect;

	KeInitializeEvent(&DisconEvent, NotificationEvent, FALSE);
	pIrp = TdiBuildInternalDeviceControlIrp(TDI_DISCONNECT,pDeviceObject,pFileObject,&DisconEvent,&IoStatusBlock);
	if (pIrp == NULL)
	{
		return STATUS_INSUFFICIENT_RESOURCES;
	}

	TdiBuildDisconnect(pIrp,pDeviceObject,pFileObject,NULL,&CompleteEvent,NULL,TDI_DISCONNECT_ABORT,&RequestDisconnect,&RequestDisconnect);
	IoSetCompletionRoutine(pIrp, TDICompletionRoutine, &DisconEvent, TRUE, TRUE, TRUE);
	status = IoCallDriver(pDeviceObject, pIrp);
	if (status == STATUS_PENDING)
	{
		KeWaitForSingleObject(&DisconEvent, Executive, KernelMode, FALSE, 0);
	}

	if ((status != STATUS_SUCCESS) &&(status != STATUS_PENDING))
	{
		return STATUS_UNSUCCESSFUL;
	}


	if (pFileInfo != NULL)
	{
		ExFreePool(pFileInfo);
		pFileInfo = NULL;
	}


	status = pIrp->IoStatus.Status;
	return status;
}
void FreeConnection()
{
	if (pFileInfo != NULL)
	{
		ExFreePool(pFileInfo);
		pFileInfo = NULL;
	}
}