TDI简单的TCP网络通信

#include <ntifs.h>#include <ntddk.h>#include <Ntstrsafe.h>#include <tdi.h>#include <tdikrnl.h>#include <stdlib.h>#include "KernelSocket.h"#define SENDBUFFLEN4096#define RECVBUFFERLEN32768#define DELAY_ONE_MICROSECOND (-10)#define DELAY_ONE_MILLISECOND (DELAY_ONE_MICROSECOND*1000)// 建立并连接TDINTSTATUS 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_OBJECTpDeviceObject = NULL;PFILE_OBJECTpFileObject = NULL;PFILE_FULL_EA_INFORMATIONpFileInfo = 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操作忙，是不完整的}// 建立并连接TDINTSTATUS OpenTDIConnection(char* szIpAddress, unsigned short Port){NTSTATUSstatus;UNICODE_STRINGTdiTransportDeviceName;OBJECT_ATTRIBUTES           TdiAttributes;HANDLETdiAddressHandle;HANDLETdiEndpointHandle;IO_STATUS_BLOCK             IoStatusBlock;PTA_IP_ADDRESSpAddress;CONNECTION_CONTEXTconnectionContext = NULL;ULONGeaSize;PIRPpIrp;PVOIDpAddressFileObject;KEVENT                      irpCompleteEvent;KEVENTconnectionEvent;TA_IP_ADDRESScontrollerTaIpAddress = {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_INFORMATIONpEaBuffer = (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 portpAddress->Address[0].Address[0].in_addr = 0; // local addressRtlZeroMemory(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 packetstatus = 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){NTSTATUSstatus = STATUS_UNSUCCESSFUL;KEVENTSendEvent;PIRPpIrp=NULL;IO_STATUS_BLOCKIoStatusBlock;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 PacketKeInitializeEvent(&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){NTSTATUSstatus;KEVENTRecvEvent;PIRPpIrp=NULL;IO_STATUS_BLOCKIoStatusBlock;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 {//分配recvKeInitializeEvent(&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(){NTSTATUSstatus;KEVENTDisconEvent;PIRPpIrp;IO_STATUS_BLOCKIoStatusBlock;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;}}