tditest.c

ReactOS是一些高手根据Windows XP的内核编写出的类XP。内核实现机理和API函数调用几乎相同。甚至可以兼容XP的程序。喜欢研究系统内核的人可以看一看。

 *     Buffer          = Pointer to buffer with data to send
 *     BufferSize      = Length of Buffer
 * RETURNS:
 *     Status of operation
 */
{
	PIRP Irp;
	PMDL Mdl;
	PDEVICE_OBJECT DeviceObject;
	PTDI_CONNECTION_INFORMATION ConnectInfo;
	PTA_IP_ADDRESS TA;
	PTDI_ADDRESS_IP IpAddress;
	IO_STATUS_BLOCK Iosb;
	NTSTATUS Status;

	DeviceObject = IoGetRelatedDeviceObject(TransportObject);
	ConnectInfo  = (PTDI_CONNECTION_INFORMATION)
		ExAllocatePool(NonPagedPool,
		sizeof(TDI_CONNECTION_INFORMATION) +
		sizeof(TA_IP_ADDRESS));

	if (!ConnectInfo)
		return STATUS_INSUFFICIENT_RESOURCES;

	RtlZeroMemory(ConnectInfo, sizeof(TDI_CONNECTION_INFORMATION) + sizeof(TA_IP_ADDRESS));

	ConnectInfo->RemoteAddressLength = sizeof(TA_IP_ADDRESS);
	ConnectInfo->RemoteAddress       = (PUCHAR) ((ULONG)ConnectInfo + sizeof(TDI_CONNECTION_INFORMATION));

	TA = (PTA_IP_ADDRESS)(ConnectInfo->RemoteAddress);
	TA->TAAddressCount           = 1;
	TA->Address[0].AddressLength = sizeof(TDI_ADDRESS_IP);
	TA->Address[0].AddressType   = TDI_ADDRESS_TYPE_IP;
	IpAddress           = (PTDI_ADDRESS_IP)(TA->Address[0].Address);
	IpAddress->sin_port = WH2N(Port);
	IpAddress->in_addr  = DH2N(Address);
	Irp = TdiBuildInternalDeviceControlIrp(
		TDI_SEND_DATAGRAM,   /* Sub function */
		DeviceObject,        /* Device object */
		TransportObject,     /* File object */
		NULL,                /* Event */
		NULL);               /* Return buffer */

	if (!Irp)
		{
			TDI_DbgPrint(MIN_TRACE, ("TdiBuildInternalDeviceControlIrp() failed.\n"));
			ExFreePool(ConnectInfo);
			return STATUS_INSUFFICIENT_RESOURCES;
		}

	Mdl = IoAllocateMdl(
		Buffer,     /* Virtual address of buffer */
		BufferSize, /* Length of buffer */
		FALSE,      /* Not secondary */
		FALSE,      /* Don't charge quota */
		NULL);      /* Don't use IRP */

	if (!Mdl)
		{
			TDI_DbgPrint(MIN_TRACE, ("IoAllocateMdl() failed.\n"));
			IoFreeIrp(Irp);
			ExFreePool(ConnectInfo);
			return STATUS_INSUFFICIENT_RESOURCES;
		}

#ifdef _MSC_VER
	try
		{
#endif
			MmProbeAndLockPages(Mdl, KernelMode, IoModifyAccess);
#ifdef _MSC_VER
		}
	except(EXCEPTION_EXECUTE_HANDLER)
		{
			TDI_DbgPrint(MIN_TRACE, ("MmProbeAndLockPages() failed.\n"));
			IoFreeMdl(Mdl);
			IoFreeIrp(Irp);
			ExFreePool(ConnectInfo);
			return STATUS_UNSUCCESSFUL;
		}
#endif

	TdiBuildSendDatagram(
		Irp,               /* I/O Request Packet */
		DeviceObject,      /* Device object */
		TransportObject,   /* File object */
		NULL,              /* Completion routine */
		NULL,              /* Completion context */
		Mdl,               /* Descriptor for data buffer */
		BufferSize,        /* Size of data to send */
		ConnectInfo);      /* Connection information */

	Status = TdiCall(Irp, DeviceObject, &Iosb, FALSE);

	ExFreePool(ConnectInfo);

	return Status;
}


NTSTATUS TdiReceiveDatagram(
    PFILE_OBJECT TransportObject,
    USHORT Port,
    PULONG Address,
    PUCHAR Buffer,
    PULONG BufferSize)
