nganhkhoa
87a61a625a
Add a simple device io routine to print routine code Fix variables due to recent Windows update to 19569
480 lines
14 KiB
C
480 lines
14 KiB
C
/*++
|
|
|
|
Module Name:
|
|
|
|
simplewsk.c
|
|
|
|
Abstract:
|
|
|
|
Wrapper library for WSK functions
|
|
|
|
Author:
|
|
|
|
MaD, 12-May-2009
|
|
|
|
--*/
|
|
|
|
#include "simplewsk.h"
|
|
|
|
static WSK_REGISTRATION g_WskRegistration;
|
|
static WSK_PROVIDER_NPI g_WskProvider;
|
|
static WSK_CLIENT_DISPATCH g_WskDispatch = {MAKE_WSK_VERSION(1, 0), 0, NULL};
|
|
|
|
enum { DEINITIALIZED, DEINITIALIZING, INITIALIZING, INITIALIZED };
|
|
|
|
static LONG g_SocketsState = DEINITIALIZED;
|
|
|
|
static NTSTATUS NTAPI CompletionRoutine(__in PDEVICE_OBJECT DeviceObject,
|
|
__in PIRP Irp,
|
|
__in PKEVENT CompletionEvent) {
|
|
ASSERT(CompletionEvent);
|
|
|
|
UNREFERENCED_PARAMETER(Irp);
|
|
UNREFERENCED_PARAMETER(DeviceObject);
|
|
|
|
KeSetEvent(CompletionEvent, IO_NO_INCREMENT, FALSE);
|
|
return STATUS_MORE_PROCESSING_REQUIRED;
|
|
}
|
|
|
|
static NTSTATUS InitWskData(__out PIRP* pIrp, __out PKEVENT CompletionEvent) {
|
|
ASSERT(pIrp);
|
|
ASSERT(CompletionEvent);
|
|
|
|
*pIrp = IoAllocateIrp(1, FALSE);
|
|
if (!*pIrp) {
|
|
KdPrint(("InitWskData(): IoAllocateIrp() failed\n"));
|
|
return STATUS_INSUFFICIENT_RESOURCES;
|
|
}
|
|
|
|
KeInitializeEvent(CompletionEvent, SynchronizationEvent, FALSE);
|
|
IoSetCompletionRoutine(*pIrp, CompletionRoutine, CompletionEvent, TRUE, TRUE,
|
|
TRUE);
|
|
return STATUS_SUCCESS;
|
|
}
|
|
|
|
static NTSTATUS InitWskBuffer(__in PVOID Buffer, __in ULONG BufferSize,
|
|
__out PWSK_BUF WskBuffer) {
|
|
NTSTATUS Status = STATUS_SUCCESS;
|
|
|
|
ASSERT(Buffer);
|
|
ASSERT(BufferSize);
|
|
ASSERT(WskBuffer);
|
|
|
|
WskBuffer->Offset = 0;
|
|
WskBuffer->Length = BufferSize;
|
|
|
|
WskBuffer->Mdl = IoAllocateMdl(Buffer, BufferSize, FALSE, FALSE, NULL);
|
|
if (!WskBuffer->Mdl) {
|
|
KdPrint(("InitWskBuffer(): IoAllocateMdl() failed\n"));
|
|
return STATUS_INSUFFICIENT_RESOURCES;
|
|
}
|
|
|
|
__try {
|
|
MmProbeAndLockPages(WskBuffer->Mdl, KernelMode, IoWriteAccess);
|
|
} __except (EXCEPTION_EXECUTE_HANDLER) {
|
|
KdPrint(("InitWskBuffer(): MmProbeAndLockPages(%p) failed\n", Buffer));
|
|
IoFreeMdl(WskBuffer->Mdl);
|
|
Status = STATUS_ACCESS_VIOLATION;
|
|
}
|
|
|
|
return Status;
|
|
}
|
|
|
|
static VOID FreeWskBuffer(__in PWSK_BUF WskBuffer) {
|
|
ASSERT(WskBuffer);
|
|
|
|
MmUnlockPages(WskBuffer->Mdl);
|
|
IoFreeMdl(WskBuffer->Mdl);
|
|
}
|
|
|
|
//
|
|
// Library initialization routine
|
|
//
|
|
|
|
NTSTATUS NTAPI WSKStartup() {
|
|
WSK_CLIENT_NPI WskClient = {0};
|
|
NTSTATUS Status = STATUS_UNSUCCESSFUL;
|
|
|
|
if (InterlockedCompareExchange(&g_SocketsState, INITIALIZING,
|
|
DEINITIALIZED) != DEINITIALIZED)
|
|
return STATUS_ALREADY_REGISTERED;
|
|
|
|
WskClient.ClientContext = NULL;
|
|
WskClient.Dispatch = &g_WskDispatch;
|
|
|
|
Status = WskRegister(&WskClient, &g_WskRegistration);
|
|
if (!NT_SUCCESS(Status)) {
|
|
KdPrint(("WskRegister() failed with status 0x%08X\n", Status));
|
|
InterlockedExchange(&g_SocketsState, DEINITIALIZED);
|
|
return Status;
|
|
}
|
|
|
|
Status =
|
|
WskCaptureProviderNPI(&g_WskRegistration, WSK_NO_WAIT, &g_WskProvider);
|
|
if (!NT_SUCCESS(Status)) {
|
|
KdPrint(("WskCaptureProviderNPI() failed with status 0x%08X\n", Status));
|
|
WskDeregister(&g_WskRegistration);
|
|
InterlockedExchange(&g_SocketsState, DEINITIALIZED);
|
|
return Status;
|
|
}
|
|
|
|
InterlockedExchange(&g_SocketsState, INITIALIZED);
|
|
return STATUS_SUCCESS;
|
|
}
|
|
|
|
//
|
|
// Library deinitialization routine
|
|
//
|
|
|
|
VOID NTAPI WSKCleanup() {
|
|
if (InterlockedCompareExchange(&g_SocketsState, INITIALIZED,
|
|
DEINITIALIZING) != INITIALIZED)
|
|
return;
|
|
|
|
WskReleaseProviderNPI(&g_WskRegistration);
|
|
WskDeregister(&g_WskRegistration);
|
|
|
|
InterlockedExchange(&g_SocketsState, DEINITIALIZED);
|
|
}
|
|
|
|
PWSK_SOCKET
|
|
NTAPI
|
|
CreateSocket(__in ADDRESS_FAMILY AddressFamily, __in USHORT SocketType,
|
|
__in ULONG Protocol, __in ULONG Flags) {
|
|
KEVENT CompletionEvent = {0};
|
|
PIRP Irp = NULL;
|
|
PWSK_SOCKET WskSocket = NULL;
|
|
NTSTATUS Status = STATUS_UNSUCCESSFUL;
|
|
|
|
if (g_SocketsState != INITIALIZED) return NULL;
|
|
|
|
Status = InitWskData(&Irp, &CompletionEvent);
|
|
if (!NT_SUCCESS(Status)) {
|
|
KdPrint(
|
|
("CreateSocket(): InitWskData() failed with status 0x%08X\n", Status));
|
|
return NULL;
|
|
}
|
|
|
|
Status = g_WskProvider.Dispatch->WskSocket(
|
|
g_WskProvider.Client, AddressFamily, SocketType, Protocol, Flags, NULL,
|
|
NULL, NULL, NULL, NULL, Irp);
|
|
if (Status == STATUS_PENDING) {
|
|
KeWaitForSingleObject(&CompletionEvent, Executive, KernelMode, FALSE, NULL);
|
|
Status = Irp->IoStatus.Status;
|
|
}
|
|
|
|
WskSocket =
|
|
NT_SUCCESS(Status) ? (PWSK_SOCKET)Irp->IoStatus.Information : NULL;
|
|
|
|
IoFreeIrp(Irp);
|
|
return (PWSK_SOCKET)WskSocket;
|
|
}
|
|
|
|
NTSTATUS
|
|
NTAPI
|
|
CloseSocket(__in PWSK_SOCKET WskSocket) {
|
|
KEVENT CompletionEvent = {0};
|
|
PIRP Irp = NULL;
|
|
NTSTATUS Status = STATUS_UNSUCCESSFUL;
|
|
|
|
if (g_SocketsState != INITIALIZED || !WskSocket)
|
|
return STATUS_INVALID_PARAMETER;
|
|
|
|
Status = InitWskData(&Irp, &CompletionEvent);
|
|
if (!NT_SUCCESS(Status)) {
|
|
KdPrint(
|
|
("CloseSocket(): InitWskData() failed with status 0x%08X\n", Status));
|
|
return Status;
|
|
}
|
|
|
|
Status = ((PWSK_PROVIDER_BASIC_DISPATCH)WskSocket->Dispatch)
|
|
->WskCloseSocket(WskSocket, Irp);
|
|
if (Status == STATUS_PENDING) {
|
|
KeWaitForSingleObject(&CompletionEvent, Executive, KernelMode, FALSE, NULL);
|
|
Status = Irp->IoStatus.Status;
|
|
}
|
|
|
|
IoFreeIrp(Irp);
|
|
return Status;
|
|
}
|
|
|
|
NTSTATUS
|
|
NTAPI
|
|
Connect(__in PWSK_SOCKET WskSocket, __in PSOCKADDR RemoteAddress) {
|
|
KEVENT CompletionEvent = {0};
|
|
PIRP Irp = NULL;
|
|
NTSTATUS Status = STATUS_UNSUCCESSFUL;
|
|
|
|
if (g_SocketsState != INITIALIZED || !WskSocket || !RemoteAddress)
|
|
return STATUS_INVALID_PARAMETER;
|
|
|
|
Status = InitWskData(&Irp, &CompletionEvent);
|
|
if (!NT_SUCCESS(Status)) {
|
|
KdPrint(("Connect(): InitWskData() failed with status 0x%08X\n", Status));
|
|
return Status;
|
|
}
|
|
|
|
Status = ((PWSK_PROVIDER_CONNECTION_DISPATCH)WskSocket->Dispatch)
|
|
->WskConnect(WskSocket, RemoteAddress, 0, Irp);
|
|
if (Status == STATUS_PENDING) {
|
|
KeWaitForSingleObject(&CompletionEvent, Executive, KernelMode, FALSE, NULL);
|
|
Status = Irp->IoStatus.Status;
|
|
}
|
|
|
|
IoFreeIrp(Irp);
|
|
return Status;
|
|
}
|
|
|
|
PWSK_SOCKET
|
|
NTAPI
|
|
SocketConnect(__in USHORT SocketType, __in ULONG Protocol,
|
|
__in PSOCKADDR RemoteAddress, __in PSOCKADDR LocalAddress) {
|
|
KEVENT CompletionEvent = {0};
|
|
PIRP Irp = NULL;
|
|
NTSTATUS Status = STATUS_UNSUCCESSFUL;
|
|
PWSK_SOCKET WskSocket = NULL;
|
|
|
|
if (g_SocketsState != INITIALIZED || !RemoteAddress || !LocalAddress)
|
|
return NULL;
|
|
|
|
Status = InitWskData(&Irp, &CompletionEvent);
|
|
if (!NT_SUCCESS(Status)) {
|
|
KdPrint(("InitWskData() failed with status 0x%08X\n", Status));
|
|
return NULL;
|
|
}
|
|
|
|
Status = g_WskProvider.Dispatch->WskSocketConnect(
|
|
g_WskProvider.Client, SocketType, Protocol, LocalAddress, RemoteAddress,
|
|
0, NULL, NULL, NULL, NULL, NULL, Irp);
|
|
if (Status == STATUS_PENDING) {
|
|
KeWaitForSingleObject(&CompletionEvent, Executive, KernelMode, FALSE, NULL);
|
|
Status = Irp->IoStatus.Status;
|
|
}
|
|
|
|
WskSocket =
|
|
NT_SUCCESS(Status) ? (PWSK_SOCKET)Irp->IoStatus.Information : NULL;
|
|
|
|
IoFreeIrp(Irp);
|
|
return WskSocket;
|
|
}
|
|
|
|
LONG NTAPI Send(__in PWSK_SOCKET WskSocket, __in PVOID Buffer,
|
|
__in ULONG BufferSize, __in ULONG Flags) {
|
|
KEVENT CompletionEvent = {0};
|
|
PIRP Irp = NULL;
|
|
WSK_BUF WskBuffer = {0};
|
|
LONG BytesSent = SOCKET_ERROR;
|
|
NTSTATUS Status = STATUS_UNSUCCESSFUL;
|
|
|
|
if (g_SocketsState != INITIALIZED || !WskSocket || !Buffer || !BufferSize)
|
|
return SOCKET_ERROR;
|
|
|
|
Status = InitWskBuffer(Buffer, BufferSize, &WskBuffer);
|
|
if (!NT_SUCCESS(Status)) {
|
|
KdPrint(("Send(): InitWskData() failed with status 0x%08X\n", Status));
|
|
return SOCKET_ERROR;
|
|
}
|
|
|
|
Status = InitWskData(&Irp, &CompletionEvent);
|
|
if (!NT_SUCCESS(Status)) {
|
|
KdPrint(("Send(): InitWskData() failed with status 0x%08X\n", Status));
|
|
FreeWskBuffer(&WskBuffer);
|
|
return SOCKET_ERROR;
|
|
}
|
|
|
|
Status = ((PWSK_PROVIDER_CONNECTION_DISPATCH)WskSocket->Dispatch)
|
|
->WskSend(WskSocket, &WskBuffer, Flags, Irp);
|
|
if (Status == STATUS_PENDING) {
|
|
KeWaitForSingleObject(&CompletionEvent, Executive, KernelMode, FALSE, NULL);
|
|
Status = Irp->IoStatus.Status;
|
|
}
|
|
|
|
BytesSent =
|
|
NT_SUCCESS(Status) ? (LONG)Irp->IoStatus.Information : SOCKET_ERROR;
|
|
|
|
IoFreeIrp(Irp);
|
|
FreeWskBuffer(&WskBuffer);
|
|
return BytesSent;
|
|
}
|
|
|
|
LONG NTAPI SendTo(__in PWSK_SOCKET WskSocket, __in PVOID Buffer,
|
|
__in ULONG BufferSize, __in_opt PSOCKADDR RemoteAddress) {
|
|
KEVENT CompletionEvent = {0};
|
|
PIRP Irp = NULL;
|
|
WSK_BUF WskBuffer = {0};
|
|
LONG BytesSent = SOCKET_ERROR;
|
|
NTSTATUS Status = STATUS_UNSUCCESSFUL;
|
|
|
|
if (g_SocketsState != INITIALIZED || !WskSocket || !Buffer || !BufferSize)
|
|
return SOCKET_ERROR;
|
|
|
|
Status = InitWskBuffer(Buffer, BufferSize, &WskBuffer);
|
|
if (!NT_SUCCESS(Status)) {
|
|
KdPrint(("SendTo(): InitWskData() failed with status 0x%08X\n", Status));
|
|
return SOCKET_ERROR;
|
|
}
|
|
|
|
Status = InitWskData(&Irp, &CompletionEvent);
|
|
if (!NT_SUCCESS(Status)) {
|
|
KdPrint(("SendTo(): InitWskData() failed with status 0x%08X\n", Status));
|
|
FreeWskBuffer(&WskBuffer);
|
|
return SOCKET_ERROR;
|
|
}
|
|
|
|
Status =
|
|
((PWSK_PROVIDER_DATAGRAM_DISPATCH)WskSocket->Dispatch)
|
|
->WskSendTo(WskSocket, &WskBuffer, 0, RemoteAddress, 0, NULL, Irp);
|
|
if (Status == STATUS_PENDING) {
|
|
KeWaitForSingleObject(&CompletionEvent, Executive, KernelMode, FALSE, NULL);
|
|
Status = Irp->IoStatus.Status;
|
|
}
|
|
|
|
BytesSent =
|
|
NT_SUCCESS(Status) ? (LONG)Irp->IoStatus.Information : SOCKET_ERROR;
|
|
|
|
IoFreeIrp(Irp);
|
|
FreeWskBuffer(&WskBuffer);
|
|
return BytesSent;
|
|
}
|
|
|
|
LONG NTAPI Receive(__in PWSK_SOCKET WskSocket, __out PVOID Buffer,
|
|
__in ULONG BufferSize, __in ULONG Flags) {
|
|
KEVENT CompletionEvent = {0};
|
|
PIRP Irp = NULL;
|
|
WSK_BUF WskBuffer = {0};
|
|
LONG BytesReceived = SOCKET_ERROR;
|
|
NTSTATUS Status = STATUS_UNSUCCESSFUL;
|
|
|
|
if (g_SocketsState != INITIALIZED || !WskSocket || !Buffer || !BufferSize)
|
|
return SOCKET_ERROR;
|
|
|
|
Status = InitWskBuffer(Buffer, BufferSize, &WskBuffer);
|
|
if (!NT_SUCCESS(Status)) {
|
|
KdPrint(("Receive(): InitWskData() failed with status 0x%08X\n", Status));
|
|
return SOCKET_ERROR;
|
|
}
|
|
|
|
Status = InitWskData(&Irp, &CompletionEvent);
|
|
if (!NT_SUCCESS(Status)) {
|
|
KdPrint(("Receive(): InitWskData() failed with status 0x%08X\n", Status));
|
|
FreeWskBuffer(&WskBuffer);
|
|
return SOCKET_ERROR;
|
|
}
|
|
|
|
Status = ((PWSK_PROVIDER_CONNECTION_DISPATCH)WskSocket->Dispatch)
|
|
->WskReceive(WskSocket, &WskBuffer, Flags, Irp);
|
|
if (Status == STATUS_PENDING) {
|
|
KeWaitForSingleObject(&CompletionEvent, Executive, KernelMode, FALSE, NULL);
|
|
Status = Irp->IoStatus.Status;
|
|
}
|
|
|
|
BytesReceived =
|
|
NT_SUCCESS(Status) ? (LONG)Irp->IoStatus.Information : SOCKET_ERROR;
|
|
|
|
IoFreeIrp(Irp);
|
|
FreeWskBuffer(&WskBuffer);
|
|
return BytesReceived;
|
|
}
|
|
|
|
LONG NTAPI ReceiveFrom(__in PWSK_SOCKET WskSocket, __out PVOID Buffer,
|
|
__in ULONG BufferSize, __out_opt PSOCKADDR RemoteAddress,
|
|
__out_opt PULONG ControlFlags) {
|
|
KEVENT CompletionEvent = {0};
|
|
PIRP Irp = NULL;
|
|
WSK_BUF WskBuffer = {0};
|
|
LONG BytesReceived = SOCKET_ERROR;
|
|
NTSTATUS Status = STATUS_UNSUCCESSFUL;
|
|
|
|
if (g_SocketsState != INITIALIZED || !WskSocket || !Buffer || !BufferSize)
|
|
return SOCKET_ERROR;
|
|
|
|
Status = InitWskBuffer(Buffer, BufferSize, &WskBuffer);
|
|
if (!NT_SUCCESS(Status)) {
|
|
KdPrint(
|
|
("ReceiveFrom(): InitWskData() failed with status 0x%08X\n", Status));
|
|
return SOCKET_ERROR;
|
|
}
|
|
|
|
Status = InitWskData(&Irp, &CompletionEvent);
|
|
if (!NT_SUCCESS(Status)) {
|
|
KdPrint(
|
|
("ReceiveFrom(): InitWskData() failed with status 0x%08X\n", Status));
|
|
FreeWskBuffer(&WskBuffer);
|
|
return SOCKET_ERROR;
|
|
}
|
|
|
|
Status = ((PWSK_PROVIDER_DATAGRAM_DISPATCH)WskSocket->Dispatch)
|
|
->WskReceiveFrom(WskSocket, &WskBuffer, 0, RemoteAddress, 0,
|
|
NULL, ControlFlags, Irp);
|
|
if (Status == STATUS_PENDING) {
|
|
KeWaitForSingleObject(&CompletionEvent, Executive, KernelMode, FALSE, NULL);
|
|
Status = Irp->IoStatus.Status;
|
|
}
|
|
|
|
BytesReceived =
|
|
NT_SUCCESS(Status) ? (LONG)Irp->IoStatus.Information : SOCKET_ERROR;
|
|
|
|
IoFreeIrp(Irp);
|
|
FreeWskBuffer(&WskBuffer);
|
|
return BytesReceived;
|
|
}
|
|
|
|
NTSTATUS
|
|
NTAPI
|
|
Bind(__in PWSK_SOCKET WskSocket, __in PSOCKADDR LocalAddress) {
|
|
KEVENT CompletionEvent = {0};
|
|
PIRP Irp = NULL;
|
|
NTSTATUS Status = STATUS_UNSUCCESSFUL;
|
|
|
|
if (g_SocketsState != INITIALIZED || !WskSocket || !LocalAddress)
|
|
return STATUS_INVALID_PARAMETER;
|
|
|
|
Status = InitWskData(&Irp, &CompletionEvent);
|
|
if (!NT_SUCCESS(Status)) {
|
|
KdPrint(("Bind(): InitWskData() failed with status 0x%08X\n", Status));
|
|
return Status;
|
|
}
|
|
|
|
Status = ((PWSK_PROVIDER_CONNECTION_DISPATCH)WskSocket->Dispatch)
|
|
->WskBind(WskSocket, LocalAddress, 0, Irp);
|
|
if (Status == STATUS_PENDING) {
|
|
KeWaitForSingleObject(&CompletionEvent, Executive, KernelMode, FALSE, NULL);
|
|
Status = Irp->IoStatus.Status;
|
|
}
|
|
|
|
IoFreeIrp(Irp);
|
|
return Status;
|
|
}
|
|
|
|
PWSK_SOCKET
|
|
NTAPI
|
|
Accept(__in PWSK_SOCKET WskSocket, __out_opt PSOCKADDR LocalAddress,
|
|
__out_opt PSOCKADDR RemoteAddress) {
|
|
KEVENT CompletionEvent = {0};
|
|
PIRP Irp = NULL;
|
|
NTSTATUS Status = STATUS_UNSUCCESSFUL;
|
|
PWSK_SOCKET AcceptedSocket = NULL;
|
|
|
|
if (g_SocketsState != INITIALIZED || !WskSocket) return NULL;
|
|
|
|
Status = InitWskData(&Irp, &CompletionEvent);
|
|
if (!NT_SUCCESS(Status)) {
|
|
KdPrint(("Accept(): InitWskData() failed with status 0x%08X\n", Status));
|
|
return NULL;
|
|
}
|
|
|
|
Status = ((PWSK_PROVIDER_LISTEN_DISPATCH)WskSocket->Dispatch)
|
|
->WskAccept(WskSocket, 0, NULL, NULL, LocalAddress,
|
|
RemoteAddress, Irp);
|
|
if (Status == STATUS_PENDING) {
|
|
KeWaitForSingleObject(&CompletionEvent, Executive, KernelMode, FALSE, NULL);
|
|
Status = Irp->IoStatus.Status;
|
|
}
|
|
|
|
AcceptedSocket =
|
|
NT_SUCCESS(Status) ? (PWSK_SOCKET)Irp->IoStatus.Information : NULL;
|
|
|
|
IoFreeIrp(Irp);
|
|
return AcceptedSocket;
|
|
}
|