TDI簡單的TCP網路通訊
阿新 • • 發佈:2019-02-12
#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; } }