/*
 * FUNCTION: Receives a datagram
 * ARGUMENTS:
 *     TransportObject = Pointer to transport object
 *     Port            = Port to receive on
 *     Address         = Address of buffer to place remote address
 *     Buffer          = Address of buffer to place received data
 *     BufferSize      = Address of buffer with length of Buffer (updated)
 * RETURNS:
 *     Status of operation
 */
{
	PTDI_CONNECTION_INFORMATION ReceiveInfo;
	PTDI_CONNECTION_INFORMATION ReturnInfo;
	PTA_IP_ADDRESS ReturnAddress;
	PDEVICE_OBJECT DeviceObject;
	PTDI_ADDRESS_IP IpAddress;
	IO_STATUS_BLOCK Iosb;
	PVOID MdlBuffer;
	NTSTATUS Status;
	PIRP Irp;
	PMDL Mdl;

	DeviceObject = IoGetRelatedDeviceObject(TransportObject);
	if (!DeviceObject)
		return STATUS_INVALID_PARAMETER;

	ReceiveInfo = (PTDI_CONNECTION_INFORMATION) ExAllocatePool(NonPagedPool,
		sizeof(TDI_CONNECTION_INFORMATION) +
		sizeof(TDI_CONNECTION_INFORMATION) +
		sizeof(TA_IP_ADDRESS));

	if (!ReceiveInfo)
		return STATUS_INSUFFICIENT_RESOURCES;

	MdlBuffer = ExAllocatePool(PagedPool, *BufferSize);
	if (!MdlBuffer)
		return STATUS_INSUFFICIENT_RESOURCES;

	RtlZeroMemory(ReceiveInfo, sizeof(TDI_CONNECTION_INFORMATION) + sizeof(TDI_CONNECTION_INFORMATION) +
		sizeof(TA_IP_ADDRESS));

	RtlCopyMemory(MdlBuffer, Buffer, *BufferSize);

	/* Receive from any address */
	ReceiveInfo->RemoteAddressLength = 0;
	ReceiveInfo->RemoteAddress       = NULL;

	ReturnInfo = (PTDI_CONNECTION_INFORMATION) ((ULONG)ReceiveInfo + sizeof(TDI_CONNECTION_INFORMATION));
	ReturnInfo->RemoteAddressLength = sizeof(TA_IP_ADDRESS);
	ReturnInfo->RemoteAddress       = (PUCHAR) ((ULONG)ReturnInfo + sizeof(TDI_CONNECTION_INFORMATION));

	ReturnAddress = (PTA_IP_ADDRESS)(ReturnInfo->RemoteAddress);
	ReturnAddress->TAAddressCount           = 1;
	ReturnAddress->Address[0].AddressLength = sizeof(TDI_ADDRESS_IP);
	ReturnAddress->Address[0].AddressType   = TDI_ADDRESS_TYPE_IP;

	IpAddress = (PTDI_ADDRESS_IP)(ReturnAddress->Address[0].Address);
	IpAddress->sin_port = WH2N(Port);
	IpAddress->in_addr  = DH2N(LocalAddress);

	Irp = TdiBuildInternalDeviceControlIrp(
		TDI_RECEIVE_DATAGRAM,    /* Sub function */
		DeviceObject,            /* Device object */
		TransportObject,         /* File object */
		NULL,                    /* Event */
		NULL);                   /* Return buffer */

	if (!Irp)
		{
			ExFreePool(MdlBuffer);
			ExFreePool(ReceiveInfo);
			return STATUS_INSUFFICIENT_RESOURCES;
		}

	Mdl = IoAllocateMdl(
		MdlBuffer,      /* Virtual address */
		*BufferSize,    /* Length of buffer */
		FALSE,          /* Not secondary */
		FALSE,          /* Don't charge quota */
		NULL);          /* Don't use IRP */

	if (!Mdl)
		{
			IoFreeIrp(Irp);
			ExFreePool(MdlBuffer);
			ExFreePool(ReceiveInfo);
			return STATUS_INSUFFICIENT_RESOURCES;
		}

#ifdef _MSC_VER
	try
		{
#endif
			MmProbeAndLockPages(Mdl, KernelMode, IoModifyAccess);
#ifdef _MSC_VER
		}
	except (EXCEPTION_EXECUTE_HANDLER)
		{
			TDI_DbgPrint(MIN_TRACE, ("MmProbeAndLockPages() failed.\n"));
			IoFreeMdl(Mdl);
			IoFreeIrp(Irp);
			ExFreePool(MdlBuffer);
			ExFreePool(ReceiveInfo);
			return STATUS_INSUFFICIENT_RESOURCES;
		}
#endif

	TdiBuildReceiveDatagram(
		Irp,                    /* I/O Request Packet */
		DeviceObject,           /* Device object */
		TransportObject,        /* File object */
		NULL,                   /* Completion routine */
		NULL,                   /* Completion context */
		Mdl,                    /* Data buffer */
		*BufferSize,            /* Size of data buffer */
		ReceiveInfo,            /* Connection information */
		ReturnInfo,             /* Connection information */
		TDI_RECEIVE_NORMAL);    /* Flags */

	Status = TdiCall(Irp, DeviceObject, &Iosb, TRUE);

	if (NT_SUCCESS(Status))
		{
			RtlCopyMemory(Buffer, MdlBuffer, Iosb.Information);
			*BufferSize = Iosb.Information;
			*Address    = DN2H(IpAddress->in_addr);
		}

	ExFreePool(MdlBuffer);
	ExFreePool(ReceiveInfo);

	return Status;
}


VOID TdiSendThread(
    PVOID Context)
/*
 * FUNCTION: Send thread
 * ARGUMENTS:
 *     Context = Pointer to context information
 * NOTES:
 *     Transmits an UDP packet every two seconds to ourselves on the chosen port
 */
{
	KEVENT Event;
	PKEVENT Events[2];
	LARGE_INTEGER Timeout;
	NTSTATUS Status = STATUS_SUCCESS;
	UCHAR Data[40]  = "Testing one, two, three, ...";

	if (!OpenError)
		{
			Timeout.QuadPart = 10000000L;           /* Second factor */
			Timeout.QuadPart *= 2;                  /* Number of seconds */
			Timeout.QuadPart = -(Timeout.QuadPart); /* Relative time */

			KeInitializeEvent(&Event, SynchronizationEvent, FALSE);

			Events[0] = &StopEvent;
			Events[1] = &Event;

			while (NT_SUCCESS(Status))
				{
					/* Wait until timeout or stop flag is set */
 					KeWaitForMultipleObjects( 2, (PVOID)Events, WaitAny, Executive, KernelMode, FALSE, &Timeout, NULL);

					if (KeReadStateEvent(&StopEvent) != 0)
						{
							TDI_DbgPrint(MAX_TRACE, ("Received terminate signal...\n"));
							break;
						}

					DbgPrint("Sending data - '%s'\n", Data);

					Status = TdiSendDatagram(TdiTransportObject, TEST_PORT, LocalAddress, Data, sizeof(Data));

					if (!NT_SUCCESS(Status))
						DbgPrint("Failed sending data (Status = 0x%X)\n", Status);
				}
		}

	TDI_DbgPrint(MAX_TRACE, ("Terminating send thread...\n"));

	PsTerminateSystemThread(STATUS_SUCCESS);
}


VOID TdiReceiveThread(
    PVOID Context)
/*
 * FUNCTION: Receive thread
 * ARGUMENTS:
 *     Context = Pointer to context information
 * NOTES:
 *     Waits until an UDP packet is received on the chosen endpoint and displays the data
 */
{
	ULONG Address;
	UCHAR Data[40];
	ULONG Size;
	NTSTATUS Status = STATUS_SUCCESS;

	if (!OpenError)
		{
			while (NT_SUCCESS(Status))
				{
					Size = sizeof(Data);
					RtlZeroMemory(Data, Size);

					Status = TdiReceiveDatagram(TdiTransportObject, TEST_PORT, &Address, Data, &Size);

					if (NT_SUCCESS(Status))
						{
							DbgPrint("Received data - '%s'\n", Data);
						}
					else
						if (Status != STATUS_CANCELLED)
							{
								TDI_DbgPrint(MIN_TRACE, ("Receive error (Status = 0x%X).\n", Status));
							}
						else
							{
								TDI_DbgPrint(MAX_TRACE, ("IRP was cancelled.\n"));
							}
				}
		}

	TDI_DbgPrint(MAX_TRACE, ("Terminating receive thread...\n"));

	PsTerminateSystemThread(STATUS_SUCCESS);
}


VOID TdiOpenThread(
    PVOID Context)
/*
 * FUNCTION: Open thread
 * ARGUMENTS:
 *     Context = Pointer to context information (event)
 */
{
	NTSTATUS Status;

	TDI_DbgPrint(MAX_TRACE, ("Called.\n"));

	OpenError = TRUE;

	Status = TdiOpenTransport(UDP_DEVICE_NAME, TEST_PORT, &TdiTransport, &TdiTransportObject);

	if (NT_SUCCESS(Status))
		{
			Status = TdiQueryAddress(TdiTransportObject, &LocalAddress);

			if (NT_SUCCESS(Status))
				{
					OpenError = FALSE;
					DbgPrint("Using local IP address 0x%X\n", LocalAddress);
				}
			else
				{
					TDI_DbgPrint(MIN_TRACE, ("Unable to determine local IP address.\n"));
				}
			}
	else
		TDI_DbgPrint(MIN_TRACE, ("Cannot open transport (Status = 0x%X).\n", Status));

	TDI_DbgPrint(MAX_TRACE, ("Setting close event.\n"));

	KeSetEvent((PKEVENT)Context, 0, FALSE);

	TDI_DbgPrint(MIN_TRACE, ("Leaving.\n"));
}


VOID TdiUnload(
    PDRIVER_OBJECT DriverObject)
/*
 * FUNCTION: Unload routine
 * ARGUMENTS:
 *     DriverObject = Pointer to a driver object for this driver
 */
{
	PVOID ReceiveThreadObject = 0;
	PVOID SendThreadObject = 0;

	TDI_DbgPrint(MAX_TRACE, ("Setting stop flag\n"));

	/* Get pointers to the thread objects */
	ObReferenceObjectByHandle(SendThread, THREAD_ALL_ACCESS, NULL, KernelMode, &SendThreadObject, NULL);
	ObReferenceObjectByHandle(ReceiveThread, THREAD_ALL_ACCESS, NULL, KernelMode, &ReceiveThreadObject, NULL);

	KeSetEvent(&StopEvent, 0, FALSE);

	/* Wait for send thread to stop */
	KeWaitForSingleObject(SendThreadObject, Executive, KernelMode, FALSE, NULL);

	/* Wait for receive thread to stop */
	KeWaitForSingleObject(ReceiveThreadObject, Executive, KernelMode, FALSE, NULL);

	/* Close device */
	TdiCloseDevice(TdiTransport, TdiTransportObject);
}


NTSTATUS
#ifndef _MSC_VER
STDCALL
#endif
DriverEntry(
    PDRIVER_OBJECT DriverObject,
    PUNICODE_STRING RegistryPath)
/*
 * FUNCTION: Main driver entry point
 * ARGUMENTS:
 *     DriverObject = Pointer to a driver object for this driver
 *     RegistryPath = Registry node for configuration parameters
 * RETURNS:
 *     Status of driver initialization
 */
{
	KEVENT Event;
	NTSTATUS Status;
	WORK_QUEUE_ITEM WorkItem;

	KeInitializeEvent(&StopEvent, NotificationEvent, FALSE);

	/* Call TdiOpenThread() */
	KeInitializeEvent(&Event, SynchronizationEvent, FALSE);
	ExInitializeWorkItem(&WorkItem, TdiOpenThread, &Event);
 	ExQueueWorkItem(&WorkItem, DelayedWorkQueue);
	KeWaitForSingleObject(&Event, Executive, KernelMode, TRUE, NULL);

	/* Create a UDP send thread that sends a dgram every 2 seconds */
	Status = PsCreateSystemThread(
		&SendThread,                      /* Thread handle */
		0,                                /* Desired access */
		NULL,                             /* Object attributes */
		NULL,                             /* Process handle */
		NULL,                             /* Client id */
		(PKSTART_ROUTINE)TdiSendThread,   /* Start routine */
		NULL);                            /* Start context */

	if (!NT_SUCCESS(Status))
		{
			TDI_DbgPrint(MIN_TRACE, ("PsCreateSystemThread() failed for send thread (Status = 0x%X).\n", Status));
			return STATUS_INSUFFICIENT_RESOURCES;
		}

	/* Create a UDP receive thread */
	Status = PsCreateSystemThread(
		&ReceiveThread,                       /* Thread handle */
		0,                                    /* Desired access */
		NULL,                                 /* Object attributes */
		NULL,                                 /* Process handle */
		NULL,                                 /* Client id */
		(PKSTART_ROUTINE)TdiReceiveThread,    /* Start routine */
		NULL);                                /* Start context */

	if (!NT_SUCCESS(Status))
		{
			TDI_DbgPrint(MIN_TRACE, ("PsCreateSystemThread() failed for receive thread (Status = 0x%X).\n", Status));
			ZwClose(SendThread);
			return STATUS_INSUFFICIENT_RESOURCES;
		}

	DriverObject->DriverUnload = (PDRIVER_UNLOAD)TdiUnload;

	return STATUS_SUCCESS;
}

/* EOF */