From c9d0914e6484ac5ec3467e9e0a3d2750e7c88201 Mon Sep 17 00:00:00 2001 From: Martin Drab Date: Wed, 12 Oct 2022 04:13:13 -0700 Subject: [PATCH 01/18] WSK: Fixes related to the repository and MSVS requirements - WSK project added to the main solution, - removed Debug configurations (to be in sync with other projects) Signed-off-by: Martin Drab --- buildAll.bat | 2 ++ viosock/buildAll.bat | 2 ++ viosock/cleanAll.bat | 4 ++++ viosock/wsk/LICENSE | 28 ++++++++++++++++++++++++++++ viosock/wsk/wsk.vcxproj | 12 ------------ 5 files changed, 36 insertions(+), 12 deletions(-) create mode 100644 viosock/wsk/LICENSE diff --git a/buildAll.bat b/buildAll.bat index 50a1c143e..dd18e79ae 100755 --- a/buildAll.bat +++ b/buildAll.bat @@ -26,6 +26,8 @@ call tools\build.bat vioserial\sys\vioser.vcxproj "Win11_SDV" %* if errorlevel 1 goto :fail call tools\build.bat viosock\sys\viosock.vcxproj "Win11_SDV" %* if errorlevel 1 goto :fail +call tools\build.bat viosock\wsk\wsk.vcxproj "Win11_SDV" %* +if errorlevel 1 goto :fail call tools\build.bat viofs\pci\viofs.vcxproj "Win11_SDV" %* if errorlevel 1 goto :fail call tools\build.bat vioinput\hidpassthrough\hidpassthrough.vcxproj "Win11_SDV" %* diff --git a/viosock/buildAll.bat b/viosock/buildAll.bat index a250f4999..26581edd2 100644 --- a/viosock/buildAll.bat +++ b/viosock/buildAll.bat @@ -4,3 +4,5 @@ if errorlevel 1 goto :eof call ..\tools\build.bat viosock.sln "Win10 Win11" %* if errorlevel 1 goto :eof call ..\tools\build.bat sys\viosock.vcxproj "Win11_SDV" %* +if errorlevel 1 goto :eof +call ..\tools\build.bat wsk\wsk.vcxproj "Win11_SDV" %* diff --git a/viosock/cleanAll.bat b/viosock/cleanAll.bat index 7ce4b642c..47f7edb45 100644 --- a/viosock/cleanAll.bat +++ b/viosock/cleanAll.bat @@ -13,6 +13,10 @@ pushd sys call cleanAll.bat popd +pushd wsk +call cleanAll.bat +popd + pushd viosock-test call cleanAll.bat popd diff --git a/viosock/wsk/LICENSE b/viosock/wsk/LICENSE new file mode 100644 index 000000000..95c4900e4 --- /dev/null +++ b/viosock/wsk/LICENSE @@ -0,0 +1,28 @@ +Copyright 2019 Virtuozzo, Inc. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions +are met: + +Redistributions of source code must retain the above copyright +notice, this list of conditions and the following disclaimer. + +Redistributions in binary form must reproduce the above copyright +notice, this list of conditions and the following disclaimer in the +documentation and/or other materials provided with the distribution. + +Neither the name of the copyright holder nor the names of its +contributors may be used to endorse or promote products derived from +this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/viosock/wsk/wsk.vcxproj b/viosock/wsk/wsk.vcxproj index a89026973..909a388c7 100644 --- a/viosock/wsk/wsk.vcxproj +++ b/viosock/wsk/wsk.vcxproj @@ -1,18 +1,6 @@ - - Win10 Debug - ARM64 - - - Win10 Debug - Win32 - - - Win10 Debug - x64 - Win10 Release Win32 From 2e6e47a0873c5855bc077df4d3baa1d5a782f1b1 Mon Sep 17 00:00:00 2001 From: Martin Drab Date: Sun, 24 Jul 2022 16:16:36 +0200 Subject: [PATCH 02/18] WSK: Add more convenient debug print macros These macros will be used mostly when entering into or exiting from a function, sending routine name nad values of its arguments to a potential kernel debugger. These macros are NOPs on Release build, however, their usage can be invaluable when debugging driver crashes. Signed-off-by: Martin Drab --- viosock/inc/debug-utils.h | 83 +++++++++++++++++++++++++++++++++ viosock/wsk/precomp.h | 2 + viosock/wsk/wsk.vcxproj | 1 + viosock/wsk/wsk.vcxproj.filters | 3 ++ 4 files changed, 89 insertions(+) create mode 100644 viosock/inc/debug-utils.h diff --git a/viosock/inc/debug-utils.h b/viosock/inc/debug-utils.h new file mode 100644 index 000000000..ce8c319c1 --- /dev/null +++ b/viosock/inc/debug-utils.h @@ -0,0 +1,83 @@ +/* + * Exports definition for virtio socket WSK interface + * + * Copyright (c) 2021 Virtuozzo International GmbH + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions + * are met : + * 1. Redistributions of source code must retain the above copyright + * notice, this list of conditions and the following disclaimer. + * 2. Redistributions in binary form must reproduce the above copyright + * notice, this list of conditions and the following disclaimer in the + * documentation and / or other materials provided with the distribution. + * 3. Neither the names of the copyright holders nor the names of their contributors + * may be used to endorse or promote products derived from this software + * without specific prior written permission. + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS ``AS IS'' AND + * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + * ARE DISCLAIMED.IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS + * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) + * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT + * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY + * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF + * SUCH DAMAGE. + */ + +#ifndef __DEBUG_UTILS_H__ +#define __DEBUG_UTILS_H__ + + + +#if defined(DBG) || defined(_DEBUG) + +#define DEBUG_ENTER_FUNCTION(Format, ...) \ + DbgPrintEx(DPFLTR_DEFAULT_ID, DPFLTR_TRACE_LEVEL, "[0x%x:0x%x:%u]: " __FUNCTION__ "(" Format ")\n", \ + PtrToUlong(PsGetCurrentProcessId()), PtrToUlong(PsGetCurrentThreadId()), KeGetCurrentIrql(), __VA_ARGS__) + +#define DEBUG_ENTER_FUNCTION_NO_ARGS() \ + DbgPrintEx(DPFLTR_DEFAULT_ID, DPFLTR_TRACE_LEVEL, "[0x%x:0x%x:%u]: " __FUNCTION__ "()\n", \ + PtrToUlong(PsGetCurrentProcessId()), PtrToUlong(PsGetCurrentThreadId()), KeGetCurrentIrql()) + +#define DEBUG_EXIT_FUNCTION(Format, ...) \ + DbgPrintEx(DPFLTR_DEFAULT_ID, DPFLTR_TRACE_LEVEL, "[0x%x:0x%x:%u]: " __FUNCTION__ "(-):" Format "\n", \ + PtrToUlong(PsGetCurrentProcessId()), PtrToUlong(PsGetCurrentThreadId()), KeGetCurrentIrql(), __VA_ARGS__) + +#define DEBUG_EXIT_FUNCTION_VOID() \ + DbgPrintEx(DPFLTR_DEFAULT_ID, DPFLTR_TRACE_LEVEL, "[0x%x:0x%x:%u]: " __FUNCTION__ "(-)\n", \ + PtrToUlong(PsGetCurrentProcessId()), PtrToUlong(PsGetCurrentThreadId()), KeGetCurrentIrql()) + +#else + +#define DEBUG_ENTER_FUNCTION(Format, ...) do { } while (FALSE) + +#define DEBUG_ENTER_FUNCTION_NO_ARGS() do { } while (FALSE) + +#define DEBUG_EXIT_FUNCTION(Format, ...) do { } while (FALSE) + +#define DEBUG_EXIT_FUNCTION_VOID() do { } while (FALSE) + +#endif + + +#define DEBUG_ERROR(Format, ...) \ + DbgPrintEx(DPFLTR_DEFAULT_ID, DPFLTR_ERROR_LEVEL, "[0x%x:0x%x:%u]: " __FUNCTION__ " (" __FILE__ ":%u): ERROR: " Format "\n", \ + PtrToUlong(PsGetCurrentProcessId()), PtrToUlong(PsGetCurrentThreadId()), KeGetCurrentIrql(), __LINE__, __VA_ARGS__) + +#define DEBUG_WARNING(Format, ...) \ + DbgPrintEx(DPFLTR_DEFAULT_ID, DPFLTR_WARNING_LEVEL, "[0x%x:0x%x:%u]: " __FUNCTION__ " (" __FILE__ ":%u): WARNING: " Format "\n", \ + PtrToUlong(PsGetCurrentProcessId()), PtrToUlong(PsGetCurrentThreadId()), KeGetCurrentIrql(), __LINE__, __VA_ARGS__) + +#define DEBUG_TRACE(Format, ...) \ + DbgPrintEx(DPFLTR_DEFAULT_ID, DPFLTR_TRACE_LEVEL, "[0x%x:0x%x:%u]: " __FUNCTION__ " (" __FILE__ ":%u): TRACE: " Format "\n", \ + PtrToUlong(PsGetCurrentProcessId()), PtrToUlong(PsGetCurrentThreadId()), KeGetCurrentIrql(), __LINE__, __VA_ARGS__) + +#define DEBUG_INFO(Format, ...) \ + DbgPrintEx(DPFLTR_DEFAULT_ID, DPFLTR_INFO_LEVEL, "[0x%x:0x%x:%u]: " __FUNCTION__ ": INFO: " Format "\n", \ + PtrToUlong(PsGetCurrentProcessId()), PtrToUlong(PsGetCurrentThreadId()), KeGetCurrentIrql(), __VA_ARGS__) + + +#endif diff --git a/viosock/wsk/precomp.h b/viosock/wsk/precomp.h index fefa6b7c0..19a740e66 100644 --- a/viosock/wsk/precomp.h +++ b/viosock/wsk/precomp.h @@ -26,6 +26,7 @@ * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF * SUCH DAMAGE. */ + #include #include #include @@ -39,5 +40,6 @@ #include #include +#include "..\inc\debug-utils.h" #include "..\sys\public.h" #include "..\inc\vio_sockets.h" diff --git a/viosock/wsk/wsk.vcxproj b/viosock/wsk/wsk.vcxproj index 909a388c7..8adf75063 100644 --- a/viosock/wsk/wsk.vcxproj +++ b/viosock/wsk/wsk.vcxproj @@ -27,6 +27,7 @@ + diff --git a/viosock/wsk/wsk.vcxproj.filters b/viosock/wsk/wsk.vcxproj.filters index b214ee26a..39d2d539c 100644 --- a/viosock/wsk/wsk.vcxproj.filters +++ b/viosock/wsk/wsk.vcxproj.filters @@ -21,6 +21,9 @@ Header Files + + Header Files + From 3dcadeb5a00b6dcef9189281f3a221632f315cdb Mon Sep 17 00:00:00 2001 From: Martin Drab Date: Sun, 24 Jul 2022 17:02:11 +0200 Subject: [PATCH 03/18] WSK: Module initialization, Wsk(De)Register, Wsk(Capture|Release)ProviderNPI Unlike the Windows WSK implementation, this library needs to be initialized through VioWskModuleInit before calling its WsRegister equivalent. This approach allows the client driver to provide the library with an optional device object that is then used to prevent IRP completion routines and workitems used by the library to execute after the device deletion. If the client application decides not to provide such a device object, the library can live with it but no unload protection is then enforced. Also, the client driver must provide pointer to its DRIVER_OBJECT which is used during WskCaptureProviderNPI to connect to the VIOSOCK driver. Instead of having the VIOSOCK device name hard-coded, the library takes advantage of IoRegisterPlugPlayNotification to search for the GUID_DEVINTERFACE_VIOSOCK device interface (which is enabled by the VIOSOCK driver) dynamically. Signed-off-by: Martin Drab --- viosock/inc/vio_wsk.h | 16 ++ viosock/wsk/viowsk-internal.h | 12 ++ viosock/wsk/viowsk.c | 297 ++++++++++++++++++++++++++++---- viosock/wsk/viowsk.h | 5 + viosock/wsk/wsk.vcxproj | 1 + viosock/wsk/wsk.vcxproj.filters | 3 + 6 files changed, 304 insertions(+), 30 deletions(-) create mode 100644 viosock/wsk/viowsk-internal.h diff --git a/viosock/inc/vio_wsk.h b/viosock/inc/vio_wsk.h index ab2c4128b..ef6044e3b 100644 --- a/viosock/inc/vio_wsk.h +++ b/viosock/inc/vio_wsk.h @@ -70,6 +70,22 @@ VioWskQueryProviderCharacteristics( _Out_ PWSK_PROVIDER_CHARACTERISTICS WskProviderCharacteristics ); +_Must_inspect_result_ +NTSTATUS +VioWskModuleInit( + _In_ PDRIVER_OBJECT DriverObject, + _In_ PUNICODE_STRING RegistryPath, + _In_opt_ PDEVICE_OBJECT DeviceObject +); + +VOID +VioWskModuleFinit( + VOID +); + + + + #ifdef __cplusplus } #endif diff --git a/viosock/wsk/viowsk-internal.h b/viosock/wsk/viowsk-internal.h new file mode 100644 index 000000000..ec90c1aa2 --- /dev/null +++ b/viosock/wsk/viowsk-internal.h @@ -0,0 +1,12 @@ + +#ifndef __VIOWSK_INTERNAL_H__ +#define __VIOWSK_INTERNAL_H__ + + + +extern PDRIVER_OBJECT _viowskDriverObject; +extern PDEVICE_OBJECT _viowskDeviceObject; + + + +#endif diff --git a/viosock/wsk/viowsk.c b/viosock/wsk/viowsk.c index 55adfd37c..3eca3a086 100644 --- a/viosock/wsk/viowsk.c +++ b/viosock/wsk/viowsk.c @@ -28,6 +28,8 @@ */ #include "precomp.h" +#include "..\inc\debug-utils.h" +#include "viowsk-internal.h" #include "viowsk.h" #include "..\inc\vio_wsk.h" @@ -36,6 +38,164 @@ #pragma alloc_text (PAGE, VioWskDeregister) #endif + +typedef struct _VIOSOCK_WAIT_CONTEXT { + UNICODE_STRING SymbolicLinkName; + KEVENT Event; + NTSTATUS Result; +} VIOSOCK_WAIT_CONTEXT, * PVIOSOCK_WAIT_CONTEXT; + + +PDRIVER_OBJECT _viowskDriverObject = NULL; +PDEVICE_OBJECT _viowskDeviceObject = NULL; + +static const WSK_PROVIDER_DISPATCH _providerDispatch = { + MAKE_WSK_VERSION(VIOWSK_PROVIDER_VERSION, 0), + 0, + VioWskSocket, + VioWskSocketConnect, + VioWskControlClient, + VioWskGetAddressInfo, + VioWskFreeAddressInfo, + VioWskGetNameInfo, +}; + + + +static +NTSTATUS +_NotifyCallback( + _In_ PVOID NotificationStructure, + _Inout_ PVOID Context +) +{ + NTSTATUS Status = STATUS_UNSUCCESSFUL; + PVIOSOCK_WAIT_CONTEXT Ctx = (PVIOSOCK_WAIT_CONTEXT)Context; + PDEVICE_INTERFACE_CHANGE_NOTIFICATION NotifyInfo = (PDEVICE_INTERFACE_CHANGE_NOTIFICATION)NotificationStructure; + DEBUG_ENTER_FUNCTION("NotificationStructure=0x%p; Context=0x%p", NotificationStructure, Context); + + Status = STATUS_SUCCESS; + if (IsEqualGUID(&NotifyInfo->Event, &GUID_DEVICE_INTERFACE_ARRIVAL) && + NotifyInfo->SymbolicLinkName != NULL) + { + Ctx->SymbolicLinkName = *NotifyInfo->SymbolicLinkName; + Ctx->SymbolicLinkName.MaximumLength = Ctx->SymbolicLinkName.Length; + Ctx->SymbolicLinkName.Buffer = ExAllocatePoolUninitialized(PagedPool, Ctx->SymbolicLinkName.MaximumLength, VIOSOCK_WSK_MEMORY_TAG); + if (Ctx->SymbolicLinkName.Buffer == NULL) + Status = STATUS_INSUFFICIENT_RESOURCES; + + if (NT_SUCCESS(Status)) + { + RtlCopyMemory(Ctx->SymbolicLinkName.Buffer, NotifyInfo->SymbolicLinkName->Buffer, Ctx->SymbolicLinkName.Length); + } + + Ctx->Result = Status; + KeSetEvent(&Ctx->Event, IO_NO_INCREMENT, FALSE); + Status = STATUS_SUCCESS; + } + + DEBUG_EXIT_FUNCTION("0x%x", Status); + return Status; +} + +_Must_inspect_result_ +static +NTSTATUS +_VioSockDriverConnect( + _Inout_ PVIOWSK_REG_CONTEXT WskContext, + _In_ DWORD Timeout, + _In_ PDRIVER_OBJECT DriverObject) +{ + PVOID NotifyHandle = NULL; + LARGE_INTEGER Timeout100ns; + PLARGE_INTEGER pTimeout = NULL; + VIOSOCK_WAIT_CONTEXT WaitContext; + NTSTATUS Status = STATUS_UNSUCCESSFUL; + DEBUG_ENTER_FUNCTION("WskContext=0x%p; Timeout=%u; DriverObject=0x%p", WskContext, Timeout, DriverObject); + + switch (Timeout) { + case WSK_NO_WAIT: + Timeout100ns.QuadPart = 0; + pTimeout = &Timeout100ns; + break; + case WSK_INFINITE_WAIT: + break; + default: + Timeout100ns.QuadPart = Timeout; + Timeout100ns.QuadPart *= -10000; + pTimeout = &Timeout100ns; + break; + } + + KeInitializeEvent(&WaitContext.Event, NotificationEvent, FALSE); + WaitContext.Result = STATUS_UNSUCCESSFUL; + Status = IoRegisterPlugPlayNotification( + EventCategoryDeviceInterfaceChange, + PNPNOTIFY_DEVICE_INTERFACE_INCLUDE_EXISTING_INTERFACES, + (PVOID)&GUID_DEVINTERFACE_VIOSOCK, + DriverObject, + _NotifyCallback, + &WaitContext, + &NotifyHandle); + + if (!NT_SUCCESS(Status)) + { + DEBUG_ERROR("IoRegisterPlugPlayNotification: 0x%x", Status); + goto Exit; + } + + Status = KeWaitForSingleObject(&WaitContext.Event, Executive, KernelMode, FALSE, pTimeout); +#if (NTDDI_VERSION >= NTDDI_WIN7) + IoUnregisterPlugPlayNotificationEx(NotifyHandle); +#else // if (NTDDI_VERSION >= NTDDI_WIN7) + IoUnregisterPlugPlayNotification(NotifyHandle); +#endif // if (NTDDI_VERSION < NTDDI_WIN7) + switch (Status) { + case STATUS_WAIT_0: + Status = WaitContext.Result; + if (!NT_SUCCESS(Status)) + { + DEBUG_ERROR("The interface arrival notificaiton routine failed: 0x%x", Status); + break; + } + + Status = IoGetDeviceObjectPointer(&WaitContext.SymbolicLinkName, FILE_READ_ATTRIBUTES, &WskContext->VIOSockMainFileObject, &WskContext->VIOSockDevice); + if (!NT_SUCCESS(Status)) + { + DEBUG_ERROR("IoGetDeviceObjectPointer: 0x%x", Status); + ExFreePoolWithTag(WaitContext.SymbolicLinkName.Buffer, VIOSOCK_WSK_MEMORY_TAG); + break; + } + + WskContext->VIOSockLinkName = WaitContext.SymbolicLinkName; + break; + case STATUS_TIMEOUT: + Status = STATUS_DEVICE_NOT_READY; + break; + } + +Exit: + DEBUG_EXIT_FUNCTION("0x%x", Status); + return Status; +} + +static +VOID +_VIOSockDriverDisconnect( + _Inout_ PVIOWSK_REG_CONTEXT RegContext) +{ + DEBUG_ENTER_FUNCTION("RegContext=0x%p", RegContext); + + ObDereferenceObject(RegContext->VIOSockMainFileObject); + RegContext->VIOSockMainFileObject = NULL; + ExFreePoolWithTag(RegContext->VIOSockLinkName.Buffer, VIOSOCK_WSK_MEMORY_TAG); + + DEBUG_EXIT_FUNCTION_VOID(); + return; +} + + + _Must_inspect_result_ NTSTATUS VioWskRegister( @@ -44,19 +204,24 @@ VioWskRegister( ) { PVIOWSK_REG_CONTEXT pContext; + NTSTATUS Status = STATUS_UNSUCCESSFUL; + DEBUG_ENTER_FUNCTION("WskClientNpi=0x%p; WskRegistration=0x%p", WskClientNpi, WskRegistration); PAGED_CODE(); - if (!WskClientNpi || !WskRegistration) - return STATUS_INVALID_PARAMETER; - pContext = ExAllocatePoolUninitialized(NonPagedPoolNx, sizeof(VIOWSK_REG_CONTEXT), VIOSOCK_WSK_MEMORY_TAG); if (!pContext) - return STATUS_INSUFFICIENT_RESOURCES; + { + Status = STATUS_INSUFFICIENT_RESOURCES; + goto Exit; + } RtlZeroMemory(pContext, sizeof(VIOWSK_REG_CONTEXT)); + Status = ExInitializeResourceLite(&pContext->NPILock); + if (!NT_SUCCESS(Status)) + goto FreeContext; pContext->ClientContext = WskClientNpi->ClientContext; if (WskClientNpi->Dispatch) @@ -64,7 +229,12 @@ VioWskRegister( WskRegistration->ReservedRegistrationContext = pContext; - return STATUS_SUCCESS; +FreeContext: + if (!NT_SUCCESS(Status)) + ExFreePoolWithTag(pContext, VIOSOCK_WSK_MEMORY_TAG); +Exit: + DEBUG_EXIT_FUNCTION("0x%x", Status); + return Status; } VOID @@ -72,14 +242,19 @@ VioWskDeregister( _In_ PWSK_REGISTRATION WskRegistration ) { - if (!WskRegistration) - return; + PVIOWSK_REG_CONTEXT pContext = NULL; + DEBUG_ENTER_FUNCTION("WskRegistration=0x%p", WskRegistration); - if (WskRegistration->ReservedRegistrationContext) + pContext = WskRegistration->ReservedRegistrationContext; + if (pContext != NULL) { - ExFreePool(WskRegistration->ReservedRegistrationContext); + ExDeleteResourceLite(&pContext->NPILock); + ExFreePoolWithTag(pContext, VIOSOCK_WSK_MEMORY_TAG); WskRegistration->ReservedRegistrationContext = NULL; } + + DEBUG_EXIT_FUNCTION_VOID(); + return; } _Must_inspect_result_ @@ -90,29 +265,37 @@ VioWskCaptureProviderNPI( _Out_ PWSK_PROVIDER_NPI WskProviderNpi ) { - WSK_PROVIDER_DISPATCH *Dispatch; - - UNREFERENCED_PARAMETER(WaitTimeout); + NTSTATUS Status = STATUS_UNSUCCESSFUL; + PVIOWSK_REG_CONTEXT regContext = (PVIOWSK_REG_CONTEXT)WskRegistration->ReservedRegistrationContext; + DEBUG_ENTER_FUNCTION("WskRegistration=0x%p; WaitTimeout=%u; WskProviderNpi=0x%p", WskRegistration, WaitTimeout, WskProviderNpi); - if (!WskProviderNpi || !WskProviderNpi->Dispatch || !WskRegistration) - return STATUS_INVALID_PARAMETER; - - Dispatch = (WSK_PROVIDER_DISPATCH *)WskProviderNpi->Dispatch; - - //TODO: open viosock and handle WaitTimeout - WskProviderNpi->Client = (PWSK_CLIENT)WskRegistration; + if (KeGetCurrentIrql() >= APC_LEVEL) + { + Status = STATUS_INVALID_PARAMETER; + goto Exit; + } - Dispatch->Version = MAKE_WSK_VERSION(VIOWSK_PROVIDER_VERSION, 0); - Dispatch->Reserved = 0; + Status = STATUS_SUCCESS; + KeEnterCriticalRegion(); + ExAcquireResourceExclusiveLite(®Context->NPILock, TRUE); + if (InterlockedIncrement(®Context->NPICount) == 1) + { + Status = _VioSockDriverConnect(regContext, WaitTimeout, _viowskDriverObject); + if (!NT_SUCCESS(Status)) + InterlockedDecrement(®Context->NPICount); + } - Dispatch->WskSocket = VioWskSocket; - Dispatch->WskSocketConnect = VioWskSocketConnect; - Dispatch->WskControlClient = VioWskControlClient; - Dispatch->WskGetAddressInfo = VioWskGetAddressInfo; - Dispatch->WskFreeAddressInfo = VioWskFreeAddressInfo; - Dispatch->WskGetNameInfo = VioWskGetNameInfo; + ExReleaseResourceLite(®Context->NPILock); + KeLeaveCriticalRegion(); + if (NT_SUCCESS(Status)) + { + WskProviderNpi->Client = (PWSK_CLIENT)WskRegistration; + WskProviderNpi->Dispatch = &_providerDispatch; + } - return STATUS_SUCCESS; +Exit: + DEBUG_EXIT_FUNCTION("0x%x", Status); + return Status; } VOID @@ -120,8 +303,19 @@ VioWskReleaseProviderNPI( _In_ PWSK_REGISTRATION WskRegistration ) { - UNREFERENCED_PARAMETER(WskRegistration); - //TODO: close viosock + PVIOWSK_REG_CONTEXT regContext = (PVIOWSK_REG_CONTEXT)WskRegistration->ReservedRegistrationContext; + DEBUG_ENTER_FUNCTION("WskRegistration=0x%p", WskRegistration); + + KeEnterCriticalRegion(); + ExAcquireResourceExclusiveLite(®Context->NPILock, TRUE); + if (InterlockedDecrement(®Context->NPICount) == 0) + _VIOSockDriverDisconnect(regContext); + + ExReleaseResourceLite(®Context->NPILock); + KeLeaveCriticalRegion(); + + DEBUG_EXIT_FUNCTION_VOID(); + return; } _Must_inspect_result_ @@ -140,3 +334,46 @@ VioWskQueryProviderCharacteristics( return STATUS_SUCCESS; } +_Must_inspect_result_ +NTSTATUS +VioWskModuleInit( + _In_ PDRIVER_OBJECT DriverObject, + _In_ PUNICODE_STRING RegistryPath, + _In_opt_ PDEVICE_OBJECT DeviceObject +) +{ + NTSTATUS Status = STATUS_UNSUCCESSFUL; + DEBUG_ENTER_FUNCTION("DriverObject=0x%p; RegistryPath=\"%wZ\"", DriverObject, RegistryPath); + + UNREFERENCED_PARAMETER(RegistryPath); + + Status = STATUS_SUCCESS; + ObReferenceObject(DriverObject); + _viowskDriverObject = DriverObject; + if (DeviceObject) { + ObReferenceObject(DeviceObject); + _viowskDeviceObject = DeviceObject; + } + + DEBUG_EXIT_FUNCTION("0x%x", Status); + return Status; +} + +VOID +VioWskModuleFinit( + VOID +) +{ + DEBUG_ENTER_FUNCTION_NO_ARGS(); + + if (_viowskDeviceObject) { + ObDereferenceObject(_viowskDeviceObject); + _viowskDeviceObject = NULL; + } + + ObDereferenceObject(_viowskDriverObject); + _viowskDriverObject = NULL; + + DEBUG_EXIT_FUNCTION_VOID(); + return; +} diff --git a/viosock/wsk/viowsk.h b/viosock/wsk/viowsk.h index e6e7ba7c5..86e243530 100644 --- a/viosock/wsk/viowsk.h +++ b/viosock/wsk/viowsk.h @@ -40,6 +40,11 @@ typedef struct _VIOWSK_REG_CONTEXT { PVOID ClientContext; WSK_CLIENT_DISPATCH ClientDispatch; + volatile LONG NPICount; + ERESOURCE NPILock; + UNICODE_STRING VIOSockLinkName; + PDEVICE_OBJECT VIOSockDevice; + PFILE_OBJECT VIOSockMainFileObject; }VIOWSK_REG_CONTEXT,*PVIOWSK_REG_CONTEXT; #define VIOWSK_PROVIDER_VERSION 1 diff --git a/viosock/wsk/wsk.vcxproj b/viosock/wsk/wsk.vcxproj index 8adf75063..2892759cd 100644 --- a/viosock/wsk/wsk.vcxproj +++ b/viosock/wsk/wsk.vcxproj @@ -29,6 +29,7 @@ + diff --git a/viosock/wsk/wsk.vcxproj.filters b/viosock/wsk/wsk.vcxproj.filters index 39d2d539c..a865e460c 100644 --- a/viosock/wsk/wsk.vcxproj.filters +++ b/viosock/wsk/wsk.vcxproj.filters @@ -24,6 +24,9 @@ Header Files + + Header Files + From 2b8c9370edbfa5a5d6b4c132d5d10b79174447f9 Mon Sep 17 00:00:00 2001 From: Martin Drab Date: Sun, 24 Jul 2022 17:50:21 +0200 Subject: [PATCH 04/18] WSK: Address translation routines Implements WskGetAddressInfo, WskFreeAddressInfo and WskGetNameInfo which can be used to translate AF_VSOCK addresses between their string and binary forms. Also, some necessary logic for IRP processing is being added as well. This is necessary because the WSK interface uses IRPs to signal the client driver that certain WSK operation was completed. Apart from that (and from providing operation results in their `IoStatus` field), IRPs are not used by the interface (they carry neither operation information, nor data associated with it). It is worth to note that the whole concept of address translation is not that useful for AF_VSOCK since there are no such things as DNS servers and domain names. ref: [WskGetAddressInfo](https://learn.microsoft.com/en-us/windows-hardware/drivers/ddi/wsk/nc-wsk-pfn_wsk_get_address_info) ref: [WskFreeAddressInfo](https://learn.microsoft.com/en-us/windows-hardware/drivers/ddi/wsk/nc-wsk-pfn_wsk_free_address_info) ref: [WskGetNameInfo](https://learn.microsoft.com/en-us/windows-hardware/drivers/ddi/wsk/nc-wsk-pfn_wsk_get_name_info) Signed-off-by: Martin Drab --- viosock/wsk/provider.c | 132 ++++++++++++++++++++++--- viosock/wsk/viowsk.h | 2 + viosock/wsk/wsk-utils.c | 167 ++++++++++++++++++++++++++++++++ viosock/wsk/wsk-utils.h | 78 +++++++++++++++ viosock/wsk/wsk.vcxproj | 2 + viosock/wsk/wsk.vcxproj.filters | 6 ++ 6 files changed, 373 insertions(+), 14 deletions(-) create mode 100644 viosock/wsk/wsk-utils.c create mode 100644 viosock/wsk/wsk-utils.h diff --git a/viosock/wsk/provider.c b/viosock/wsk/provider.c index 28e4b1fc7..b2bcc26f1 100644 --- a/viosock/wsk/provider.c +++ b/viosock/wsk/provider.c @@ -29,9 +29,14 @@ #include "precomp.h" #include "viowsk.h" +#include "..\inc\debug-utils.h" +#include "wsk-utils.h" #include "..\inc\vio_wsk.h" #ifdef ALLOC_PRAGMA +#pragma alloc_text (PAGE, VioWskGetAddressInfo) +#pragma alloc_text (PAGE, VioWskFreeAddressInfo) +#pragma alloc_text (PAGE, VioWskGetNameInfo) #endif static @@ -233,6 +238,7 @@ VioWskControlClient( return VioWskCompleteIrp(Irp, STATUS_NOT_IMPLEMENTED, 0); } +_At_(*Result, __drv_allocatesMem(Mem)) NTSTATUS WSKAPI VioWskGetAddressInfo( @@ -242,25 +248,82 @@ VioWskGetAddressInfo( _In_opt_ ULONG NameSpace, _In_opt_ GUID *Provider, _In_opt_ PADDRINFOEXW Hints, - _Outptr_ PADDRINFOEXW *Result, + _Outptr_ PADDRINFOEXW *Result, _In_opt_ PEPROCESS OwningProcess, _In_opt_ PETHREAD OwningThread, _Inout_ PIRP Irp - ) +) { + ULONG Cid = 0; + ULONG Port = 0; + ULONG_PTR SizeReturned = 0; + PADDRINFOEXW AddrInfo = NULL; + PSOCKADDR_VM VMAddr = NULL; + NTSTATUS Status = STATUS_UNSUCCESSFUL; + DEBUG_ENTER_FUNCTION("Client=0x%p; NodeName=\"%wZ\"; ServiceName=\"%wZ\"; NameSpace=%u; Provider=0x%p; Hints=0x%p; Result=0x%p; OwningProcess=0x%p; OwningThread=0x%p; Irp=0x%p", Client, NodeName, ServiceName, NameSpace, Provider, Hints, Result, OwningProcess, OwningThread, Irp); + + PAGED_CODE(); UNREFERENCED_PARAMETER(Client); - UNREFERENCED_PARAMETER(NodeName); - UNREFERENCED_PARAMETER(ServiceName); UNREFERENCED_PARAMETER(NameSpace); UNREFERENCED_PARAMETER(Provider); UNREFERENCED_PARAMETER(Hints); - UNREFERENCED_PARAMETER(Result); UNREFERENCED_PARAMETER(OwningProcess); UNREFERENCED_PARAMETER(OwningThread); - return VioWskCompleteIrp(Irp, STATUS_NOT_IMPLEMENTED, 0); + Status = STATUS_SUCCESS; + if (NodeName) + { + Status = VioWskStringToAddressPart(NodeName, &Cid); + if (!NT_SUCCESS(Status)) + goto CompleteIrp; + } + + if (ServiceName) + { + Status = VioWskStringToAddressPart(ServiceName, &Port); + if (!NT_SUCCESS(Status)) + goto CompleteIrp; + } + + VMAddr = ExAllocatePoolUninitialized(NonPagedPool, sizeof(SOCKADDR_VM), VIOSOCK_WSK_MEMORY_TAG); + if (!VMAddr) + { + Status = STATUS_INSUFFICIENT_RESOURCES; + goto CompleteIrp; + } + + memset(VMAddr, 0, sizeof(SOCKADDR_VM)); + VMAddr->svm_family = AF_VSOCK; + VMAddr->svm_cid = Cid; + VMAddr->svm_port = Port; + AddrInfo = ExAllocatePoolUninitialized(PagedPool, sizeof(ADDRINFOEXW), VIOSOCK_WSK_MEMORY_TAG); + if (!AddrInfo) + { + Status = STATUS_INSUFFICIENT_RESOURCES; + goto FreeVMAddr; + } + + memset(AddrInfo, 0, sizeof(ADDRINFOEXW)); + AddrInfo->ai_family = AF_VSOCK; + AddrInfo->ai_socktype = SOCK_STREAM; + AddrInfo->ai_addr = (struct sockaddr*)VMAddr; + AddrInfo->ai_addrlen = sizeof(SOCKADDR_VM); + *Result = AddrInfo; + SizeReturned = sizeof(ADDRINFOEXW); + VMAddr = NULL; + +FreeVMAddr: + if (VMAddr) + ExFreePoolWithTag(VMAddr, VIOSOCK_WSK_MEMORY_TAG); +CompleteIrp: + if (Irp) + VioWskIrpComplete(NULL, Irp, Status, SizeReturned); + + DEBUG_EXIT_FUNCTION("0x%x, *Result=0x%p", Status, *Result); + return Status; } +_At_(AddrInfo, __drv_freesMem(Mem)) VOID WSKAPI VioWskFreeAddressInfo( @@ -268,9 +331,23 @@ VioWskFreeAddressInfo( _In_ PADDRINFOEXW AddrInfo ) { + + PADDRINFOEXW Prev = NULL; + DEBUG_ENTER_FUNCTION("Client=0x%p; AddrInfo=0x%p", Client, AddrInfo); + + PAGED_CODE(); UNREFERENCED_PARAMETER(Client); - UNREFERENCED_PARAMETER(AddrInfo); + while (AddrInfo != NULL) + { + Prev = AddrInfo; + AddrInfo = AddrInfo->ai_next; + ExFreePoolWithTag(Prev->ai_addr, VIOSOCK_WSK_MEMORY_TAG); + ExFreePoolWithTag(Prev, VIOSOCK_WSK_MEMORY_TAG); + } + + DEBUG_EXIT_FUNCTION_VOID(); + return; } NTSTATUS @@ -285,16 +362,43 @@ VioWskGetNameInfo( _In_opt_ PEPROCESS OwningProcess, _In_opt_ PETHREAD OwningThread, _Inout_ PIRP Irp - ) +) { + PSOCKADDR_VM VMAddr = NULL; + NTSTATUS Status = STATUS_UNSUCCESSFUL; + DEBUG_ENTER_FUNCTION("Client=0x%p; SockAddr=0x%p; SockAddrLength=%u; NodeName=0x%p; ServiceName=0x%p; Flags=0x%x; OwningProcess=0x%p; OwningThread=0x%p; Irp=0x%p", Client, SockAddr, SockAddrLength, NodeName, ServiceName, Flags, OwningProcess, OwningThread, Irp); + + PAGED_CODE(); UNREFERENCED_PARAMETER(Client); - UNREFERENCED_PARAMETER(NodeName); - UNREFERENCED_PARAMETER(ServiceName); - UNREFERENCED_PARAMETER(SockAddr); - UNREFERENCED_PARAMETER(SockAddrLength); UNREFERENCED_PARAMETER(Flags); UNREFERENCED_PARAMETER(OwningProcess); UNREFERENCED_PARAMETER(OwningThread); - return VioWskCompleteIrp(Irp, STATUS_NOT_IMPLEMENTED, 0); -} \ No newline at end of file + Status = STATUS_SUCCESS; + VMAddr = (PSOCKADDR_VM)SockAddr; + if ((NodeName == NULL && ServiceName == NULL) || + SockAddrLength != sizeof(SOCKADDR_VM) || + VMAddr->svm_family != AF_VSOCK) { + Status = STATUS_INVALID_PARAMETER; + goto CompleteIrp; + } + + if (NodeName != NULL) { + Status = VioWskAddressPartToString(VMAddr->svm_cid, NodeName); + if (!NT_SUCCESS(Status)) + goto CompleteIrp; + } + + if (ServiceName != NULL) { + Status = VioWskAddressPartToString(VMAddr->svm_port, ServiceName); + if (!NT_SUCCESS(Status)) + goto CompleteIrp; + } + +CompleteIrp: + if (Irp) + Status = VioWskIrpComplete(NULL, Irp, Status, 0); + + DEBUG_EXIT_FUNCTION("0x%x, Nodename=\"%wZ\", Servicename=\"%wZ\"", Status, NodeName, ServiceName); + return Status; +} diff --git a/viosock/wsk/viowsk.h b/viosock/wsk/viowsk.h index 86e243530..cd908e9c6 100644 --- a/viosock/wsk/viowsk.h +++ b/viosock/wsk/viowsk.h @@ -152,6 +152,8 @@ VioWskCompleteIrp( typedef struct _VIOWSK_SOCKET { WSK_SOCKET WskSocket; + volatile LONG RequestCount; + IO_REMOVE_LOCK CloseRemoveLock; PVOID SocketContext; ULONG Type; WSK_CLIENT_LISTEN_DISPATCH ListenDispatch; diff --git a/viosock/wsk/wsk-utils.c b/viosock/wsk/wsk-utils.c new file mode 100644 index 000000000..7b00e5ea4 --- /dev/null +++ b/viosock/wsk/wsk-utils.c @@ -0,0 +1,167 @@ +/* + * Provider NPI functions + * + * Copyright (c) 2021 Virtuozzo International GmbH + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions + * are met : + * 1. Redistributions of source code must retain the above copyright + * notice, this list of conditions and the following disclaimer. + * 2. Redistributions in binary form must reproduce the above copyright + * notice, this list of conditions and the following disclaimer in the + * documentation and / or other materials provided with the distribution. + * 3. Neither the names of the copyright holders nor the names of their contributors + * may be used to endorse or promote products derived from this software + * without specific prior written permission. + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS ``AS IS'' AND + * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + * ARE DISCLAIMED.IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS + * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) + * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT + * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY + * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF + * SUCH DAMAGE. + */ + +#include "precomp.h" +#include "viowsk.h" +#include "..\inc\debug-utils.h" +#include "..\inc\vio_wsk.h" +#include "viowsk-internal.h" +#include "wsk-utils.h" + +#ifdef ALLOC_PRAGMA +#endif + + + + +_Must_inspect_result_ +NTSTATUS +VioWskIrpAcquire( + _In_opt_ PVIOWSK_SOCKET Socket, + _Inout_ PIRP Irp +) +{ + LONG reqCount = 0; + NTSTATUS Status = STATUS_UNSUCCESSFUL; + DEBUG_ENTER_FUNCTION("Socket=0x%p; Irp=0x%p", Socket, Irp); + + Status = STATUS_SUCCESS; + + if (Socket) + { + Status = IoAcquireRemoveLock(&Socket->CloseRemoveLock, Irp); + if (NT_SUCCESS(Status)) + reqCount = InterlockedIncrement(&Socket->RequestCount); + + if (!NT_SUCCESS(Status)) + Status = STATUS_FILE_FORCED_CLOSED; + } + + DEBUG_EXIT_FUNCTION("0x%x, reqCount=%i", Status, reqCount); + return Status; +} + + +void +VioWskIrpRelease( + _In_opt_ PVIOWSK_SOCKET Socket, + _In_ PIRP Irp +) +{ + LONG reqCount = 0; + DEBUG_ENTER_FUNCTION("Socket=0x%p; Irp=0x%p", Socket, Irp); + + if (Socket) + { + IoReleaseRemoveLock(&Socket->CloseRemoveLock, Irp); + reqCount = InterlockedDecrement(&Socket->RequestCount); + } + + DEBUG_EXIT_FUNCTION("void, reqCount=%i", reqCount); + return; +} + + +NTSTATUS +VioWskIrpComplete( + _Inout_opt_ PVIOWSK_SOCKET Socket, + _In_ PIRP Irp, + _In_ NTSTATUS Status, + _In_ ULONG_PTR Information +) +{ + DEBUG_ENTER_FUNCTION("Socket=0x%p; Irp=0x%p; Status=0x%x; Information=%zu", Socket, Irp, Status, Information); + + IoSetCancelRoutine(Irp, NULL); + if (Socket) + VioWskIrpRelease(Socket, Irp); + + Irp->IoStatus.Information = Information; + Irp->IoStatus.Status = Status; + IoSetNextIrpStackLocation(Irp); + IoCompleteRequest(Irp, IO_NO_INCREMENT); + + DEBUG_EXIT_FUNCTION("0x%x", Status); + return Status; +} + + +_Must_inspect_result_ +NTSTATUS +VioWskAddressPartToString( + _In_ ULONG Value, + _Out_ PUNICODE_STRING String +) +{ + NTSTATUS Status = STATUS_UNSUCCESSFUL; + DEBUG_ENTER_FUNCTION("Value=%u; String=0x%p", Value, String); + + String->Length = 0; + Status = RtlIntegerToUnicodeString(Value, 0, String); + if (!NT_SUCCESS(Status)) + goto Exit; + + if (String->Length == String->MaximumLength) { + Status = STATUS_BUFFER_OVERFLOW; + goto Exit; + } + + String->Buffer[String->Length / sizeof(wchar_t)] = L'\0'; + +Exit: + DEBUG_EXIT_FUNCTION("0x%x, *String=\"%wZ\"", Status, String); + return Status; +} + +_Must_inspect_result_ +NTSTATUS +VioWskStringToAddressPart( + _In_ PUNICODE_STRING String, + _Out_ PULONG Value +) +{ + NTSTATUS Status = STATUS_UNSUCCESSFUL; + DEBUG_ENTER_FUNCTION("String=\"%wZ\"; Value=0x%p", String, Value); + + *Value = 0; + if (String->Length > 0) { + if (String->Buffer[0] < L'0' || String->Buffer[0] > '9') { + Status = STATUS_INVALID_PARAMETER; + goto Exit; + } + + Status = RtlUnicodeStringToInteger(String, 0, Value); + if (!NT_SUCCESS(Status)) + goto Exit; + } + +Exit: + DEBUG_EXIT_FUNCTION("0x%x, *Value=%u", Status, *Value); + return Status; +} diff --git a/viosock/wsk/wsk-utils.h b/viosock/wsk/wsk-utils.h new file mode 100644 index 000000000..0ecdc8253 --- /dev/null +++ b/viosock/wsk/wsk-utils.h @@ -0,0 +1,78 @@ +/* + * Provider NPI functions + * + * Copyright (c) 2021 Virtuozzo International GmbH + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions + * are met : + * 1. Redistributions of source code must retain the above copyright + * notice, this list of conditions and the following disclaimer. + * 2. Redistributions in binary form must reproduce the above copyright + * notice, this list of conditions and the following disclaimer in the + * documentation and / or other materials provided with the distribution. + * 3. Neither the names of the copyright holders nor the names of their contributors + * may be used to endorse or promote products derived from this software + * without specific prior written permission. + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS ``AS IS'' AND + * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + * ARE DISCLAIMED.IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS + * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) + * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT + * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY + * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF + * SUCH DAMAGE. + */ + +#ifndef _WSK_UTILS_H +#define _WSK_UTILS_H + +#ifdef _MSC_VER +#pragma once +#endif //_MSC_VER + + +#include "..\inc\vio_wsk.h" +#include "viowsk.h" + + +_Must_inspect_result_ +NTSTATUS +VioWskIrpAcquire( + _In_opt_ PVIOWSK_SOCKET Socket, + _Inout_ PIRP Irp +); + +void +VioWskIrpRelease( + _In_opt_ PVIOWSK_SOCKET Socket, + _In_ PIRP Irp +); + +NTSTATUS +VioWskIrpComplete( + _Inout_opt_ PVIOWSK_SOCKET Socket, + _In_ PIRP Irp, + _In_ NTSTATUS Status, + _In_ ULONG_PTR Information +); + +_Must_inspect_result_ +NTSTATUS +VioWskAddressPartToString( + _In_ ULONG Value, + _Out_ PUNICODE_STRING String +); + +_Must_inspect_result_ +NTSTATUS +VioWskStringToAddressPart( + _In_ PUNICODE_STRING String, + _Out_ PULONG Value +); + + +#endif diff --git a/viosock/wsk/wsk.vcxproj b/viosock/wsk/wsk.vcxproj index 2892759cd..f38fc703d 100644 --- a/viosock/wsk/wsk.vcxproj +++ b/viosock/wsk/wsk.vcxproj @@ -31,11 +31,13 @@ + + {96FDD976-0035-4E24-A61B-E93BED675101} diff --git a/viosock/wsk/wsk.vcxproj.filters b/viosock/wsk/wsk.vcxproj.filters index a865e460c..2ef681aa0 100644 --- a/viosock/wsk/wsk.vcxproj.filters +++ b/viosock/wsk/wsk.vcxproj.filters @@ -27,6 +27,9 @@ Header Files + + Header Files + @@ -38,5 +41,8 @@ Source Files + + Source Files + \ No newline at end of file From 12333f5ffd54fa62a84abe5cac144380f3045272 Mon Sep 17 00:00:00 2001 From: Martin Drab Date: Sun, 24 Jul 2022 10:13:13 -0700 Subject: [PATCH 05/18] WSK: Ability to create and close sockets Signed-off-by: Martin Drab --- viosock/wsk/provider.c | 179 +++++++--------------- viosock/wsk/socket-internal.c | 260 ++++++++++++++++++++++++++++++++ viosock/wsk/socket.c | 40 ++++- viosock/wsk/viowsk-internal.h | 68 +++++++++ viosock/wsk/viowsk.c | 2 +- viosock/wsk/viowsk.h | 4 + viosock/wsk/wsk-workitem.c | 187 +++++++++++++++++++++++ viosock/wsk/wsk-workitem.h | 110 ++++++++++++++ viosock/wsk/wsk.vcxproj | 3 + viosock/wsk/wsk.vcxproj.filters | 9 ++ 10 files changed, 731 insertions(+), 131 deletions(-) create mode 100644 viosock/wsk/socket-internal.c create mode 100644 viosock/wsk/wsk-workitem.c create mode 100644 viosock/wsk/wsk-workitem.h diff --git a/viosock/wsk/provider.c b/viosock/wsk/provider.c index b2bcc26f1..486552756 100644 --- a/viosock/wsk/provider.c +++ b/viosock/wsk/provider.c @@ -31,6 +31,8 @@ #include "viowsk.h" #include "..\inc\debug-utils.h" #include "wsk-utils.h" +#include "viowsk-internal.h" +#include "wsk-workitem.h" #include "..\inc\vio_wsk.h" #ifdef ALLOC_PRAGMA @@ -39,100 +41,6 @@ #pragma alloc_text (PAGE, VioWskGetNameInfo) #endif -static -NTSTATUS -VioWskSocketInternal( - _In_ PWSK_CLIENT Client, - _In_ ULONG Flags, - _In_opt_ PVOID SocketContext, - _In_opt_ CONST VOID *Dispatch, - _In_opt_ PEPROCESS OwningProcess, - _In_opt_ PETHREAD OwningThread, - _In_opt_ PSECURITY_DESCRIPTOR SecurityDescriptor, - _Out_ PVIOWSK_SOCKET *pNewSocket, - _Out_ PBOOLEAN CompleteIrp -) -{ - PWSK_CLIENT_LISTEN_DISPATCH pListenDispatch = NULL; - PWSK_CLIENT_CONNECTION_DISPATCH pConnectionDispatch = NULL; - PVIOWSK_SOCKET pSocket; - NTSTATUS Status; - PVOID pProviderDispatch; - SIZE_T ProviderDispatchSize; - - UNREFERENCED_PARAMETER(Client); - UNREFERENCED_PARAMETER(OwningProcess); - UNREFERENCED_PARAMETER(OwningThread); - UNREFERENCED_PARAMETER(SecurityDescriptor); - - ASSERT(pNewSocket&&CompleteIrp); - - *pNewSocket = NULL; - *CompleteIrp = FALSE; - - switch (Flags) - { - case WSK_FLAG_BASIC_SOCKET: - pProviderDispatch = &gBasicDispatch; - ProviderDispatchSize = sizeof(gBasicDispatch); - break; - case WSK_FLAG_LISTEN_SOCKET: - pListenDispatch = (PWSK_CLIENT_LISTEN_DISPATCH)Dispatch; - pProviderDispatch = &gListenDispatch; - ProviderDispatchSize = sizeof(gListenDispatch); - break; - case WSK_FLAG_CONNECTION_SOCKET: - pConnectionDispatch = (PWSK_CLIENT_CONNECTION_DISPATCH)Dispatch; - pProviderDispatch = &gConnectionDispatch; - ProviderDispatchSize = sizeof(gConnectionDispatch); - break; - case WSK_FLAG_DATAGRAM_SOCKET: - *CompleteIrp = TRUE; - return STATUS_NOT_SUPPORTED; - -#if (NTDDI_VERSION >= NTDDI_WIN10_RS2) - case WSK_FLAG_STREAM_SOCKET: - if (Dispatch) - { - pListenDispatch = (PWSK_CLIENT_LISTEN_DISPATCH)((PWSK_CLIENT_STREAM_DISPATCH)Dispatch)->Listen; - pConnectionDispatch = (PWSK_CLIENT_CONNECTION_DISPATCH)((PWSK_CLIENT_STREAM_DISPATCH)Dispatch)->Connect; - } - pProviderDispatch = &gStreamDispatch; - ProviderDispatchSize = sizeof(gStreamDispatch); - break; -#endif // if (NTDDI_VERSION >= NTDDI_WIN10_RS2) - default: - return STATUS_INVALID_PARAMETER; - } - - *CompleteIrp = TRUE; - - pSocket = ExAllocatePoolUninitialized(NonPagedPoolNx, sizeof(VIOWSK_SOCKET), VIOSOCK_WSK_MEMORY_TAG); - if (pSocket) - { - RtlZeroMemory(pSocket, sizeof(VIOWSK_SOCKET)); - - pSocket->SocketContext = SocketContext; - pSocket->Type = Flags; - - if (pListenDispatch) - RtlCopyMemory(&pSocket->ListenDispatch, pListenDispatch, sizeof(*pListenDispatch)); - - if (pConnectionDispatch) - RtlCopyMemory(&pSocket->ConnectionDispatch, pConnectionDispatch, sizeof(*pConnectionDispatch)); - - RtlCopyMemory(&pSocket->ProviderDispatch, pProviderDispatch, ProviderDispatchSize); - pSocket->WskSocket.Dispatch = &pSocket->ProviderDispatch; - - //TODO: Open socket - *pNewSocket = pSocket; - Status = STATUS_SUCCESS; - } - else - Status = STATUS_INSUFFICIENT_RESOURCES; - - return Status; -} NTSTATUS WSKAPI @@ -150,25 +58,46 @@ VioWskSocket( _Inout_ PIRP Irp ) { - PVIOWSK_SOCKET pSocket; - NTSTATUS Status; - BOOLEAN CompleteIrp; - - if (!Client || !Irp) - return STATUS_INVALID_PARAMETER; - - if (AddressFamily != AF_VSOCK || SocketType != SOCK_STREAM || Protocol != 0) - return VioWskCompleteIrp(Irp, STATUS_NOT_SUPPORTED, 0); - - Status = VioWskSocketInternal(Client, Flags, SocketContext, Dispatch, OwningProcess, - OwningThread, SecurityDescriptor, &pSocket, &CompleteIrp); + PWSK_WORKITEM WorkItem = NULL; + PVIOWSK_SOCKET pSocket = NULL; + NTSTATUS Status = STATUS_UNSUCCESSFUL; + DEBUG_ENTER_FUNCTION("Client=0x%p; AddressFamily=%u; SocketType=%u; Protocol=%u; Flags=0x%x; SocketContext=0x%p; Dispatch=0x%p; OwningProcess=0x%p; OwningThread=0x%p; SecurityDescriptor=0x%p; Irp=0x%p", Client, AddressFamily, SocketType, Protocol, Flags, SocketContext, Dispatch, OwningProcess, OwningThread, SecurityDescriptor, Irp); + + _At_((void*)Irp->IoStatus.Information, __drv_allocatesMem(Mem)) + + if (KeGetCurrentIrql() > PASSIVE_LEVEL) { + WorkItem = WskWorkItemAlloc(wskwitSocket, Irp); + if (!WorkItem) { + Status = STATUS_INSUFFICIENT_RESOURCES; + goto CompleteIrp; + } + + WorkItem->Specific.Socket.AddressFamily = AddressFamily; + WorkItem->Specific.Socket.Client = Client; + WorkItem->Specific.Socket.Dispatch = Dispatch; + WorkItem->Specific.Socket.Flags = Flags; + WorkItem->Specific.Socket.OwningProcess = OwningProcess; + WorkItem->Specific.Socket.OwningThread = OwningThread; + WorkItem->Specific.Socket.Protocol = Protocol; + WorkItem->Specific.Socket.SecurityDescriptor = SecurityDescriptor; + WorkItem->Specific.Socket.SocketContext = SocketContext; + WorkItem->Specific.Socket.SocketType = SocketType; + WskWorkItemQueue(WorkItem); + Status = STATUS_PENDING; + goto Exit; + } - if (CompleteIrp) - VioWskCompleteIrp(Irp, Status, (ULONG_PTR)pSocket); + Status = VioWskSocketInternal(Client, NULL, Flags, SocketContext, Dispatch, OwningProcess, + OwningThread, SecurityDescriptor, &pSocket); - return Status; +CompleteIrp: + VioWskIrpComplete(NULL, Irp, Status, (ULONG_PTR)pSocket); +Exit: + DEBUG_EXIT_FUNCTION("0x%x", Status); + return Status; } +_At_(Irp->IoStatus.Information, __drv_allocatesMem(Mem)) NTSTATUS WSKAPI VioWskSocketConnect( @@ -186,31 +115,25 @@ VioWskSocketConnect( _Inout_ PIRP Irp ) { - PVIOWSK_SOCKET pSocket; - NTSTATUS Status; - BOOLEAN CompleteIrp; + PVIOWSK_SOCKET pSocket = NULL; + NTSTATUS Status = STATUS_UNSUCCESSFUL; + UNREFERENCED_PARAMETER(Client); + UNREFERENCED_PARAMETER(SocketType); + UNREFERENCED_PARAMETER(Protocol); UNREFERENCED_PARAMETER(LocalAddress); UNREFERENCED_PARAMETER(RemoteAddress); UNREFERENCED_PARAMETER(Flags); + UNREFERENCED_PARAMETER(SocketContext); + UNREFERENCED_PARAMETER(Dispatch); + UNREFERENCED_PARAMETER(OwningProcess); + UNREFERENCED_PARAMETER(OwningThread); + UNREFERENCED_PARAMETER(SecurityDescriptor); - if (!Client || !Irp) - return STATUS_INVALID_PARAMETER; - - if (SocketType != SOCK_STREAM || Protocol != 0) - return VioWskCompleteIrp(Irp, STATUS_NOT_SUPPORTED, 0); - - Status = VioWskSocketInternal(Client, WSK_FLAG_CONNECTION_SOCKET, SocketContext, - Dispatch, OwningProcess, OwningThread, SecurityDescriptor, &pSocket, &CompleteIrp); - - if (NT_SUCCESS(Status)) - { - //TODO: bind and connect - } - - if (CompleteIrp) - VioWskCompleteIrp(Irp, Status, (ULONG_PTR)pSocket); + Status = STATUS_NOT_IMPLEMENTED; + VioWskCompleteIrp(Irp, Status, (ULONG_PTR)pSocket); + DEBUG_EXIT_FUNCTION("0x%x", Status); return Status; } diff --git a/viosock/wsk/socket-internal.c b/viosock/wsk/socket-internal.c new file mode 100644 index 000000000..a4185daef --- /dev/null +++ b/viosock/wsk/socket-internal.c @@ -0,0 +1,260 @@ +/* + * Exports definition for virtio socket WSK interface + * + * Copyright (c) 2021 Virtuozzo International GmbH + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions + * are met : + * 1. Redistributions of source code must retain the above copyright + * notice, this list of conditions and the following disclaimer. + * 2. Redistributions in binary form must reproduce the above copyright + * notice, this list of conditions and the following disclaimer in the + * documentation and / or other materials provided with the distribution. + * 3. Neither the names of the copyright holders nor the names of their contributors + * may be used to endorse or promote products derived from this software + * without specific prior written permission. + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS ``AS IS'' AND + * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + * ARE DISCLAIMED.IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS + * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) + * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT + * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY + * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF + * SUCH DAMAGE. + */ + +#include "precomp.h" +#include "viowsk.h" +#include "..\inc\debug-utils.h" +#include "wsk-utils.h" +#include "viowsk-internal.h" + +#ifdef ALLOC_PRAGMA +#pragma alloc_text (PAGE, VioWskSocketInternal) +#pragma alloc_text (PAGE, VioWskCloseSocketInternal) +#endif + + +_Must_inspect_result_ +static +NTSTATUS +_VioSocketCreate( + _In_ PWSK_CLIENT Client, + _In_opt_ PVIOWSK_SOCKET ListenSocket, + _In_ ULONG SocketType, + _Out_ PHANDLE SocketFileHandle, + _Out_ PFILE_OBJECT *SocketFileObject +) +{ + NTSTATUS Status = STATUS_UNSUCCESSFUL; + PVIOWSK_REG_CONTEXT pContext = NULL; + PWSK_REGISTRATION reg = NULL; + HANDLE hFile = NULL; + OBJECT_ATTRIBUTES ObjectAttributes; + IO_STATUS_BLOCK Iosb = { 0 }; + VIRTIO_VSOCK_PARAMS SockParams = { 0 }; + UCHAR EaBuffer[sizeof(FILE_FULL_EA_INFORMATION) + sizeof(SockParams)] = { 0 }; + PFILE_FULL_EA_INFORMATION pEaBuffer = (PFILE_FULL_EA_INFORMATION)EaBuffer; + DEBUG_ENTER_FUNCTION("Client=0x%p; ListenSocket=0x%p; Sockettype=%u; SocketFileHandle=0x%p; SocketFileObject=0x%p", Client, ListenSocket, SocketType, SocketFileHandle, SocketFileObject); + + PAGED_CODE(); + reg = (PWSK_REGISTRATION)Client; + pContext = (PVIOWSK_REG_CONTEXT)reg->ReservedRegistrationContext; + RtlSecureZeroMemory(&SockParams, sizeof(SockParams)); + if (ListenSocket != NULL) + SockParams.Socket = (ULONGLONG)ListenSocket->FileHandle; + + SockParams.Type = SocketType; + pEaBuffer->EaNameLength = sizeof(FILE_FULL_EA_INFORMATION) - + FIELD_OFFSET(FILE_FULL_EA_INFORMATION, EaName) - sizeof(UCHAR); + pEaBuffer->EaValueLength = sizeof(SockParams); + RtlCopyMemory(&EaBuffer[sizeof(FILE_FULL_EA_INFORMATION)], &SockParams, sizeof(SockParams)); + + InitializeObjectAttributes(&ObjectAttributes, &pContext->VIOSockLinkName, OBJ_CASE_INSENSITIVE, NULL, NULL); + + Status = ZwCreateFile(&hFile, FILE_GENERIC_READ | FILE_GENERIC_WRITE, &ObjectAttributes, &Iosb, 0, 0, + FILE_SHARE_READ | FILE_SHARE_WRITE, FILE_OPEN, FILE_NON_DIRECTORY_FILE, pEaBuffer, + pEaBuffer ? sizeof(EaBuffer) : 0); + + if (NT_SUCCESS(Status)) + { + Status = ObReferenceObjectByHandle(hFile, FILE_GENERIC_READ | FILE_GENERIC_WRITE, *IoFileObjectType, KernelMode, (PVOID*)SocketFileObject, NULL); + if (NT_SUCCESS(Status)) + *SocketFileHandle = hFile; + } + + DEBUG_EXIT_FUNCTION("0x%x, *SocketFileHandle=0x%p, *SocketFileObject=0x%p", Status, *SocketFileHandle, *SocketFileObject); + return Status; +} + + + +_Must_inspect_result_ +NTSTATUS +VioWskSocketInternal( + _In_ PWSK_CLIENT Client, + _In_opt_ PVIOWSK_SOCKET ListenSocket, + _In_ ULONG Flags, + _In_opt_ PVOID SocketContext, + _In_opt_ CONST VOID* Dispatch, + _In_opt_ PEPROCESS OwningProcess, + _In_opt_ PETHREAD OwningThread, + _In_opt_ PSECURITY_DESCRIPTOR SecurityDescriptor, + _Out_ PVIOWSK_SOCKET* pNewSocket +) +{ + PWSK_CLIENT_LISTEN_DISPATCH pListenDispatch = NULL; + PWSK_CLIENT_CONNECTION_DISPATCH pConnectionDispatch = NULL; + PVIOWSK_SOCKET pSocket = NULL; + NTSTATUS Status = STATUS_UNSUCCESSFUL; + PVOID pProviderDispatch = NULL; + SIZE_T ProviderDispatchSize = 0; + VIRTIO_VSOCK_CONFIG SocketConfig; + PIRP ConfigIrp = NULL; + KEVENT ConfigEvent; + IO_STATUS_BLOCK ConfigIosb; + DEBUG_ENTER_FUNCTION("Client=0x%p; ListenSocket=0x%p; Flags=0x%x; SocketContext=0x%p; Dispatch=0x%p; OwningProcess=0x%p; OwningThread=0x%p; SecurityDescriptor=0x%p; pNewSocket=0x%p", Client, ListenSocket, Flags, SocketContext, Dispatch, OwningProcess, OwningThread, SecurityDescriptor, pNewSocket); + + PAGED_CODE(); + UNREFERENCED_PARAMETER(OwningProcess); + UNREFERENCED_PARAMETER(OwningThread); + UNREFERENCED_PARAMETER(SecurityDescriptor); + + *pNewSocket = NULL; + Status = STATUS_SUCCESS; + + switch (Flags) + { + case WSK_FLAG_BASIC_SOCKET: + pProviderDispatch = &gBasicDispatch; + ProviderDispatchSize = sizeof(gBasicDispatch); + break; + case WSK_FLAG_LISTEN_SOCKET: + pListenDispatch = (PWSK_CLIENT_LISTEN_DISPATCH)Dispatch; + pProviderDispatch = &gListenDispatch; + ProviderDispatchSize = sizeof(gListenDispatch); + break; + case WSK_FLAG_CONNECTION_SOCKET: + pConnectionDispatch = (PWSK_CLIENT_CONNECTION_DISPATCH)Dispatch; + pProviderDispatch = &gConnectionDispatch; + ProviderDispatchSize = sizeof(gConnectionDispatch); + break; + case WSK_FLAG_DATAGRAM_SOCKET: + Status = STATUS_NOT_SUPPORTED; + break; +#if (NTDDI_VERSION >= NTDDI_WIN10_RS2) + case WSK_FLAG_STREAM_SOCKET: + if (Dispatch) + { + pListenDispatch = (PWSK_CLIENT_LISTEN_DISPATCH)((PWSK_CLIENT_STREAM_DISPATCH)Dispatch)->Listen; + pConnectionDispatch = (PWSK_CLIENT_CONNECTION_DISPATCH)((PWSK_CLIENT_STREAM_DISPATCH)Dispatch)->Connect; + } + + pProviderDispatch = &gStreamDispatch; + ProviderDispatchSize = sizeof(gStreamDispatch); + break; +#endif // if (NTDDI_VERSION >= NTDDI_WIN10_RS2) + default: + Status = STATUS_INVALID_PARAMETER; + break; + } + + if (!NT_SUCCESS(Status)) + goto Exit; + + pSocket = ExAllocatePoolUninitialized(NonPagedPoolNx, sizeof(VIOWSK_SOCKET), VIOSOCK_WSK_MEMORY_TAG); + if (!pSocket) { + Status = STATUS_INSUFFICIENT_RESOURCES; + goto Exit; + } + + memset(pSocket, 0, sizeof(VIOWSK_SOCKET)); + pSocket->Client = Client; + pSocket->SocketContext = SocketContext; + pSocket->Type = Flags; + if (pListenDispatch) + RtlCopyMemory(&pSocket->ListenDispatch, pListenDispatch, sizeof(*pListenDispatch)); + + if (pConnectionDispatch) + RtlCopyMemory(&pSocket->ConnectionDispatch, pConnectionDispatch, sizeof(*pConnectionDispatch)); + + if (pProviderDispatch) + RtlCopyMemory(&pSocket->ProviderDispatch, pProviderDispatch, ProviderDispatchSize); + + pSocket->WskSocket.Dispatch = &pSocket->ProviderDispatch; + Status = _VioSocketCreate(Client, ListenSocket, SOCK_STREAM, &pSocket->FileHandle, &pSocket->FileObject); + if (!NT_SUCCESS(Status)) + goto FreeSocket; + + KeInitializeEvent(&ConfigEvent, NotificationEvent, FALSE); + ConfigIrp = IoBuildDeviceIoControlRequest(IOCTL_GET_CONFIG, pContext->VIOSockDevice, NULL, 0, &SocketConfig, sizeof(SocketConfig), FALSE, &ConfigEvent, &ConfigIosb); + if (!ConfigIrp) + { + Status = STATUS_INSUFFICIENT_RESOURCES; + goto CloseSocket; + } + + ObReferenceObject(pSocket->FileObject); + ConfigIrp->Tail.Overlay.OriginalFileObject = pSocket->FileObject; + IoGetNextIrpStackLocation(ConfigIrp)->FileObject = pSocket->FileObject; + Status = IoCallDriver(pContext->VIOSockDevice, ConfigIrp); + if (Status == STATUS_PENDING) { + KeWaitForSingleObject(&ConfigEvent, Executive, KernelMode, FALSE, NULL); + Status = ConfigIosb.Status; + } + + if (NT_SUCCESS(Status) && ConfigIosb.Information < sizeof(SocketConfig)) + Status = STATUS_INVALID_PARAMETER; + + if (!NT_SUCCESS(Status)) + goto CloseSocket; + + pSocket->GuestId = SocketConfig.guest_cid; + IoInitializeRemoveLock(&pSocket->CloseRemoveLock, VIOSOCK_WSK_MEMORY_TAG, 0, 0x7FFFFFFF); + *pNewSocket = pSocket; + pSocket = NULL; + +CloseSocket: + if (pSocket) + { + ZwClose(pSocket->FileHandle); + ObDereferenceObject(pSocket->FileObject); + } +FreeSocket: + if (pSocket) + ExFreePoolWithTag(pSocket, VIOSOCK_WSK_MEMORY_TAG); +Exit: + DEBUG_EXIT_FUNCTION("0x%x, *pNewSocket=0x%p", Status, *pNewSocket); + return Status; +} + + +NTSTATUS +VioWskCloseSocketInternal( + _Inout_ PVIOWSK_SOCKET Socket, + _In_opt_ PVOID ReleaseTag +) +{ + HANDLE fileHandle = NULL; + NTSTATUS Status = STATUS_UNSUCCESSFUL; + DEBUG_ENTER_FUNCTION("Socket=0x%p; ReleaseTag=0x%p", Socket, ReleaseTag); + + PAGED_CODE(); + fileHandle = InterlockedExchangePointer(&Socket->FileHandle, NULL); + if (fileHandle) + { + Status = ZwClose(fileHandle); + if (ReleaseTag) + IoReleaseRemoveLockAndWait(&Socket->CloseRemoveLock, ReleaseTag); + + ObDereferenceObject(Socket->FileObject); + ExFreePoolWithTag(Socket, VIOSOCK_WSK_MEMORY_TAG); + } + + DEBUG_EXIT_FUNCTION("0x%x", Status); + return Status; +} diff --git a/viosock/wsk/socket.c b/viosock/wsk/socket.c index 133e0b06a..7aafea06d 100644 --- a/viosock/wsk/socket.c +++ b/viosock/wsk/socket.c @@ -28,7 +28,11 @@ */ #include "precomp.h" +#include "..\inc\debug-utils.h" #include "viowsk.h" +#include "wsk-utils.h" +#include "viowsk-internal.h" +#include "wsk-workitem.h" #include "..\inc\vio_wsk.h" NTSTATUS @@ -280,9 +284,41 @@ VioWskCloseSocket( _Inout_ PIRP Irp ) { - UNREFERENCED_PARAMETER(Socket); + PWSK_WORKITEM WorkItem = NULL; + NTSTATUS Status = STATUS_UNSUCCESSFUL; + PVIOWSK_SOCKET pSocket = CONTAINING_RECORD(Socket, VIOWSK_SOCKET, WskSocket); + DEBUG_ENTER_FUNCTION("Socket=0x%p; Irp=0x%p", Socket, Irp); - return VioWskCompleteIrp(Irp, STATUS_NOT_IMPLEMENTED, 0); + if (KeGetCurrentIrql() > PASSIVE_LEVEL) + { + WorkItem = WskWorkItemAlloc(wskwitCloseSocket, Irp); + if (!WorkItem) + { + Status = STATUS_INSUFFICIENT_RESOURCES; + goto CompleteIrp; + } + + WorkItem->Specific.CloseSocket.Socket = Socket; + WskWorkItemQueue(WorkItem); + Status = STATUS_PENDING; + goto Exit; + } + + Status = VioWskIrpAcquire(pSocket, Irp); + if (!NT_SUCCESS(Status)) + { + pSocket = NULL; + goto CompleteIrp; + } + + Status = VioWskCloseSocketInternal(pSocket, Irp); + pSocket = NULL; + +CompleteIrp: + VioWskIrpComplete(pSocket, Irp, Status, 0); +Exit: + DEBUG_EXIT_FUNCTION("0x%x", Status); + return Status; } NTSTATUS diff --git a/viosock/wsk/viowsk-internal.h b/viosock/wsk/viowsk-internal.h index ec90c1aa2..c8f1fe5de 100644 --- a/viosock/wsk/viowsk-internal.h +++ b/viosock/wsk/viowsk-internal.h @@ -1,3 +1,31 @@ +/* + * Exports definition for virtio socket WSK interface + * + * Copyright (c) 2021 Virtuozzo International GmbH + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions + * are met : + * 1. Redistributions of source code must retain the above copyright + * notice, this list of conditions and the following disclaimer. + * 2. Redistributions in binary form must reproduce the above copyright + * notice, this list of conditions and the following disclaimer in the + * documentation and / or other materials provided with the distribution. + * 3. Neither the names of the copyright holders nor the names of their contributors + * may be used to endorse or promote products derived from this software + * without specific prior written permission. + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS ``AS IS'' AND + * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + * ARE DISCLAIMED.IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS + * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) + * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT + * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY + * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF + * SUCH DAMAGE. + */ #ifndef __VIOWSK_INTERNAL_H__ #define __VIOWSK_INTERNAL_H__ @@ -8,5 +36,45 @@ extern PDRIVER_OBJECT _viowskDriverObject; extern PDEVICE_OBJECT _viowskDeviceObject; +_Must_inspect_result_ +NTSTATUS +VioWskSocketInternal( + _In_ PWSK_CLIENT Client, + _In_opt_ PVIOWSK_SOCKET ListenSocket, + _In_ ULONG Flags, + _In_opt_ PVOID SocketContext, + _In_opt_ CONST VOID* Dispatch, + _In_opt_ PEPROCESS OwningProcess, + _In_opt_ PETHREAD OwningThread, + _In_opt_ PSECURITY_DESCRIPTOR SecurityDescriptor, + _Out_ PVIOWSK_SOCKET* pNewSocket +); + +NTSTATUS +VioWskCloseSocketInternal( + _Inout_ PVIOWSK_SOCKET Socket, + _In_opt_ PVOID ReleleaseTag +); + +NTSTATUS +WSKAPI +VioWskCloseSocket( + _In_ PWSK_SOCKET Socket, + _Inout_ PIRP Irp +); + +NTSTATUS +WSKAPI +VioWskAccept( + _In_ PWSK_SOCKET ListenSocket, + _Reserved_ ULONG Flags, + _In_opt_ PVOID AcceptSocketContext, + _In_opt_ CONST WSK_CLIENT_CONNECTION_DISPATCH* AcceptSocketDispatch, + _Out_opt_ PSOCKADDR LocalAddress, + _Out_opt_ PSOCKADDR RemoteAddress, + _Inout_ PIRP Irp +); + + #endif diff --git a/viosock/wsk/viowsk.c b/viosock/wsk/viowsk.c index 3eca3a086..012d67499 100644 --- a/viosock/wsk/viowsk.c +++ b/viosock/wsk/viowsk.c @@ -29,9 +29,9 @@ #include "precomp.h" #include "..\inc\debug-utils.h" -#include "viowsk-internal.h" #include "viowsk.h" #include "..\inc\vio_wsk.h" +#include "viowsk-internal.h" #ifdef ALLOC_PRAGMA #pragma alloc_text (PAGE, VioWskRegister) diff --git a/viosock/wsk/viowsk.h b/viosock/wsk/viowsk.h index cd908e9c6..4e9851e5f 100644 --- a/viosock/wsk/viowsk.h +++ b/viosock/wsk/viowsk.h @@ -152,10 +152,14 @@ VioWskCompleteIrp( typedef struct _VIOWSK_SOCKET { WSK_SOCKET WskSocket; + PWSK_CLIENT Client; volatile LONG RequestCount; IO_REMOVE_LOCK CloseRemoveLock; PVOID SocketContext; ULONG Type; + ULONG GuestId; + HANDLE FileHandle; + PFILE_OBJECT FileObject; WSK_CLIENT_LISTEN_DISPATCH ListenDispatch; WSK_CLIENT_CONNECTION_DISPATCH ConnectionDispatch; union diff --git a/viosock/wsk/wsk-workitem.c b/viosock/wsk/wsk-workitem.c new file mode 100644 index 000000000..4c8f858cc --- /dev/null +++ b/viosock/wsk/wsk-workitem.c @@ -0,0 +1,187 @@ +/* + * Exports definition for virtio socket WSK interface + * + * Copyright (c) 2021 Virtuozzo International GmbH + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions + * are met : + * 1. Redistributions of source code must retain the above copyright + * notice, this list of conditions and the following disclaimer. + * 2. Redistributions in binary form must reproduce the above copyright + * notice, this list of conditions and the following disclaimer in the + * documentation and / or other materials provided with the distribution. + * 3. Neither the names of the copyright holders nor the names of their contributors + * may be used to endorse or promote products derived from this software + * without specific prior written permission. + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS ``AS IS'' AND + * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + * ARE DISCLAIMED.IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS + * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) + * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT + * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY + * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF + * SUCH DAMAGE. + */ + +#include "precomp.h" +#include "..\inc\debug-utils.h" +#include "viowsk.h" +#include "..\inc\vio_wsk.h" +#include "viowsk-internal.h" +#include "wsk-workitem.h" + + +#pragma warning(disable : 4996) + + +static +void +_ProcessWorkItem( + _In_ PWSK_WORKITEM WorkItem +) +{ + DEBUG_ENTER_FUNCTION("WorkItem=0x%p", WorkItem); + + switch (WorkItem->Type) + { + case wskwitSocket: + VioWskSocket(WorkItem->Specific.Socket.Client, + WorkItem->Specific.Socket.AddressFamily, + WorkItem->Specific.Socket.SocketType, + WorkItem->Specific.Socket.Protocol, + WorkItem->Specific.Socket.Flags, + WorkItem->Specific.Socket.SocketContext, + WorkItem->Specific.Socket.Dispatch, + WorkItem->Specific.Socket.OwningProcess, + WorkItem->Specific.Socket.OwningThread, + WorkItem->Specific.Socket.SecurityDescriptor, + WorkItem->Irp); + break; + case wskwitSocketAndConnect: + VioWskSocketConnect( + WorkItem->Specific.SocketConnect.Client, + WorkItem->Specific.SocketConnect.SocketType, + WorkItem->Specific.SocketConnect.Protocol, + WorkItem->Specific.SocketConnect.LocalAddress, + WorkItem->Specific.SocketConnect.RemoteAddress, + WorkItem->Specific.SocketConnect.Flags, + WorkItem->Specific.SocketConnect.SocketContext, + WorkItem->Specific.SocketConnect.Dispatch, + WorkItem->Specific.SocketConnect.OwningProcess, + WorkItem->Specific.SocketConnect.OwningThread, + WorkItem->Specific.SocketConnect.SecurityDescriptor, + WorkItem->Irp); + break; + case wskwitCloseSocket: + VioWskCloseSocket( + WorkItem->Specific.CloseSocket.Socket, + WorkItem->Irp); + break; + case wskwitAccept: + VioWskAccept( + WorkItem->Specific.Accept.ListenSocket, + WorkItem->Specific.Accept.Flags, + WorkItem->Specific.Accept.AcceptSocketContext, + WorkItem->Specific.Accept.AcceptSocketDispatch, + WorkItem->Specific.Accept.LocalAddress, + WorkItem->Specific.Accept.RemoteAddress, + WorkItem->Irp); + break; + default: + ASSERT(FALSE); + break; + } + + DEBUG_EXIT_FUNCTION_VOID(); + return; +} + +static +void +_IoWskRoutine( + _In_ PVOID IoObject, + _In_opt_ PVOID Context, + _In_ PIO_WORKITEM IoWorkItem +) +{ + PWSK_WORKITEM WorkItem = (PWSK_WORKITEM)Context; + DEBUG_ENTER_FUNCTION("IoObject=0x%p; Context=0x%p; IoWorkItem=0x%p", IoObject, Context, IoWorkItem); + + __analysis_assert(WorkItem); + UNREFERENCED_PARAMETER(IoObject); + _ProcessWorkItem(WorkItem); + IoUninitializeWorkItem(IoWorkItem); + ExFreePoolWithTag(WorkItem, VIOSOCK_WSK_MEMORY_TAG); + + DEBUG_EXIT_FUNCTION_VOID(); + return; +} + +static +void +_ExWskRoutine( + _In_opt_ PVOID Context +) +{ + PWSK_WORKITEM WorkItem = (PWSK_WORKITEM)Context; + DEBUG_ENTER_FUNCTION("Context=0x%p", Context); + + __analysis_assert(WorkItem); + _ProcessWorkItem(WorkItem); + ExFreePoolWithTag(WorkItem, VIOSOCK_WSK_MEMORY_TAG); + + DEBUG_EXIT_FUNCTION_VOID(); + return; +} + + + +PWSK_WORKITEM +WskWorkItemAlloc( + _In_ EWSKWorkItemType Type, + _In_ PIRP Irp +) +{ + PWSK_WORKITEM Ret = NULL; + SIZE_T TotalSize = sizeof(WSK_WORKITEM); + BOOLEAN IoWorkItem = (_viowskDeviceObject != NULL); + DEBUG_ENTER_FUNCTION("Type%u; Irp=0x%p", Type, Irp); + + if (IoWorkItem) + TotalSize += (IoSizeofWorkItem() - sizeof(Ret->ExWorkItem)); + + Ret = ExAllocatePoolWithTag(NonPagedPool,TotalSize , VIOSOCK_WSK_MEMORY_TAG); + if (Ret) + { + memset(Ret, 0, TotalSize); + Ret->Type = Type; + Ret->IoMethod = IoWorkItem; + Ret->Irp = Irp; + if (Ret->IoMethod) + IoInitializeWorkItem(_viowskDeviceObject, (PIO_WORKITEM)Ret->IoWorkItem); + else ExInitializeWorkItem(&Ret->ExWorkItem, _ExWskRoutine, Ret); + } + + DEBUG_EXIT_FUNCTION("0x%p", Ret); + return Ret; +} + + +void +WskWorkItemQueue( + _In_ PWSK_WORKITEM WorkItem +) +{ + DEBUG_ENTER_FUNCTION("WorkItem=0x%p", WorkItem); + + if (WorkItem->IoMethod) + IoQueueWorkItemEx((PIO_WORKITEM)WorkItem->IoWorkItem, _IoWskRoutine, DelayedWorkQueue, WorkItem); + else ExQueueWorkItem(&WorkItem->ExWorkItem, DelayedWorkQueue); + + DEBUG_EXIT_FUNCTION_VOID(); + return; +} diff --git a/viosock/wsk/wsk-workitem.h b/viosock/wsk/wsk-workitem.h new file mode 100644 index 000000000..3beaf945e --- /dev/null +++ b/viosock/wsk/wsk-workitem.h @@ -0,0 +1,110 @@ +/* + * Exports definition for virtio socket WSK interface + * + * Copyright (c) 2021 Virtuozzo International GmbH + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions + * are met : + * 1. Redistributions of source code must retain the above copyright + * notice, this list of conditions and the following disclaimer. + * 2. Redistributions in binary form must reproduce the above copyright + * notice, this list of conditions and the following disclaimer in the + * documentation and / or other materials provided with the distribution. + * 3. Neither the names of the copyright holders nor the names of their contributors + * may be used to endorse or promote products derived from this software + * without specific prior written permission. + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS ``AS IS'' AND + * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + * ARE DISCLAIMED.IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS + * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) + * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT + * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY + * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF + * SUCH DAMAGE. + */ + +#ifndef __WSK_WORKITEM_H__ +#define __WSK_WORKITEM_H__ + + +#include +#include "viowsk.h" + + + +typedef enum _EWSKWorkItemType { + wskwitUndefined, + wskwitSocket, + wskwitSocketAndConnect, + wskwitCloseSocket, + wskwitAccept, + wskwitMax, +} EWSKWorkItemType, *PEWSKWorkItemType; + +typedef struct _WSK_WORKITEM { + EWSKWorkItemType Type; + BOOLEAN IoMethod; + PIRP Irp; + union { + struct { + PWSK_CLIENT Client; + ADDRESS_FAMILY AddressFamily; + USHORT SocketType; + ULONG Protocol; + ULONG Flags; + PVOID SocketContext; + const VOID *Dispatch; + PEPROCESS OwningProcess; + PETHREAD OwningThread; + PSECURITY_DESCRIPTOR SecurityDescriptor; + } Socket; + struct { + PWSK_CLIENT Client; + USHORT SocketType; + ULONG Protocol; + PSOCKADDR LocalAddress; + PSOCKADDR RemoteAddress; + ULONG Flags; + PVOID SocketContext; + const WSK_CLIENT_CONNECTION_DISPATCH *Dispatch; + PEPROCESS OwningProcess; + PETHREAD OwningThread; + PSECURITY_DESCRIPTOR SecurityDescriptor; + } SocketConnect; + struct { + PWSK_SOCKET Socket; + } CloseSocket; + struct { + PWSK_SOCKET ListenSocket; + ULONG Flags; + PVOID AcceptSocketContext; + CONST WSK_CLIENT_CONNECTION_DISPATCH *AcceptSocketDispatch; + PSOCKADDR LocalAddress; + PSOCKADDR RemoteAddress; + } Accept; + } Specific; + union { + WORK_QUEUE_ITEM ExWorkItem; + unsigned char IoWorkItem[1]; + }; +} WSK_WORKITEM, *PWSK_WORKITEM; + + +PWSK_WORKITEM +WskWorkItemAlloc( + _In_ EWSKWorkItemType Type, + _In_ PIRP Irp +); + +void +WskWorkItemQueue( + _In_ PWSK_WORKITEM WorkItem +); + + + +#endif diff --git a/viosock/wsk/wsk.vcxproj b/viosock/wsk/wsk.vcxproj index f38fc703d..b99fb6430 100644 --- a/viosock/wsk/wsk.vcxproj +++ b/viosock/wsk/wsk.vcxproj @@ -32,12 +32,15 @@ + + + {96FDD976-0035-4E24-A61B-E93BED675101} diff --git a/viosock/wsk/wsk.vcxproj.filters b/viosock/wsk/wsk.vcxproj.filters index 2ef681aa0..7c6c3a32a 100644 --- a/viosock/wsk/wsk.vcxproj.filters +++ b/viosock/wsk/wsk.vcxproj.filters @@ -30,6 +30,9 @@ Header Files + + Header Files + @@ -44,5 +47,11 @@ Source Files + + Source Files + + + Source Files + \ No newline at end of file From f90ad5ad18f75bf749bb42c1a1f8a56fe5d24208 Mon Sep 17 00:00:00 2001 From: Martin Drab Date: Sun, 24 Jul 2022 12:20:21 -0700 Subject: [PATCH 06/18] WSK: Add logic for sending IOCTL requests Most of the work done by this WSK implementation consists of packing client driver's calls into IOCTL and read/write request and passing them to the VIOSOCK driver. This commit introduces routines for creating IOCTL requests and sending them to VIOSOCK. Also, a concept of WSK Completion Contexts is introduces. The purpose of these contexts is to transfer information between code sending IRPs and their completion routines. Also, these contexts can be used to implement one WSK interface call as a sequence of requests to the VIOSOCK driver. Signed-off-by: Martin Drab --- viosock/wsk/socket-internal.c | 47 +++++-- viosock/wsk/wsk-completion.c | 232 ++++++++++++++++++++++++++++++++ viosock/wsk/wsk-completion.h | 84 ++++++++++++ viosock/wsk/wsk-utils.c | 225 +++++++++++++++++++++++++++++++ viosock/wsk/wsk-utils.h | 32 +++++ viosock/wsk/wsk.vcxproj | 2 + viosock/wsk/wsk.vcxproj.filters | 6 + 7 files changed, 616 insertions(+), 12 deletions(-) create mode 100644 viosock/wsk/wsk-completion.c create mode 100644 viosock/wsk/wsk-completion.h diff --git a/viosock/wsk/socket-internal.c b/viosock/wsk/socket-internal.c index a4185daef..737bf8e29 100644 --- a/viosock/wsk/socket-internal.c +++ b/viosock/wsk/socket-internal.c @@ -92,6 +92,25 @@ _VioSocketCreate( } +static +NTSTATUS +_ConfigIrpComplete( + PDEVICE_OBJECT DeviceObject, + PIRP Irp, + PVOID Context +) +{ + NTSTATUS status = STATUS_UNSUCCESSFUL; + PKEVENT event = (PKEVENT)Context; + DEBUG_ENTER_FUNCTION("DeviceObject=0x%p; Irp=0x%p; Context=0x%p", DeviceObject, Irp, Context); + + KeSetEvent(event, IO_NO_INCREMENT, FALSE); + status = STATUS_MORE_PROCESSING_REQUIRED; + + DEBUG_EXIT_FUNCTION("0x%x", status); + return status; +} + _Must_inspect_result_ NTSTATUS @@ -116,7 +135,6 @@ VioWskSocketInternal( VIRTIO_VSOCK_CONFIG SocketConfig; PIRP ConfigIrp = NULL; KEVENT ConfigEvent; - IO_STATUS_BLOCK ConfigIosb; DEBUG_ENTER_FUNCTION("Client=0x%p; ListenSocket=0x%p; Flags=0x%x; SocketContext=0x%p; Dispatch=0x%p; OwningProcess=0x%p; OwningThread=0x%p; SecurityDescriptor=0x%p; pNewSocket=0x%p", Client, ListenSocket, Flags, SocketContext, Dispatch, OwningProcess, OwningThread, SecurityDescriptor, pNewSocket); PAGED_CODE(); @@ -191,33 +209,38 @@ VioWskSocketInternal( goto FreeSocket; KeInitializeEvent(&ConfigEvent, NotificationEvent, FALSE); - ConfigIrp = IoBuildDeviceIoControlRequest(IOCTL_GET_CONFIG, pContext->VIOSockDevice, NULL, 0, &SocketConfig, sizeof(SocketConfig), FALSE, &ConfigEvent, &ConfigIosb); + IoInitializeRemoveLock(&pSocket->CloseRemoveLock, VIOSOCK_WSK_MEMORY_TAG, 0, 0x7FFFFFFF); + ConfigIrp = IoAllocateIrp(1, FALSE); if (!ConfigIrp) { Status = STATUS_INSUFFICIENT_RESOURCES; goto CloseSocket; } - ObReferenceObject(pSocket->FileObject); - ConfigIrp->Tail.Overlay.OriginalFileObject = pSocket->FileObject; - IoGetNextIrpStackLocation(ConfigIrp)->FileObject = pSocket->FileObject; - Status = IoCallDriver(pContext->VIOSockDevice, ConfigIrp); - if (Status == STATUS_PENDING) { + IoSetCompletionRoutine(ConfigIrp, _ConfigIrpComplete, &ConfigEvent, TRUE, TRUE, TRUE); + Status = VioWskIrpAcquire(pSocket, ConfigIrp); + if (!NT_SUCCESS(Status)) + goto FreeConfigIrp; + + Status = VioWskSocketIOCTL(pSocket, IOCTL_GET_CONFIG, NULL, 0, &SocketConfig, sizeof(SocketConfig), ConfigIrp, NULL); + if (Status == STATUS_PENDING) + { KeWaitForSingleObject(&ConfigEvent, Executive, KernelMode, FALSE, NULL); - Status = ConfigIosb.Status; + Status = ConfigIrp->IoStatus.Status; } - - if (NT_SUCCESS(Status) && ConfigIosb.Information < sizeof(SocketConfig)) + + if (NT_SUCCESS(Status) && ConfigIrp->IoStatus.Information < sizeof(SocketConfig)) Status = STATUS_INVALID_PARAMETER; if (!NT_SUCCESS(Status)) - goto CloseSocket; + goto FreeConfigIrp; pSocket->GuestId = SocketConfig.guest_cid; - IoInitializeRemoveLock(&pSocket->CloseRemoveLock, VIOSOCK_WSK_MEMORY_TAG, 0, 0x7FFFFFFF); *pNewSocket = pSocket; pSocket = NULL; +FreeConfigIrp: + IoFreeIrp(ConfigIrp); CloseSocket: if (pSocket) { diff --git a/viosock/wsk/wsk-completion.c b/viosock/wsk/wsk-completion.c new file mode 100644 index 000000000..5010e8929 --- /dev/null +++ b/viosock/wsk/wsk-completion.c @@ -0,0 +1,232 @@ +/* + * Provider NPI functions + * + * Copyright (c) 2021 Virtuozzo International GmbH + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions + * are met : + * 1. Redistributions of source code must retain the above copyright + * notice, this list of conditions and the following disclaimer. + * 2. Redistributions in binary form must reproduce the above copyright + * notice, this list of conditions and the following disclaimer in the + * documentation and / or other materials provided with the distribution. + * 3. Neither the names of the copyright holders nor the names of their contributors + * may be used to endorse or promote products derived from this software + * without specific prior written permission. + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS ``AS IS'' AND + * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + * ARE DISCLAIMED.IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS + * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) + * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT + * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY + * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF + * SUCH DAMAGE. + */ + +#include "precomp.h" +#include "..\inc\debug-utils.h" +#include "wsk-utils.h" +#include "viowsk-internal.h" +#include "wsk-completion.h" + + + + +static +void +WskCancelIrp( + _In_ PDEVICE_OBJECT DeviceObject, + _In_ PIRP Irp +) +{ + PIRP CurrentIrp = NULL; + DEBUG_ENTER_FUNCTION("DeviceObject=0x%p; Irp=0x%p", DeviceObject, Irp); + + UNREFERENCED_PARAMETER(DeviceObject); + + IoReleaseCancelSpinLock(Irp->CancelIrql); + CurrentIrp = InterlockedExchangePointer(Irp->Tail.Overlay.DriverContext + 1, NULL); + if (CurrentIrp) + IoCancelIrp(CurrentIrp); + + DEBUG_EXIT_FUNCTION_VOID(); + return; +} + + +static +NTSTATUS +WskGeneralIrpCompletion( + _In_ PDEVICE_OBJECT DeviceObject, + _In_ PIRP Irp, + _In_ PVOID Context +) +{ + EWSKState opState; + PVIOSOCKET_COMPLETION_CONTEXT Ctx = (PVIOSOCKET_COMPLETION_CONTEXT)Context; + DEBUG_ENTER_FUNCTION("DeviceObject=0x%p; Irp=0x%p; Context=0x%p", DeviceObject, Irp, Context); + + if (Ctx->MasterIrp) + { + InterlockedExchangePointer(Ctx->MasterIrp->Tail.Overlay.DriverContext + 1, NULL); + if (Ctx->MasterIrp->Cancel) + Irp->IoStatus.Status = STATUS_CANCELLED; + } + + UNREFERENCED_PARAMETER(DeviceObject); + opState = Ctx->State; + if (NT_SUCCESS(Irp->IoStatus.Status)) { + switch (opState) + { + case wsksSingleIOCTL: + memcpy(Irp->UserBuffer, Irp->AssociatedIrp.SystemBuffer, Irp->IoStatus.Information); + opState = wsksFinished; + break; + default: + opState = wsksFinished; + break; + } + } + + IO_STATUS_BLOCK irpStatus; + + irpStatus = Irp->IoStatus; + VioWskIrpFree(Irp, Ctx->DeviceObject, TRUE); + if (!NT_SUCCESS(irpStatus.Status) || + opState == wsksFinished) { + if (Ctx->IoStatusBlock) + *Ctx->IoStatusBlock = irpStatus; + + if (Ctx->MasterIrp) + { + if (!Ctx->UseIOSBInformation) + Ctx->IOSBInformation = irpStatus.Information; + + VioWskIrpComplete(Ctx->Socket, Ctx->MasterIrp, irpStatus.Status, Ctx->IOSBInformation); + } + } + + WskCompContextDereference(Ctx); + + DEBUG_EXIT_FUNCTION("0x%x", STATUS_MORE_PROCESSING_REQUIRED); + return STATUS_MORE_PROCESSING_REQUIRED; +} + + +static +void +WskCompContextFree( + PVIOSOCKET_COMPLETION_CONTEXT CompContext +) +{ + DEBUG_ENTER_FUNCTION("CompContext=0x%p", CompContext); + + ExFreePoolWithTag(CompContext, VIOSOCK_WSK_MEMORY_TAG); + + DEBUG_EXIT_FUNCTION_VOID(); + return; +} + + +PVIOSOCKET_COMPLETION_CONTEXT +WskCompContextAlloc( + _In_ EWSKState State, + _In_ PVIOWSK_SOCKET Socket, + _In_opt_ PIRP MasterIrp, + _In_opt_ PIO_STATUS_BLOCK IoStatusBlock +) +{ + PVIOWSK_REG_CONTEXT pContext = NULL; + PWSK_REGISTRATION Registration = NULL; + PVIOSOCKET_COMPLETION_CONTEXT Ret = NULL; + DEBUG_ENTER_FUNCTION("State=%u; Socket=0x%p; MasterIrp=0x%p; IoStatusBlock=0x%p", State, Socket, MasterIrp, IoStatusBlock); + + Ret = ExAllocatePoolUninitialized(NonPagedPool, sizeof(*Ret), VIOSOCK_WSK_MEMORY_TAG); + if (!Ret) + goto Exit; + + memset(Ret, 0, sizeof(*Ret)); + InterlockedExchange(&Ret->ReferenceCount, 1); + Ret->State = State; + Ret->Socket = Socket; + Ret->MasterIrp = MasterIrp; + Ret->IoStatusBlock = IoStatusBlock; + Registration = (PWSK_REGISTRATION)Socket->Client; + pContext = (PVIOWSK_REG_CONTEXT)Registration->ReservedRegistrationContext; + Ret->DeviceObject = pContext->VIOSockDevice; + +Exit: + DEBUG_EXIT_FUNCTION("0x%p", Ret); + return Ret; +} + +void +WskCompContextReference( + _Inout_ PVIOSOCKET_COMPLETION_CONTEXT CompContext +) +{ + DEBUG_ENTER_FUNCTION("CompContext=0x%p", CompContext); + + InterlockedIncrement(&CompContext->ReferenceCount); + + DEBUG_EXIT_FUNCTION_VOID(); + return; +} + +void +WskCompContextDereference( + _Inout_ PVIOSOCKET_COMPLETION_CONTEXT CompContext +) +{ + DEBUG_ENTER_FUNCTION("CompContext=0x%p", CompContext); + + if (InterlockedDecrement(&CompContext->ReferenceCount) == 0) + WskCompContextFree(CompContext); + + DEBUG_EXIT_FUNCTION_VOID(); + return; +} + + +NTSTATUS +WskCompContextSendIrp( + _Inout_ PVIOSOCKET_COMPLETION_CONTEXT CompContext, + _In_ PIRP Irp +) +{ + NTSTATUS Status = STATUS_UNSUCCESSFUL; + PVIOWSK_REG_CONTEXT pContext = NULL; + PWSK_REGISTRATION Registration = NULL; + DEBUG_ENTER_FUNCTION("CompContext=0x%p; Irp=0x%p", CompContext, Irp); + + Status = STATUS_SUCCESS; + Registration = (PWSK_REGISTRATION)CompContext->Socket->Client; + pContext = (PVIOWSK_REG_CONTEXT)Registration->ReservedRegistrationContext; + if (_viowskDeviceObject) + { + Status = IoSetCompletionRoutineEx(_viowskDeviceObject, Irp, WskGeneralIrpCompletion, CompContext, TRUE, TRUE, TRUE); + if (!NT_SUCCESS(Status)) + goto CompleteMasterIrp; + } else IoSetCompletionRoutine(Irp, WskGeneralIrpCompletion, CompContext, TRUE, TRUE, TRUE); + + WskCompContextReference(CompContext); + if (CompContext->MasterIrp) + { + InterlockedExchangePointer(CompContext->MasterIrp->Tail.Overlay.DriverContext + 1, Irp); + IoSetCancelRoutine(CompContext->MasterIrp, WskCancelIrp); + } + + IoCallDriver(pContext->VIOSockDevice, Irp); + Status = STATUS_PENDING; + Irp = NULL; +CompleteMasterIrp: + if (Irp && CompContext->MasterIrp) + VioWskIrpComplete(CompContext->Socket, CompContext->MasterIrp, Status, 0); + + DEBUG_EXIT_FUNCTION("0x%x", Status); + return Status; +} diff --git a/viosock/wsk/wsk-completion.h b/viosock/wsk/wsk-completion.h new file mode 100644 index 000000000..54a33473c --- /dev/null +++ b/viosock/wsk/wsk-completion.h @@ -0,0 +1,84 @@ +/* + * Provider NPI functions + * + * Copyright (c) 2021 Virtuozzo International GmbH + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions + * are met : + * 1. Redistributions of source code must retain the above copyright + * notice, this list of conditions and the following disclaimer. + * 2. Redistributions in binary form must reproduce the above copyright + * notice, this list of conditions and the following disclaimer in the + * documentation and / or other materials provided with the distribution. + * 3. Neither the names of the copyright holders nor the names of their contributors + * may be used to endorse or promote products derived from this software + * without specific prior written permission. + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS ``AS IS'' AND + * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + * ARE DISCLAIMED.IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS + * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) + * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT + * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY + * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF + * SUCH DAMAGE. + */ + +#ifndef __WSK_COMPLETION_H__ +#define __WSK_COMPLETION_H__ + + + +#include "..\inc\vio_wsk.h" + + + +typedef enum _EWSKState { + wsksUndefined, + wsksSingleIOCTL, + wsksFinished, +} EWSKState, * PEWSKState; + +typedef struct _VIOSOCKET_COMPLETION_CONTEXT { + volatile LONG ReferenceCount; + PVIOWSK_SOCKET Socket; + PDEVICE_OBJECT DeviceObject; + EWSKState State; + PIRP MasterIrp; + PIO_STATUS_BLOCK IoStatusBlock; + ULONG_PTR IOSBInformation; + int UseIOSBInformation : 1; +} VIOSOCKET_COMPLETION_CONTEXT, * PVIOSOCKET_COMPLETION_CONTEXT; + + + +PVIOSOCKET_COMPLETION_CONTEXT +WskCompContextAlloc( + _In_ EWSKState State, + _In_ PVIOWSK_SOCKET Socket, + _In_opt_ PIRP MasterIrp, + _In_opt_ PIO_STATUS_BLOCK IoStatusBlock +); + +void +WskCompContextReference( + _Inout_ PVIOSOCKET_COMPLETION_CONTEXT CompContext +); + +void +WskCompContextDereference( + _Inout_ PVIOSOCKET_COMPLETION_CONTEXT CompContext +); + +NTSTATUS +WskCompContextSendIrp( + _Inout_ PVIOSOCKET_COMPLETION_CONTEXT CompContext, + _In_ PIRP Irp +); + + + +#endif diff --git a/viosock/wsk/wsk-utils.c b/viosock/wsk/wsk-utils.c index 7b00e5ea4..717a5e7b7 100644 --- a/viosock/wsk/wsk-utils.c +++ b/viosock/wsk/wsk-utils.c @@ -32,6 +32,7 @@ #include "..\inc\debug-utils.h" #include "..\inc\vio_wsk.h" #include "viowsk-internal.h" +#include "wsk-completion.h" #include "wsk-utils.h" #ifdef ALLOC_PRAGMA @@ -112,6 +113,41 @@ VioWskIrpComplete( } +void +VioWskIrpFree( + _Inout_ PIRP Irp, + _In_opt_ PDEVICE_OBJECT DeviceObject, + _In_ BOOLEAN Completion +) +{ + PDEVICE_OBJECT targetDevice = NULL; + DEBUG_ENTER_FUNCTION("Irp=0x%p; DeviceObject=0x%p; Completion=%u", Irp, DeviceObject, Completion); + + targetDevice = (Completion) ? DeviceObject : IoGetNextIrpStackLocation(Irp)->DeviceObject; + if (Irp->MdlAddress) + { + if ((targetDevice->Flags & DO_DIRECT_IO) == DO_DIRECT_IO) + MmUnlockPages(Irp->MdlAddress); + + IoFreeMdl(Irp->MdlAddress); + Irp->MdlAddress = NULL; + } + + if ((Irp->Flags & IRP_BUFFERED_IO) != 0 && + (Irp->Flags & IRP_DEALLOCATE_BUFFER) != 0) + { + Irp->Flags &= ~(IRP_BUFFERED_IO | IRP_DEALLOCATE_BUFFER); + ExFreePoolWithTag(Irp->AssociatedIrp.SystemBuffer, VIOSOCK_WSK_MEMORY_TAG); + Irp->AssociatedIrp.SystemBuffer = NULL; + } + + IoFreeIrp(Irp); + + DEBUG_EXIT_FUNCTION_VOID(); + return; +} + + _Must_inspect_result_ NTSTATUS VioWskAddressPartToString( @@ -165,3 +201,192 @@ VioWskStringToAddressPart( DEBUG_EXIT_FUNCTION("0x%x, *Value=%u", Status, *Value); return Status; } + + + +_Must_inspect_result_ +NTSTATUS +VioWskSocketIOCTL( + _In_ PVIOWSK_SOCKET Socket, + _In_ ULONG ControlCode, + _In_opt_ PVOID InputBuffer, + _In_ ULONG InputBufferLength, + _Out_opt_ PVOID OutputBuffer, + _In_ ULONG OutputBufferLength, + _Inout_opt_ PIRP Irp, + _Out_opt_ PIO_STATUS_BLOCK IoStatusBlock +) +{ + PIRP IOCTLIrp = NULL; + PVIOWSK_REG_CONTEXT pContext = NULL; + PWSK_REGISTRATION Registraction = NULL; + NTSTATUS Status = STATUS_UNSUCCESSFUL; + PVIOSOCKET_COMPLETION_CONTEXT CompContext = NULL; + DEBUG_ENTER_FUNCTION("Socket=0x%p; ControlCode=0x%x; InputBuffer=0x%p; InputBufferLength=%u; OutputBuffer=0x%p; OutputBufferLength=%u; Irp=0x%p; IoStatusBLock=0x%p", Socket, ControlCode, InputBuffer, InputBufferLength, OutputBuffer, OutputBufferLength, Irp, IoStatusBlock); + + Registraction = (PWSK_REGISTRATION)Socket->Client; + pContext = (PVIOWSK_REG_CONTEXT)Registraction->ReservedRegistrationContext; + Status = VioWskSocketBuildIOCTL(Socket, ControlCode, InputBuffer, InputBufferLength, OutputBuffer, OutputBufferLength, &IOCTLIrp); + if (!NT_SUCCESS(Status)) + goto Complete; + + CompContext = WskCompContextAlloc(wsksSingleIOCTL, Socket, Irp, IoStatusBlock); + if (CompContext == NULL) + { + Status = STATUS_INSUFFICIENT_RESOURCES; + goto FreeIOCTL; + } + + Status = WskCompContextSendIrp(CompContext, IOCTLIrp); + WskCompContextDereference(CompContext); + if (NT_SUCCESS(Status)) + IOCTLIrp = NULL; + + Irp = NULL; + +FreeIOCTL: + if (IOCTLIrp) + VioWskIrpFree(IOCTLIrp, NULL, FALSE); +Complete: + if (Irp) + VioWskIrpComplete(Socket, Irp, Status, 0); + + DEBUG_EXIT_FUNCTION("0x%x", Status); + return Status; +} + + +_Must_inspect_result_ +NTSTATUS +VioWskSocketBuildIOCTL( + _In_ PVIOWSK_SOCKET Socket, + _In_ ULONG ControlCode, + _In_opt_ PVOID InputBuffer, + _In_ ULONG InputBufferLength, + _In_opt_ PVOID OutputBuffer, + _In_ ULONG OutputBufferLength, + _Out_ PIRP *Irp +) +{ + PIRP IOCTLIrp = NULL; + ULONG method = 0; + NTSTATUS Status = STATUS_UNSUCCESSFUL; + PIO_STACK_LOCATION IrpStack = NULL; + PVIOWSK_REG_CONTEXT pContext = NULL; + PWSK_REGISTRATION Registraction = NULL; + DEBUG_ENTER_FUNCTION("Socket=0x%p; ControlCode=0x%x; InputBuffer=0x%p; InputBufferLength=%u; OutputBuffer=0x%p; OutputBufferLength=%u; Irp=0x%p", Socket, ControlCode, InputBuffer, InputBufferLength, OutputBuffer, OutputBufferLength, Irp); + + Status = STATUS_SUCCESS; + Registraction = (PWSK_REGISTRATION)Socket->Client; + pContext = (PVIOWSK_REG_CONTEXT)Registraction->ReservedRegistrationContext; + IOCTLIrp = IoAllocateIrp(pContext->VIOSockDevice->StackSize, FALSE); + if (!IOCTLIrp) + { + Status = STATUS_INSUFFICIENT_RESOURCES; + goto Exit; + } + + IOCTLIrp->Tail.Overlay.Thread = PsGetCurrentThread(); + IrpStack = IoGetNextIrpStackLocation(IOCTLIrp); + IrpStack->MajorFunction = IRP_MJ_DEVICE_CONTROL; + IrpStack->DeviceObject = pContext->VIOSockDevice; + IrpStack->FileObject = Socket->FileObject; + IrpStack->Parameters.DeviceIoControl.InputBufferLength = InputBufferLength; + IrpStack->Parameters.DeviceIoControl.OutputBufferLength = OutputBufferLength; + IrpStack->Parameters.DeviceIoControl.IoControlCode = ControlCode; + method = ControlCode & 3; + switch (method) + { + case METHOD_BUFFERED: + if (InputBufferLength != 0 || OutputBufferLength != 0) + { + IOCTLIrp->AssociatedIrp.SystemBuffer = ExAllocatePoolUninitialized( + NonPagedPool, + InputBufferLength > OutputBufferLength ? InputBufferLength : OutputBufferLength, + VIOSOCK_WSK_MEMORY_TAG); + + if (IOCTLIrp->AssociatedIrp.SystemBuffer) + { + if (InputBuffer) + memcpy(IOCTLIrp->AssociatedIrp.SystemBuffer, InputBuffer, InputBufferLength); + + IOCTLIrp->Flags = IRP_BUFFERED_IO | IRP_DEALLOCATE_BUFFER; + IOCTLIrp->UserBuffer = OutputBuffer; + if (OutputBuffer) + IOCTLIrp->Flags |= IRP_INPUT_OPERATION; + } + else Status = STATUS_INSUFFICIENT_RESOURCES; + } + else { + IOCTLIrp->Flags = 0; + IOCTLIrp->UserBuffer = NULL; + } + break; + case METHOD_IN_DIRECT: + case METHOD_OUT_DIRECT: + if (InputBuffer) + { + IOCTLIrp->AssociatedIrp.SystemBuffer = ExAllocatePoolUninitialized( + NonPagedPool, + InputBufferLength, + VIOSOCK_WSK_MEMORY_TAG); + + if (IOCTLIrp->AssociatedIrp.SystemBuffer) + { + memcpy(IOCTLIrp->AssociatedIrp.SystemBuffer, InputBuffer, InputBufferLength); + IOCTLIrp->Flags = IRP_BUFFERED_IO | IRP_DEALLOCATE_BUFFER; + } + else Status = STATUS_INSUFFICIENT_RESOURCES; + } + else IOCTLIrp->Flags = 0; + + if (NT_SUCCESS(Status) && OutputBuffer) { + IOCTLIrp->MdlAddress = IoAllocateMdl( + OutputBuffer, + OutputBufferLength, + FALSE, + FALSE, + NULL); + + if (!IOCTLIrp->MdlAddress) + { + if (InputBuffer) + ExFreePoolWithTag(IOCTLIrp->AssociatedIrp.SystemBuffer, VIOSOCK_WSK_MEMORY_TAG); + + Status = STATUS_INSUFFICIENT_RESOURCES; + } + + if (NT_SUCCESS(Status)) { + __try + { + MmProbeAndLockPages(IOCTLIrp->MdlAddress, KernelMode, (LOCK_OPERATION)((method == METHOD_IN_DIRECT) ? IoReadAccess : IoWriteAccess)); + } + __except (EXCEPTION_EXECUTE_HANDLER) { + if (IOCTLIrp->MdlAddress) + IoFreeMdl(IOCTLIrp->MdlAddress); + + if (InputBuffer) + ExFreePoolWithTag(IOCTLIrp->AssociatedIrp.SystemBuffer, VIOSOCK_WSK_MEMORY_TAG); + + Status = GetExceptionCode(); + } + } + } + break; + case METHOD_NEITHER: + IOCTLIrp->UserBuffer = OutputBuffer; + IrpStack->Parameters.DeviceIoControl.Type3InputBuffer = InputBuffer; + break; + } + + if (NT_SUCCESS(Status)) { + *Irp = IOCTLIrp; + IOCTLIrp = NULL; + } + + if (IOCTLIrp) + VioWskIrpFree(IOCTLIrp, NULL, FALSE); +Exit: + DEBUG_EXIT_FUNCTION("0x%x, *Irp=0x%p", Status, *Irp); + return Status; +} diff --git a/viosock/wsk/wsk-utils.h b/viosock/wsk/wsk-utils.h index 0ecdc8253..6092db929 100644 --- a/viosock/wsk/wsk-utils.h +++ b/viosock/wsk/wsk-utils.h @@ -60,6 +60,13 @@ VioWskIrpComplete( _In_ ULONG_PTR Information ); +void +VioWskIrpFree( + _Inout_ PIRP Irp, + _In_opt_ PDEVICE_OBJECT DeviceObject, + _In_ BOOLEAN Completion +); + _Must_inspect_result_ NTSTATUS VioWskAddressPartToString( @@ -74,5 +81,30 @@ VioWskStringToAddressPart( _Out_ PULONG Value ); +_Must_inspect_result_ +NTSTATUS +VioWskSocketIOCTL( + _In_ PVIOWSK_SOCKET Socket, + _In_ ULONG ControlCode, + _In_opt_ PVOID InputBuffer, + _In_ ULONG InputBufferLength, + _Out_opt_ PVOID OutputBuffer, + _In_ ULONG OutputBufferLength, + _Inout_opt_ PIRP Irp, + _Out_opt_ PIO_STATUS_BLOCK IoStatusBlock +); + +_Must_inspect_result_ +NTSTATUS +VioWskSocketBuildIOCTL( + _In_ PVIOWSK_SOCKET Socket, + _In_ ULONG ControlCode, + _In_opt_ PVOID InputBuffer, + _In_ ULONG InputBufferLength, + _In_opt_ PVOID OutputBuffer, + _In_ ULONG OutputBufferLength, + _Out_ PIRP* Irp +); + #endif diff --git a/viosock/wsk/wsk.vcxproj b/viosock/wsk/wsk.vcxproj index b99fb6430..463fc1814 100644 --- a/viosock/wsk/wsk.vcxproj +++ b/viosock/wsk/wsk.vcxproj @@ -31,6 +31,7 @@ + @@ -39,6 +40,7 @@ + diff --git a/viosock/wsk/wsk.vcxproj.filters b/viosock/wsk/wsk.vcxproj.filters index 7c6c3a32a..a2e1b70c1 100644 --- a/viosock/wsk/wsk.vcxproj.filters +++ b/viosock/wsk/wsk.vcxproj.filters @@ -33,6 +33,9 @@ Header Files + + Header Files + @@ -53,5 +56,8 @@ Source Files + + Source Files + \ No newline at end of file From 0a1af6feefdb60f65a3c69f3dfb6a4bc8a5a31ee Mon Sep 17 00:00:00 2001 From: Martin Drab Date: Sun, 24 Jul 2022 21:27:59 +0200 Subject: [PATCH 07/18] WSK: Implement WskGetLocalAddress The library implements this call by forwarding it to the VIOSOCK driver as an IOCTL_SOCKET_GET_SOCK_NAME request. ref: [WskGetLocalAddress](https://learn.microsoft.com/en-us/windows-hardware/drivers/ddi/wsk/nc-wsk-pfn_wsk_get_local_address) Signed-off-by: Martin Drab --- viosock/wsk/socket.c | 22 +++++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/viosock/wsk/socket.c b/viosock/wsk/socket.c index 7aafea06d..0c77b4dc0 100644 --- a/viosock/wsk/socket.c +++ b/viosock/wsk/socket.c @@ -32,6 +32,7 @@ #include "viowsk.h" #include "wsk-utils.h" #include "viowsk-internal.h" +#include "wsk-completion.h" #include "wsk-workitem.h" #include "..\inc\vio_wsk.h" @@ -383,10 +384,25 @@ VioWskGetLocalAddress( _Inout_ PIRP Irp ) { - UNREFERENCED_PARAMETER(Socket); - UNREFERENCED_PARAMETER(LocalAddress); + NTSTATUS Status = STATUS_UNSUCCESSFUL; + PVIOWSK_SOCKET pSocket = CONTAINING_RECORD(Socket, VIOWSK_SOCKET, WskSocket); + DEBUG_ENTER_FUNCTION("Socket=0x%p; LocalAddress=0x%p; Irp=0x%p", Socket, LocalAddress, Irp); - return VioWskCompleteIrp(Irp, STATUS_NOT_IMPLEMENTED, 0); + Status = VioWskIrpAcquire(pSocket, Irp); + if (!NT_SUCCESS(Status)) + { + pSocket = NULL; + goto CompleteIrp; + } + + Status = VioWskSocketIOCTL(pSocket, IOCTL_SOCKET_GET_SOCK_NAME, NULL, 0, LocalAddress, sizeof(SOCKADDR_VM), Irp, NULL); + Irp = NULL; +CompleteIrp: + if (Irp) + VioWskIrpComplete(pSocket, Irp, Status, 0); + + DEBUG_EXIT_FUNCTION("0x%x", Status); + return Status; } NTSTATUS From ad273da9886493879c35f11335f916f0dd9ba608 Mon Sep 17 00:00:00 2001 From: Martin Drab Date: Sun, 24 Jul 2022 21:31:19 +0200 Subject: [PATCH 08/18] WSK: Implement WskGetRemoteAddress The library implements this call by forwarding it to the VIOSOCK driver as an IOCTL_SOCKET_GET_PEER_NAME request. ref: [WskGetRemoteAddress](https://learn.microsoft.com/en-us/windows-hardware/drivers/ddi/wsk/nc-wsk-pfn_wsk_get_remote_address) Signed-off-by: Martin Drab --- viosock/wsk/socket.c | 21 ++++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/viosock/wsk/socket.c b/viosock/wsk/socket.c index 0c77b4dc0..db3b39d7a 100644 --- a/viosock/wsk/socket.c +++ b/viosock/wsk/socket.c @@ -430,10 +430,25 @@ VioWskGetRemoteAddress( _Inout_ PIRP Irp ) { - UNREFERENCED_PARAMETER(Socket); - UNREFERENCED_PARAMETER(RemoteAddress); + NTSTATUS Status = STATUS_UNSUCCESSFUL; + PVIOWSK_SOCKET pSocket = CONTAINING_RECORD(Socket, VIOWSK_SOCKET, WskSocket); + DEBUG_ENTER_FUNCTION("Socket=0x%p; RemoteAddress=0x%p; Irp=0x%p", Socket, RemoteAddress, Irp); - return VioWskCompleteIrp(Irp, STATUS_NOT_IMPLEMENTED, 0); + Status = VioWskIrpAcquire(pSocket, Irp); + if (!NT_SUCCESS(Status)) + { + pSocket = NULL; + goto CompleteIrp; + } + + Status = VioWskSocketIOCTL(pSocket, IOCTL_SOCKET_GET_PEER_NAME, NULL, 0, RemoteAddress, sizeof(SOCKADDR_VM), Irp, NULL); + Irp = NULL; +CompleteIrp: + if (Irp) + VioWskIrpComplete(pSocket, Irp, Status, 0); + + DEBUG_EXIT_FUNCTION("0x%x", Status); + return Status; } NTSTATUS From 8ce5cc79c33c0d7d0bd6ef1a672b261e332826e7 Mon Sep 17 00:00:00 2001 From: Martin Drab Date: Sun, 24 Jul 2022 21:33:45 +0200 Subject: [PATCH 09/18] WSK: Implement WskConnect Unlike the standard WSK interface (TCP/IP etc.), AF_VSOCK does not require to call WskBind before WskConnect. This also holds for the connect() call in usermode. The client driver can use VMADDR_CID_ANY to connect to loopback. ref: https://www.bing.com/search?q=wskconnect&cvid=6a21f40a81d8444a886dbae1cbb54e09&aqs=edge..69i57j0l8.1546j0j1&pglt=2083&FORM=ANSPA1&PC=U531 Signed-off-by: Martin Drab --- viosock/wsk/socket.c | 27 ++++++++++++++++++++++++--- 1 file changed, 24 insertions(+), 3 deletions(-) diff --git a/viosock/wsk/socket.c b/viosock/wsk/socket.c index db3b39d7a..27377a3e5 100644 --- a/viosock/wsk/socket.c +++ b/viosock/wsk/socket.c @@ -414,11 +414,32 @@ VioWskConnect( _Inout_ PIRP Irp ) { - UNREFERENCED_PARAMETER(Socket); - UNREFERENCED_PARAMETER(RemoteAddress); + NTSTATUS Status = STATUS_UNSUCCESSFUL; + SOCKADDR_VM VMRemoteAddr; + PVIOWSK_SOCKET pSocket = CONTAINING_RECORD(Socket, VIOWSK_SOCKET, WskSocket); + DEBUG_ENTER_FUNCTION("Socket=0x%p; RemoteAddress=0x%p; Flags=0x%x; Irp=0x%p", Socket, RemoteAddress, Flags, Irp); + UNREFERENCED_PARAMETER(Flags); - return VioWskCompleteIrp(Irp, STATUS_NOT_IMPLEMENTED, 0); + VMRemoteAddr = *(PSOCKADDR_VM)RemoteAddress; + if (VMRemoteAddr.svm_cid == VMADDR_CID_ANY) + VMRemoteAddr.svm_cid = pSocket->GuestId; + + Status = VioWskIrpAcquire(pSocket, Irp); + if (!NT_SUCCESS(Status)) + { + pSocket = NULL; + goto CompleteIrp; + } + + Status = VioWskSocketIOCTL(pSocket, IOCTL_SOCKET_CONNECT, &VMRemoteAddr, sizeof(VMRemoteAddr), NULL, 0, Irp, NULL); + Irp = NULL; +CompleteIrp: + if (Irp) + VioWskIrpComplete(pSocket, Irp, Status, 0); + + DEBUG_EXIT_FUNCTION("0x%x", Status); + return Status; } From 98132e61323f641c3faa4f48be4d0a94d210e805 Mon Sep 17 00:00:00 2001 From: Martin Drab Date: Sun, 24 Jul 2022 22:08:29 +0200 Subject: [PATCH 10/18] WSK: Implement WskControlSocket The implementation supports both synchronous (Irp == NULL) and asynchronous (Irp != NULL) modes of operation and can be used to get/set socket options and send IOCTLs. ref: [WskControlSocket](https://learn.microsoft.com/en-us/windows-hardware/drivers/ddi/wsk/nc-wsk-pfn_wsk_control_socket) Signed-off-by: Martin Drab --- viosock/wsk/socket.c | 168 +++++++++++++++++++++++++++++------ viosock/wsk/wsk-completion.c | 6 ++ viosock/wsk/wsk-completion.h | 2 + 3 files changed, 149 insertions(+), 27 deletions(-) diff --git a/viosock/wsk/socket.c b/viosock/wsk/socket.c index 27377a3e5..165fd0a80 100644 --- a/viosock/wsk/socket.c +++ b/viosock/wsk/socket.c @@ -39,16 +39,16 @@ NTSTATUS WSKAPI VioWskControlSocket( - _In_ PWSK_SOCKET Socket, - _In_ WSK_CONTROL_SOCKET_TYPE RequestType, - _In_ ULONG ControlCode, - _In_ ULONG Level, - _In_ SIZE_T InputSize, + _In_ PWSK_SOCKET Socket, + _In_ WSK_CONTROL_SOCKET_TYPE RequestType, + _In_ ULONG ControlCode, + _In_ ULONG Level, + _In_ SIZE_T InputSize, _In_reads_bytes_opt_(InputSize) PVOID InputBuffer, - _In_ SIZE_T OutputSize, - _Out_writes_bytes_opt_(OutputSize) PVOID OutputBuffer, - _Out_opt_ SIZE_T *OutputSizeReturned, - _Inout_opt_ PIRP Irp + _In_ SIZE_T OutputSize, + _Out_writes_bytes_opt_(OutputSize) PVOID OutputBuffer, + _Out_opt_ SIZE_T *OutputSizeReturned, + _Inout_opt_ PIRP Irp ); NTSTATUS @@ -250,32 +250,146 @@ WSK_PROVIDER_STREAM_DISPATCH gStreamDispatch = #endif // if (NTDDI_VERSION >= NTDDI_WIN10_RS2) ////////////////////////////////////////////////////////////////////////// + +static +NTSTATUS +_WskControlSocketCompletion( + _In_ PDEVICE_OBJECT DeviceObject, + _In_ PIRP Irp, + _In_ PVOID Context +) +{ + PKEVENT Event = (PKEVENT)Context; + NTSTATUS Status = STATUS_MORE_PROCESSING_REQUIRED; + DEBUG_ENTER_FUNCTION("DeviceObject=0x%p; Irp=0x%p; Context=0x%p", DeviceObject, Irp, Context); + + UNREFERENCED_PARAMETER(DeviceObject); + UNREFERENCED_PARAMETER(Irp); + + KeSetEvent(Event, IO_NO_INCREMENT, FALSE); + + DEBUG_EXIT_FUNCTION("0x%x", Status); + return Status; +} + + NTSTATUS WSKAPI VioWskControlSocket( - _In_ PWSK_SOCKET Socket, - _In_ WSK_CONTROL_SOCKET_TYPE RequestType, - _In_ ULONG ControlCode, - _In_ ULONG Level, - _In_ SIZE_T InputSize, - _In_reads_bytes_opt_(InputSize) PVOID InputBuffer, - _In_ SIZE_T OutputSize, + _In_ PWSK_SOCKET Socket, + _In_ WSK_CONTROL_SOCKET_TYPE RequestType, + _In_ ULONG ControlCode, + _In_ ULONG Level, + _In_ SIZE_T InputSize, + _In_reads_bytes_opt_(InputSize) PVOID InputBuffer, + _In_ SIZE_T OutputSize, _Out_writes_bytes_opt_(OutputSize) PVOID OutputBuffer, - _Out_opt_ SIZE_T *OutputSizeReturned, - _Inout_opt_ PIRP Irp + _Out_opt_ SIZE_T *OutputSizeReturned, + _Inout_opt_ PIRP Irp ) { - UNREFERENCED_PARAMETER(Socket); - UNREFERENCED_PARAMETER(RequestType); - UNREFERENCED_PARAMETER(ControlCode); - UNREFERENCED_PARAMETER(Level); - UNREFERENCED_PARAMETER(InputSize); - UNREFERENCED_PARAMETER(InputBuffer); - UNREFERENCED_PARAMETER(OutputSize); - UNREFERENCED_PARAMETER(OutputBuffer); + KEVENT Event; + PIRP IOCTLIrp = NULL; + PVIOSOCKET_COMPLETION_CONTEXT CompContext = NULL; + NTSTATUS Status = STATUS_UNSUCCESSFUL; + PVIOWSK_SOCKET pSocket = CONTAINING_RECORD(Socket, VIOWSK_SOCKET, WskSocket); + DEBUG_ENTER_FUNCTION("Socket=0x%p; RequestType=%u; ControlCode=0x%x; Level=%u; InputSize=%zu; InputBuffer=0x%p; OutputSize=%zu; OutputBuffer=0x%p; OutputSizeReturned=0x%p; Irp=0x%p", Socket, RequestType, ControlCode, Level, InputSize, InputBuffer, OutputSize, OutputBuffer, OutputSizeReturned, Irp); + UNREFERENCED_PARAMETER(OutputSizeReturned); - return VioWskCompleteIrp(Irp, STATUS_NOT_IMPLEMENTED, 0); + if (!Irp) + { + Irp = IoAllocateIrp(1, FALSE); + if (!Irp) + { + Status = STATUS_INSUFFICIENT_RESOURCES; + goto Exit; + } + + KeInitializeEvent(&Event, NotificationEvent, FALSE); + IoSetCompletionRoutine(Irp, _WskControlSocketCompletion, &Event, TRUE, TRUE, TRUE); + Status = VioWskControlSocket(Socket, RequestType, ControlCode, Level, InputSize, InputBuffer, OutputSize, OutputBuffer, OutputSizeReturned, Irp); + if (Status == STATUS_PENDING) + KeWaitForSingleObject(&Event, Executive, KernelMode, FALSE, NULL); + + if (OutputSizeReturned) + *OutputSizeReturned = Irp->IoStatus.Information; + + IoFreeIrp(Irp); + goto Exit; + } + + Status = VioWskIrpAcquire(pSocket, Irp); + if (!NT_SUCCESS(Status)) + { + pSocket = NULL; + goto CompleteIrp; + } + + switch (RequestType) + { + case WskSetOption: + case WskGetOption: { + ULONG ioctl = 0; + VIRTIO_VSOCK_OPT Opt; + + memset(&Opt, 0, sizeof(Opt)); + Opt.level = Level; + Opt.optname = ControlCode; + switch (RequestType) + { + case WskSetOption: + ioctl = IOCTL_SOCKET_SET_SOCK_OPT; + Opt.optval = (ULONGLONG)InputBuffer; + Opt.optlen = (int)InputSize; + break; + case WskGetOption: + ioctl = IOCTL_SOCKET_GET_SOCK_OPT; + Opt.optval = (ULONGLONG)OutputBuffer; + Opt.optlen = (int)OutputSize; + break; + } + + Status = VioWskSocketBuildIOCTL(pSocket, ioctl, &Opt, sizeof(Opt), &Opt, sizeof(Opt), &IOCTLIrp); + } break; + case WskIoctl: { + VIRTIO_VSOCK_IOCTL_IN params; + + params.dwIoControlCode = ControlCode; + params.lpvInBuffer = (ULONGLONG)InputBuffer; + params.cbInBuffer = (ULONG)InputSize; + Status = VioWskSocketBuildIOCTL(pSocket, IOCTL_SOCKET_IOCTL, ¶ms, sizeof(params), OutputBuffer, (ULONG)OutputSize, &IOCTLIrp); + } break; + default: + Status = STATUS_INVALID_PARAMETER; + break; + } + + if (!NT_SUCCESS(Status)) + goto CompleteIrp; + + CompContext = WskCompContextAlloc((RequestType == WskIoctl ? wsksSingleIOCTL : wsksFinished), pSocket, Irp, NULL); + if (!CompContext) + { + Status = STATUS_INSUFFICIENT_RESOURCES; + goto FreeIOCTLIrp; + } + + Status = WskCompContextSendIrp(CompContext, IOCTLIrp); + if (NT_SUCCESS(Status)) + IOCTLIrp = NULL; + + WskCompContextDereference(CompContext); + Irp = NULL; +FreeIOCTLIrp: + if (IOCTLIrp) + VioWskIrpFree(IOCTLIrp, NULL, FALSE); +CompleteIrp: + if (Irp) + VioWskIrpComplete(pSocket, Irp, Status, 0); +Exit: + DEBUG_EXIT_FUNCTION("0x%x", Status); + return Status; } NTSTATUS diff --git a/viosock/wsk/wsk-completion.c b/viosock/wsk/wsk-completion.c index 5010e8929..8deca4c69 100644 --- a/viosock/wsk/wsk-completion.c +++ b/viosock/wsk/wsk-completion.c @@ -101,6 +101,12 @@ WskGeneralIrpCompletion( if (Ctx->IoStatusBlock) *Ctx->IoStatusBlock = irpStatus; + if (Ctx->BytesReturned) + *Ctx->BytesReturned = irpStatus.Information; + + if (Ctx->Event) + KeSetEvent(Ctx->Event, IO_NO_INCREMENT, FALSE); + if (Ctx->MasterIrp) { if (!Ctx->UseIOSBInformation) diff --git a/viosock/wsk/wsk-completion.h b/viosock/wsk/wsk-completion.h index 54a33473c..7e110f60e 100644 --- a/viosock/wsk/wsk-completion.h +++ b/viosock/wsk/wsk-completion.h @@ -49,6 +49,8 @@ typedef struct _VIOSOCKET_COMPLETION_CONTEXT { EWSKState State; PIRP MasterIrp; PIO_STATUS_BLOCK IoStatusBlock; + PSIZE_T BytesReturned; + PKEVENT Event; ULONG_PTR IOSBInformation; int UseIOSBInformation : 1; } VIOSOCKET_COMPLETION_CONTEXT, * PVIOSOCKET_COMPLETION_CONTEXT; From e01292780bf0be7b705af97bc08dd5bb2889fa64 Mon Sep 17 00:00:00 2001 From: Martin Drab Date: Sun, 24 Jul 2022 14:50:34 -0700 Subject: [PATCH 11/18] WSK: Implement VIoWskSend Routines for building IRP_MJ_WRITE and IRP_MJ_WRITE IRPs are also par of this commit since the VIOSOCK driver implements send/recv operations through these requests. It is also necessary to deal with the fact that the WSK interface transfers send/recv buffers through WSK_BUF structures, very similar to WSABUF in usermode, that support MDL chaining. Since the VIOSOCK driver does not support MDL chaining, the WskSend initiates a sequence of IRP_MJ_WRITE requests, one for each MDL in the chain. ref: [WskSend](https://learn.microsoft.com/en-us/windows-hardware/drivers/ddi/wsk/nc-wsk-pfn_wsk_send_to) Signed-off-by: Martin Drab --- viosock/wsk/socket.c | 22 ++++- viosock/wsk/wsk-completion.c | 26 ++++++ viosock/wsk/wsk-completion.h | 9 ++ viosock/wsk/wsk-utils.c | 175 +++++++++++++++++++++++++++++++++++ viosock/wsk/wsk-utils.h | 30 ++++++ 5 files changed, 259 insertions(+), 3 deletions(-) diff --git a/viosock/wsk/socket.c b/viosock/wsk/socket.c index 165fd0a80..8f91071bf 100644 --- a/viosock/wsk/socket.c +++ b/viosock/wsk/socket.c @@ -595,11 +595,27 @@ VioWskSend( _Inout_ PIRP Irp ) { - UNREFERENCED_PARAMETER(Socket); - UNREFERENCED_PARAMETER(Buffer); + NTSTATUS Status = STATUS_UNSUCCESSFUL; + PVIOWSK_SOCKET pSocket = CONTAINING_RECORD(Socket, VIOWSK_SOCKET, WskSocket); + DEBUG_ENTER_FUNCTION("Socket=0x%p; Buffer=0x%p; Flags=0x%x; Irp=0x%p", Socket, Buffer, Flags, Irp); + UNREFERENCED_PARAMETER(Flags); - return VioWskCompleteIrp(Irp, STATUS_NOT_IMPLEMENTED, 0); + Status = VioWskIrpAcquire(pSocket, Irp); + if (!NT_SUCCESS(Status)) { + pSocket = NULL; + goto CompleteIrp; + } + + Status = VioWskSocketReadWrite(pSocket, Buffer, IRP_MJ_WRITE, Irp); + Irp = NULL; + +CompleteIrp: + if (Irp) + VioWskIrpComplete(pSocket, Irp, Status, 0); + + DEBUG_EXIT_FUNCTION("0x%x", Status); + return Status; } NTSTATUS diff --git a/viosock/wsk/wsk-completion.c b/viosock/wsk/wsk-completion.c index 8deca4c69..0a578f64a 100644 --- a/viosock/wsk/wsk-completion.c +++ b/viosock/wsk/wsk-completion.c @@ -66,6 +66,8 @@ WskGeneralIrpCompletion( _In_ PVOID Context ) { + PIRP NextIrp = NULL; + NTSTATUS NextIrpStatus = STATUS_UNSUCCESSFUL; EWSKState opState; PVIOSOCKET_COMPLETION_CONTEXT Ctx = (PVIOSOCKET_COMPLETION_CONTEXT)Context; DEBUG_ENTER_FUNCTION("DeviceObject=0x%p; Irp=0x%p; Context=0x%p", DeviceObject, Irp, Context); @@ -86,6 +88,30 @@ WskGeneralIrpCompletion( memcpy(Irp->UserBuffer, Irp->AssociatedIrp.SystemBuffer, Irp->IoStatus.Information); opState = wsksFinished; break; + case wsksSend: + if (Ctx->Specific.Transfer.NextMdl && + (Irp->IoStatus.Information == Ctx->Specific.Transfer.CurrentMdlSize)) + { + PMDL NextMdl = Ctx->Specific.Transfer.NextMdl; + + Ctx->Specific.Transfer.CurrentMdlSize = NextMdl->Next ? MmGetMdlByteCount(NextMdl) : Ctx->Specific.Transfer.LastMdlSize; + Irp->IoStatus.Status = VioWskSocketBuildReadWriteSingleMdl(Ctx->Socket, NextMdl, 0, Ctx->Specific.Transfer.CurrentMdlSize, IRP_MJ_WRITE, &NextIrp); + if (!NT_SUCCESS(Irp->IoStatus.Status)) + break; + + Ctx->Specific.Transfer.NextMdl = NextMdl->Next; + NextIrpStatus = WskCompContextSendIrp(Ctx, NextIrp); + if (!NT_SUCCESS(NextIrpStatus)) { + Irp->IoStatus.Status = NextIrpStatus; + Ctx->MasterIrp = NULL; + VioWskIrpFree(NextIrp, DeviceObject, FALSE); + } + } + else opState = wsksFinished; + + Ctx->IOSBInformation += Irp->IoStatus.Information; + Ctx->UseIOSBInformation = 1; + break; default: opState = wsksFinished; break; diff --git a/viosock/wsk/wsk-completion.h b/viosock/wsk/wsk-completion.h index 7e110f60e..66553904a 100644 --- a/viosock/wsk/wsk-completion.h +++ b/viosock/wsk/wsk-completion.h @@ -39,6 +39,8 @@ typedef enum _EWSKState { wsksUndefined, wsksSingleIOCTL, + wsksSend, + wsksReceive, wsksFinished, } EWSKState, * PEWSKState; @@ -53,6 +55,13 @@ typedef struct _VIOSOCKET_COMPLETION_CONTEXT { PKEVENT Event; ULONG_PTR IOSBInformation; int UseIOSBInformation : 1; + union { + struct { + PMDL NextMdl; + ULONG CurrentMdlSize; + ULONG LastMdlSize; + } Transfer; + } Specific; } VIOSOCKET_COMPLETION_CONTEXT, * PVIOSOCKET_COMPLETION_CONTEXT; diff --git a/viosock/wsk/wsk-utils.c b/viosock/wsk/wsk-utils.c index 717a5e7b7..7d42deff2 100644 --- a/viosock/wsk/wsk-utils.c +++ b/viosock/wsk/wsk-utils.c @@ -390,3 +390,178 @@ VioWskSocketBuildIOCTL( DEBUG_EXIT_FUNCTION("0x%x, *Irp=0x%p", Status, *Irp); return Status; } + + +_Must_inspect_result_ +NTSTATUS +VioWskSocketReadWrite( + _In_ PVIOWSK_SOCKET Socket, + const WSK_BUF *Buffers, + _In_ UCHAR MajorFunction, + _Inout_ PIRP Irp +) +{ + PIRP OpIrp = NULL; + ULONG firstMdlLength = 0; + ULONG lastMdlLength = 0; + EWSKState state = wsksUndefined; + PVIOWSK_REG_CONTEXT pContext = NULL; + NTSTATUS Status = STATUS_UNSUCCESSFUL; + PWSK_REGISTRATION Registraction = NULL; + PVIOSOCKET_COMPLETION_CONTEXT CompContext = NULL; + DEBUG_ENTER_FUNCTION("Sockect=0x%p; Buffers=0x%p; MajorFunction=%u; Irp=0x%p", Socket, Buffers, MajorFunction, Irp); + + Registraction = (PWSK_REGISTRATION)Socket->Client; + pContext = (PVIOWSK_REG_CONTEXT)Registraction->ReservedRegistrationContext; + Status = WskBufferValidate(Buffers, &firstMdlLength, &lastMdlLength); + if (!NT_SUCCESS(Status)) + goto CompleteParentIrp; + + switch (MajorFunction) + { + case IRP_MJ_READ: + state = wsksReceive; + break; + case IRP_MJ_WRITE: + state = wsksSend; + break; + default: + Status = STATUS_INVALID_PARAMETER_3; + goto CompleteParentIrp; + break; + } + + Status = VioWskSocketBuildReadWriteSingleMdl(Socket, Buffers->Mdl, Buffers->Offset, firstMdlLength, MajorFunction, &OpIrp); + if (!NT_SUCCESS(Status)) + goto CompleteParentIrp; + + CompContext = WskCompContextAlloc(state, Socket, Irp, NULL); + if (!CompContext) + goto FreeOpIrp; + + CompContext->Specific.Transfer.CurrentMdlSize = firstMdlLength; + CompContext->Specific.Transfer.LastMdlSize = lastMdlLength; + CompContext->Specific.Transfer.NextMdl = Buffers->Mdl->Next; + Status = WskCompContextSendIrp(CompContext, OpIrp); + WskCompContextDereference(CompContext); + if (NT_SUCCESS(Status)) + OpIrp = NULL; + + Irp = NULL; + +FreeOpIrp: + if (OpIrp) + VioWskIrpFree(OpIrp, NULL, FALSE); +CompleteParentIrp: + if (Irp) + VioWskIrpComplete(Socket, Irp, Status, 0); + + DEBUG_EXIT_FUNCTION("0x%x", Status); + return Status; +} + + +_Must_inspect_result_ +NTSTATUS +VioWskSocketBuildReadWriteSingleMdl( + _In_ PVIOWSK_SOCKET Socket, + _In_ PMDL Mdl, + _In_ ULONG Offset, + _In_ ULONG Length, + _In_ UCHAR MajorFunction, + _Out_ PIRP *Irp +) +{ + PIRP OpIrp = NULL; + PVOID mdlBuffer = NULL; + LARGE_INTEGER StartingOffset = { 0 }; + PIO_STACK_LOCATION IrpStack = NULL; + PVIOWSK_REG_CONTEXT pContext = NULL; + PWSK_REGISTRATION Registraction = NULL; + NTSTATUS Status = STATUS_UNSUCCESSFUL; + DEBUG_ENTER_FUNCTION("Socket=0x%p; Mdl=0x%p; MajorFunction=%u; Offset=%u; Length=%u; Irp=0x%p", Socket, Mdl, Offset, Length, MajorFunction, Irp); + + Registraction = (PWSK_REGISTRATION)Socket->Client; + pContext = (PVIOWSK_REG_CONTEXT)Registraction->ReservedRegistrationContext; + mdlBuffer = MmGetSystemAddressForMdlSafe(Mdl, NormalPagePriority); + if (!mdlBuffer) + { + Status = STATUS_INSUFFICIENT_RESOURCES; + goto Exit; + } + + OpIrp = IoBuildAsynchronousFsdRequest(MajorFunction, pContext->VIOSockDevice, (PUCHAR)mdlBuffer + Offset, Length, &StartingOffset, NULL); + if (!OpIrp) + { + Status = STATUS_INSUFFICIENT_RESOURCES; + goto Exit; + } + + IrpStack = IoGetNextIrpStackLocation(OpIrp); + IrpStack->DeviceObject = pContext->VIOSockDevice; + IrpStack->FileObject = Socket->FileObject; + *Irp = OpIrp; + Status = STATUS_SUCCESS; + +Exit: + DEBUG_EXIT_FUNCTION("0x%x, *Irp=0x%p", Status, *Irp); + return Status; +} + + +NTSTATUS +WskBufferValidate( + _In_ const WSK_BUF* Buffer, + _Out_ PULONG FirstMdlLength, + _Out_ PULONG LastMdlLength +) +{ + PMDL mdl = NULL; + ULONG offset = 0; + SIZE_T length = 0; + ULONG mdlLength = 0; + NTSTATUS status = STATUS_UNSUCCESSFUL; + DEBUG_ENTER_FUNCTION("Buffer=0x%p; FirstMdlLength=0x%p; LastMdlLength=0x%p", Buffer, FirstMdlLength, LastMdlLength); + + *FirstMdlLength = 0; + *LastMdlLength = 0; + status = STATUS_SUCCESS; + length = Buffer->Length; + offset = Buffer->Offset; + mdl = Buffer->Mdl; + if (mdl != NULL) + { + mdlLength = MmGetMdlByteCount(mdl); + if (offset <= mdlLength) + { + while (TRUE) + { + ULONG effectiveLength = mdlLength - offset; + + if (length < effectiveLength) + effectiveLength = (ULONG)length; + + if (mdl == Buffer->Mdl) + *FirstMdlLength = effectiveLength; + + mdl = mdl->Next; + length -= effectiveLength; + if (length == 0 || mdl == NULL) + { + *LastMdlLength = effectiveLength; + break; + } + + mdlLength = MmGetMdlByteCount(mdl); + offset = 0; + } + } + else status = STATUS_INVALID_PARAMETER; + } + else if (length != 0) + status = STATUS_INVALID_PARAMETER; + + DEBUG_EXIT_FUNCTION("0x%x, *FirstMdlLength=%u, *LastMdlLength=%u", status, *FirstMdlLength, *LastMdlLength); + return status; +} + diff --git a/viosock/wsk/wsk-utils.h b/viosock/wsk/wsk-utils.h index 6092db929..9b7642332 100644 --- a/viosock/wsk/wsk-utils.h +++ b/viosock/wsk/wsk-utils.h @@ -107,4 +107,34 @@ VioWskSocketBuildIOCTL( ); +_Must_inspect_result_ +NTSTATUS +VioWskSocketReadWrite( + _In_ PVIOWSK_SOCKET Socket, + const WSK_BUF *Buffers, + _In_ UCHAR MajorFunction, + _Inout_ PIRP Irp +); + + +_Must_inspect_result_ +NTSTATUS +VioWskSocketBuildReadWriteSingleMdl( + _In_ PVIOWSK_SOCKET Socket, + _In_ PMDL Mdl, + _In_ ULONG Offset, + _In_ ULONG Length, + _In_ UCHAR MajorFunction, + _Out_ PIRP* Irp +); + + +NTSTATUS +WskBufferValidate( + _In_ const WSK_BUF* Buffer, + _Out_ PULONG FirstMdlLength, + _Out_ PULONG LastMdlLength +); + + #endif From 4b9d4a95b910f1a3de0d3f949a88bceafeb8e8a9 Mon Sep 17 00:00:00 2001 From: Martin Drab Date: Sun, 24 Jul 2022 15:01:03 -0700 Subject: [PATCH 12/18] WSK: Impelemnt WskReceive Since the VIOSOCK driver does not support MDL chaining, one call to WskReceive fills at most one MDL of the WSK_BUF structure. ref: [WskReceive](https://learn.microsoft.com/en-us/windows-hardware/drivers/ddi/wsk/nc-wsk-pfn_wsk_receive) Signed-off-by: Martin Drab --- viosock/wsk/socket.c | 28 ++++++++++++++++++++++++---- viosock/wsk/wsk-completion.c | 3 +++ 2 files changed, 27 insertions(+), 4 deletions(-) diff --git a/viosock/wsk/socket.c b/viosock/wsk/socket.c index 8f91071bf..3b5505b31 100644 --- a/viosock/wsk/socket.c +++ b/viosock/wsk/socket.c @@ -627,11 +627,31 @@ VioWskReceive( _Inout_ PIRP Irp ) { - UNREFERENCED_PARAMETER(Socket); - UNREFERENCED_PARAMETER(Buffer); - UNREFERENCED_PARAMETER(Flags); + PVIOWSK_SOCKET pSocket = NULL; + NTSTATUS Status = STATUS_UNSUCCESSFUL; + DEBUG_ENTER_FUNCTION("Socket=0x%p; Buffer=0x%p; Flags=0x%x; Irp=0x%p", Socket, Buffer, Flags, Irp); - return VioWskCompleteIrp(Irp, STATUS_NOT_IMPLEMENTED, 0); + if (Flags != 0) { + Status = STATUS_NOT_SUPPORTED; + goto CompleteIrp; + } + + pSocket = CONTAINING_RECORD(Socket, VIOWSK_SOCKET, WskSocket); + Status = VioWskIrpAcquire(pSocket, Irp); + if (!NT_SUCCESS(Status)) { + pSocket = NULL; + goto CompleteIrp; + } + + Status = VioWskSocketReadWrite(pSocket, Buffer, IRP_MJ_READ, Irp); + Irp = NULL; + +CompleteIrp: + if (Irp) + VioWskIrpComplete(pSocket, Irp, Status, 0); + + DEBUG_EXIT_FUNCTION("0x%x", Status); + return Status; } NTSTATUS diff --git a/viosock/wsk/wsk-completion.c b/viosock/wsk/wsk-completion.c index 0a578f64a..64d57a240 100644 --- a/viosock/wsk/wsk-completion.c +++ b/viosock/wsk/wsk-completion.c @@ -88,6 +88,9 @@ WskGeneralIrpCompletion( memcpy(Irp->UserBuffer, Irp->AssociatedIrp.SystemBuffer, Irp->IoStatus.Information); opState = wsksFinished; break; + case wsksReceive: + opState = wsksFinished; + break; case wsksSend: if (Ctx->Specific.Transfer.NextMdl && (Irp->IoStatus.Information == Ctx->Specific.Transfer.CurrentMdlSize)) From be7bac310ff5de6ae2253b348bfb96fa9f062713 Mon Sep 17 00:00:00 2001 From: Martin Drab Date: Mon, 25 Jul 2022 00:47:51 -0700 Subject: [PATCH 13/18] WSK: Implement WskDisconnect Only the zero value for the Flags argument is supported. The client driver may use the Buffer argument to send data before disconnecting from the remote end. ref: [WskDisconnect](https://learn.microsoft.com/en-us/windows-hardware/drivers/ddi/wsk/nc-wsk-pfn_wsk_disconnect) Signed-off-by: Martin Drab --- viosock/wsk/socket.c | 66 +++++++++++++++++++++++++++++++++--- viosock/wsk/wsk-completion.c | 19 +++++++++++ viosock/wsk/wsk-completion.h | 2 ++ 3 files changed, 83 insertions(+), 4 deletions(-) diff --git a/viosock/wsk/socket.c b/viosock/wsk/socket.c index 3b5505b31..b1801d451 100644 --- a/viosock/wsk/socket.c +++ b/viosock/wsk/socket.c @@ -663,11 +663,69 @@ VioWskDisconnect( _Inout_ PIRP Irp ) { - UNREFERENCED_PARAMETER(Socket); - UNREFERENCED_PARAMETER(Buffer); - UNREFERENCED_PARAMETER(Flags); + PIRP SendIrp = NULL; + ULONG How = 2; // SD_BOTH + ULONG firstMdlLength = 0; + ULONG lastMdlLength = 0; + NTSTATUS Status = STATUS_UNSUCCESSFUL; + PVIOSOCKET_COMPLETION_CONTEXT CompContext = NULL; + PVIOWSK_SOCKET pSocket = CONTAINING_RECORD(Socket, VIOWSK_SOCKET, WskSocket); + DEBUG_ENTER_FUNCTION("Socket=0x%p; Buffer=0x%p; Flags=0x%x; Irp=0x%p", Socket, Buffer, Flags, Irp); - return VioWskCompleteIrp(Irp, STATUS_NOT_IMPLEMENTED, 0); + if (Flags != 0) + { + Status = STATUS_NOT_SUPPORTED; + pSocket = NULL; + goto CompleteIrp; + } + + Status = VioWskIrpAcquire(pSocket, Irp); + if (!NT_SUCCESS(Status)) + { + pSocket = NULL; + goto CompleteIrp; + } + + if (!Buffer || !Buffer->Mdl || Buffer->Length == 0 || (Flags & WSK_FLAG_ABORTIVE)) + { + Status = VioWskSocketIOCTL(pSocket, IOCTL_SOCKET_SHUTDOWN, &How, sizeof(How), NULL, 0, Irp, NULL); + Irp = NULL; + goto CompleteIrp; + } + + Status = WskBufferValidate(Buffer, &firstMdlLength, &lastMdlLength); + if (!NT_SUCCESS(Status)) + goto CompleteIrp; + + Status = VioWskSocketBuildReadWriteSingleMdl(pSocket, Buffer->Mdl, Buffer->Offset, firstMdlLength, IRP_MJ_WRITE, &SendIrp); + if (!NT_SUCCESS(Status)) + goto CompleteIrp; + + CompContext = WskCompContextAlloc(wsksDisconnect, pSocket, Irp, NULL); + if (!CompContext) { + Status = STATUS_INSUFFICIENT_RESOURCES; + goto FreeSendIrp; + } + + CompContext->Specific.Transfer.CurrentMdlSize = firstMdlLength; + CompContext->Specific.Transfer.LastMdlSize = lastMdlLength; + CompContext->Specific.Transfer.NextMdl = Buffer->Mdl->Next; + Status = WskCompContextSendIrp(CompContext, SendIrp); + WskCompContextDereference(CompContext); + if (NT_SUCCESS(Status)) + SendIrp = NULL; + + Irp = NULL; + +FreeSendIrp: + if (SendIrp) + VioWskIrpFree(SendIrp, NULL, FALSE); +CompleteIrp: + if (Irp) + VioWskIrpComplete(pSocket, Irp, Status, 0); + + DEBUG_EXIT_FUNCTION("0x%x", Status); + return Status; } NTSTATUS diff --git a/viosock/wsk/wsk-completion.c b/viosock/wsk/wsk-completion.c index 64d57a240..3440be77a 100644 --- a/viosock/wsk/wsk-completion.c +++ b/viosock/wsk/wsk-completion.c @@ -88,10 +88,14 @@ WskGeneralIrpCompletion( memcpy(Irp->UserBuffer, Irp->AssociatedIrp.SystemBuffer, Irp->IoStatus.Information); opState = wsksFinished; break; + case wsksDisconnected: + opState = wsksFinished; + break; case wsksReceive: opState = wsksFinished; break; case wsksSend: + case wsksDisconnect: if (Ctx->Specific.Transfer.NextMdl && (Irp->IoStatus.Information == Ctx->Specific.Transfer.CurrentMdlSize)) { @@ -110,6 +114,21 @@ WskGeneralIrpCompletion( VioWskIrpFree(NextIrp, DeviceObject, FALSE); } } + else if (opState == wsksDisconnect) { + ULONG How = 2; // SD_BOTH + + Irp->IoStatus.Status = VioWskSocketBuildIOCTL(Ctx->Socket, IOCTL_SOCKET_SHUTDOWN, &How, sizeof(How), NULL, 0, &NextIrp); + if (!NT_SUCCESS(Irp->IoStatus.Status)) + break; + + Ctx->State = wsksDisconnected; + NextIrpStatus = WskCompContextSendIrp(Ctx, NextIrp); + if (!NT_SUCCESS(NextIrpStatus)) { + Irp->IoStatus.Status = NextIrpStatus; + Ctx->MasterIrp = NULL; + VioWskIrpFree(NextIrp, NULL, FALSE); + } + } else opState = wsksFinished; Ctx->IOSBInformation += Irp->IoStatus.Information; diff --git a/viosock/wsk/wsk-completion.h b/viosock/wsk/wsk-completion.h index 66553904a..ebe7fda99 100644 --- a/viosock/wsk/wsk-completion.h +++ b/viosock/wsk/wsk-completion.h @@ -41,6 +41,8 @@ typedef enum _EWSKState { wsksSingleIOCTL, wsksSend, wsksReceive, + wsksDisconnect, + wsksDisconnected, wsksFinished, } EWSKState, * PEWSKState; From 1ba61858488b0d8c6592f1e6d89f725b7f92c36c Mon Sep 17 00:00:00 2001 From: Martin Drab Date: Mon, 25 Jul 2022 03:24:25 -0700 Subject: [PATCH 14/18] WSK: Implement WskAccept If called at IRQL = PASSIVE_LEVEL, this call blocks until a connection is accepted. This is due to the fact that the VIOSOCK driver implements the accept() operation through an IRP_MJ_CREATE (ZwCreateFile) which is usually synchronous. To make WskAccept non-blocking, the client driver can call it at higher IRQL, the accept() call into the VIOSOCK driver is then deferred to a workitem. ref: https://learn.microsoft.com/en-us/windows-hardware/drivers/ddi/wsk/nc-wsk-pfn_wsk_accept Signed-off-by: Martin Drab --- viosock/wsk/socket.c | 121 +++++++++++++++++++++++++++++++++-- viosock/wsk/wsk-completion.c | 43 +++++++++++++ viosock/wsk/wsk-completion.h | 10 ++- viosock/wsk/wsk-workitem.c | 23 ++++++- viosock/wsk/wsk-workitem.h | 4 ++ 5 files changed, 192 insertions(+), 9 deletions(-) diff --git a/viosock/wsk/socket.c b/viosock/wsk/socket.c index b1801d451..ebb6b62e5 100644 --- a/viosock/wsk/socket.c +++ b/viosock/wsk/socket.c @@ -464,14 +464,121 @@ VioWskAccept( _Inout_ PIRP Irp ) { - UNREFERENCED_PARAMETER(ListenSocket); - UNREFERENCED_PARAMETER(Flags); - UNREFERENCED_PARAMETER(AcceptSocketContext); - UNREFERENCED_PARAMETER(AcceptSocketDispatch); - UNREFERENCED_PARAMETER(LocalAddress); - UNREFERENCED_PARAMETER(RemoteAddress); + PIRP AddrIrp = NULL; + PIRP CloseIrp = NULL; + PWSK_WORKITEM CloseWorkItem = NULL; + BOOLEAN acceptSocketAcquired = FALSE; + PWSK_WORKITEM WorkItem = NULL; + PVIOSOCKET_COMPLETION_CONTEXT CompContext = NULL; + PVIOWSK_SOCKET pSocket = NULL; + NTSTATUS Status = STATUS_UNSUCCESSFUL; + PVIOWSK_SOCKET pListenSocket = CONTAINING_RECORD(ListenSocket, VIOWSK_SOCKET, WskSocket); + DEBUG_ENTER_FUNCTION("ListenSocket=0x%p; Flags=0x%x; AcceptSocketContext=0x%p; AcceptSocketDispatch=0x%p; LocalAddress=0x%p; RemoteAddress=0x%p; Irp=0x%p", ListenSocket, Flags, AcceptSocketContext, AcceptSocketDispatch, LocalAddress, RemoteAddress, Irp); - return VioWskCompleteIrp(Irp, STATUS_NOT_IMPLEMENTED, 0); + Status = VioWskIrpAcquire(pListenSocket, Irp); + if (!NT_SUCCESS(Status)) + { + pListenSocket = NULL; + goto CompleteIrp; + } + + if (KeGetCurrentIrql() > PASSIVE_LEVEL) + { + WorkItem = WskWorkItemAlloc(wskwitAccept, Irp); + if (!WorkItem) + { + Status = STATUS_INSUFFICIENT_RESOURCES; + goto CompleteIrp; + } + + WorkItem->Specific.Accept.AcceptSocketContext = AcceptSocketContext; + WorkItem->Specific.Accept.AcceptSocketDispatch = AcceptSocketDispatch; + WorkItem->Specific.Accept.Flags = Flags; + WorkItem->Specific.Accept.ListenSocket = ListenSocket; + WorkItem->Specific.Accept.LocalAddress = LocalAddress; + WorkItem->Specific.Accept.RemoteAddress = RemoteAddress; + WskWorkItemQueue(WorkItem); + Status = STATUS_PENDING; + goto Exit; + } + + CloseIrp = IoAllocateIrp(1, FALSE); + if (!CloseIrp) + { + Status = STATUS_INSUFFICIENT_RESOURCES; + goto CompleteIrp; + } + + CloseWorkItem = WskWorkItemAlloc(wskwitCloseSocket, CloseIrp); + if (!CloseWorkItem) + { + Status = STATUS_INSUFFICIENT_RESOURCES; + goto FreeCloseIrp; + } + + CloseIrp = NULL; + Status = VioWskSocketInternal(pListenSocket->Client, pListenSocket, Flags, AcceptSocketContext, AcceptSocketDispatch, NULL, NULL, NULL, &pSocket); + if (!NT_SUCCESS(Status)) + goto FreeCloseWorkItem; + + if (LocalAddress || RemoteAddress) + { + Status = VioWskIrpAcquire(pSocket, Irp); + if (!NT_SUCCESS(Status)) + goto CloseNewSocket; + + acceptSocketAcquired = TRUE; + Status = VioWskSocketBuildIOCTL(pSocket, (LocalAddress ? IOCTL_SOCKET_GET_SOCK_NAME : IOCTL_SOCKET_GET_PEER_NAME), NULL, 0, (LocalAddress ? LocalAddress : RemoteAddress), sizeof(SOCKADDR_VM), &AddrIrp); + if (!NT_SUCCESS(Status)) + { + Status = STATUS_INSUFFICIENT_RESOURCES; + goto CloseNewSocket; + } + + CloseWorkItem->Specific.CloseSocket.Socket = &pSocket->WskSocket; + CompContext = WskCompContextAlloc((LocalAddress ? wsksAcceptLocal : wsksAcceptRemote), pSocket, Irp, NULL); + if (!CompContext) + { + Status = STATUS_INSUFFICIENT_RESOURCES; + goto FreeAddrIrp; + } + + VioWskIrpRelease(pListenSocket, Irp); + Irp = NULL; + CompContext->Specific.Accept.LocalAddress = LocalAddress; + CompContext->Specific.Accept.RemoteAddress = RemoteAddress; + CompContext->Specific.Accept.Socket = &pSocket->WskSocket; + CompContext->Specific.Accept.CloseWorkItem = CloseWorkItem; + Status = WskCompContextSendIrp(CompContext, AddrIrp); + WskCompContextDereference(CompContext); + if (NT_SUCCESS(Status)) + { + CloseWorkItem = NULL; + AddrIrp = NULL; + } + } + +FreeAddrIrp: + if (AddrIrp) + VioWskIrpFree(AddrIrp, NULL, FALSE); +CloseNewSocket: + if (!NT_SUCCESS(Status)) + { + VioWskCloseSocketInternal(pSocket, (acceptSocketAcquired ? Irp : NULL)); + pSocket = NULL; + } +FreeCloseWorkItem: + if (CloseWorkItem) + WskWorkItemFree(CloseWorkItem); +FreeCloseIrp: + if (CloseIrp) + VioWskIrpFree(CloseIrp, NULL, FALSE); +CompleteIrp: + if (Irp) + VioWskIrpComplete(pListenSocket, Irp, Status, (ULONG_PTR)pSocket); +Exit: + DEBUG_EXIT_FUNCTION("0x%x", Status); + return Status; } NTSTATUS diff --git a/viosock/wsk/wsk-completion.c b/viosock/wsk/wsk-completion.c index 3440be77a..95ee456f3 100644 --- a/viosock/wsk/wsk-completion.c +++ b/viosock/wsk/wsk-completion.c @@ -88,6 +88,40 @@ WskGeneralIrpCompletion( memcpy(Irp->UserBuffer, Irp->AssociatedIrp.SystemBuffer, Irp->IoStatus.Information); opState = wsksFinished; break; + case wsksAcceptLocal: + memcpy(Ctx->Specific.Accept.LocalAddress, Irp->AssociatedIrp.SystemBuffer, sizeof(SOCKADDR_VM)); + if (Ctx->Specific.Accept.RemoteAddress) + { + Irp->IoStatus.Status = VioWskSocketBuildIOCTL(Ctx->Socket, IOCTL_SOCKET_GET_PEER_NAME, NULL, 0, Ctx->Specific.Accept.RemoteAddress, sizeof(SOCKADDR_VM), &NextIrp); + if (!NT_SUCCESS(Irp->IoStatus.Status)) + break; + + Ctx->State = wsksAcceptRemote; + NextIrpStatus = WskCompContextSendIrp(Ctx, NextIrp); + if (!NT_SUCCESS(NextIrpStatus)) + { + Irp->IoStatus.Status = NextIrpStatus; + Ctx->MasterIrp = NULL; + VioWskIrpFree(NextIrp, NULL, FALSE); + } + } + else + { + opState = wsksFinished; + WskWorkItemFree(Ctx->Specific.Accept.CloseWorkItem); + Ctx->Specific.Accept.CloseWorkItem = NULL; + Ctx->IOSBInformation = (ULONG_PTR)Ctx->Specific.Accept.Socket; + Ctx->UseIOSBInformation = 1; + } + break; + case wsksAcceptRemote: + memcpy(Ctx->Specific.Accept.RemoteAddress, Irp->AssociatedIrp.SystemBuffer, sizeof(SOCKADDR_VM)); + opState = wsksFinished; + WskWorkItemFree(Ctx->Specific.Accept.CloseWorkItem); + Ctx->Specific.Accept.CloseWorkItem = NULL; + Ctx->IOSBInformation = (ULONG_PTR)Ctx->Specific.Accept.Socket; + Ctx->UseIOSBInformation = 1; + break; case wsksDisconnected: opState = wsksFinished; break; @@ -146,6 +180,15 @@ WskGeneralIrpCompletion( VioWskIrpFree(Irp, Ctx->DeviceObject, TRUE); if (!NT_SUCCESS(irpStatus.Status) || opState == wsksFinished) { + if (Ctx->State == wsksAcceptLocal || + Ctx->State == wsksAcceptRemote) { + if (Ctx->Specific.Accept.CloseWorkItem) + { + WskWorkItemQueue(Ctx->Specific.Accept.CloseWorkItem); + Ctx->Specific.Accept.CloseWorkItem = NULL; + } + } + if (Ctx->IoStatusBlock) *Ctx->IoStatusBlock = irpStatus; diff --git a/viosock/wsk/wsk-completion.h b/viosock/wsk/wsk-completion.h index ebe7fda99..86a90f758 100644 --- a/viosock/wsk/wsk-completion.h +++ b/viosock/wsk/wsk-completion.h @@ -33,7 +33,7 @@ #include "..\inc\vio_wsk.h" - +#include "wsk-workitem.h" typedef enum _EWSKState { @@ -43,6 +43,8 @@ typedef enum _EWSKState { wsksReceive, wsksDisconnect, wsksDisconnected, + wsksAcceptLocal, + wsksAcceptRemote, wsksFinished, } EWSKState, * PEWSKState; @@ -58,6 +60,12 @@ typedef struct _VIOSOCKET_COMPLETION_CONTEXT { ULONG_PTR IOSBInformation; int UseIOSBInformation : 1; union { + struct { + PWSK_SOCKET Socket; + PSOCKADDR LocalAddress; + PSOCKADDR RemoteAddress; + PWSK_WORKITEM CloseWorkItem; + } Accept; struct { PMDL NextMdl; ULONG CurrentMdlSize; diff --git a/viosock/wsk/wsk-workitem.c b/viosock/wsk/wsk-workitem.c index 4c8f858cc..70223530a 100644 --- a/viosock/wsk/wsk-workitem.c +++ b/viosock/wsk/wsk-workitem.c @@ -30,6 +30,7 @@ #include "precomp.h" #include "..\inc\debug-utils.h" #include "viowsk.h" +#include "wsk-utils.h" #include "..\inc\vio_wsk.h" #include "viowsk-internal.h" #include "wsk-workitem.h" @@ -82,6 +83,10 @@ _ProcessWorkItem( WorkItem->Irp); break; case wskwitAccept: + { + PVIOWSK_SOCKET ListenSocket = NULL; + + ListenSocket = CONTAINING_RECORD(WorkItem->Specific.Accept.ListenSocket, VIOWSK_SOCKET, WskSocket); VioWskAccept( WorkItem->Specific.Accept.ListenSocket, WorkItem->Specific.Accept.Flags, @@ -90,7 +95,8 @@ _ProcessWorkItem( WorkItem->Specific.Accept.LocalAddress, WorkItem->Specific.Accept.RemoteAddress, WorkItem->Irp); - break; + VioWskIrpRelease(ListenSocket, WorkItem->Irp); + } break; default: ASSERT(FALSE); break; @@ -185,3 +191,18 @@ WskWorkItemQueue( DEBUG_EXIT_FUNCTION_VOID(); return; } + +void +WskWorkItemFree( + _In_ PWSK_WORKITEM WorkItem +) +{ + if (WorkItem->IoMethod) + IoUninitializeWorkItem((PIO_WORKITEM)WorkItem->IoWorkItem); + + IoFreeIrp(WorkItem->Irp); + ExFreePoolWithTag(WorkItem, VIOSOCK_WSK_MEMORY_TAG); + + DEBUG_EXIT_FUNCTION_VOID(); + return; +} diff --git a/viosock/wsk/wsk-workitem.h b/viosock/wsk/wsk-workitem.h index 3beaf945e..aa7c45aa0 100644 --- a/viosock/wsk/wsk-workitem.h +++ b/viosock/wsk/wsk-workitem.h @@ -105,6 +105,10 @@ WskWorkItemQueue( _In_ PWSK_WORKITEM WorkItem ); +void +WskWorkItemFree( + _In_ PWSK_WORKITEM WorkItem +); #endif From 750cedbc8803057db29fd6fe5714570985ba7122 Mon Sep 17 00:00:00 2001 From: Martin Drab Date: Mon, 25 Jul 2022 20:21:17 +0200 Subject: [PATCH 15/18] WSK: Implement WskBind This call translates to an IOCTL_SOCKET_BIND request to the VIOSOCK driver, followed with IOCTL_SOCKET_LISTEN one if the socket in question is a listening one. This needs to be done because the WskListen function was not part of the WSK interface until Windows 10 version 1703. Contrary to the TCP/IP WSK implementation in Windows, it is not necessary to call WskBind for non-listening sockets. ref: https://learn.microsoft.com/en-us/windows-hardware/drivers/ddi/wsk/nc-wsk-pfn_wsk_bind Signed-off-by: Martin Drab --- viosock/wsk/socket.c | 42 +++++++++++++++++++++++++++++++++--- viosock/wsk/wsk-completion.c | 21 ++++++++++++++++++ viosock/wsk/wsk-completion.h | 2 ++ 3 files changed, 62 insertions(+), 3 deletions(-) diff --git a/viosock/wsk/socket.c b/viosock/wsk/socket.c index ebb6b62e5..1023ca317 100644 --- a/viosock/wsk/socket.c +++ b/viosock/wsk/socket.c @@ -445,11 +445,47 @@ VioWskBind( _Inout_ PIRP Irp ) { - UNREFERENCED_PARAMETER(Socket); - UNREFERENCED_PARAMETER(LocalAddress); + PIRP BindIrp = NULL; + NTSTATUS Status = STATUS_UNSUCCESSFUL; + PVIOWSK_SOCKET pSocket = CONTAINING_RECORD(Socket, VIOWSK_SOCKET, WskSocket); + PVIOSOCKET_COMPLETION_CONTEXT CompContext = NULL; + DEBUG_ENTER_FUNCTION("Socket=0x%p; LocalAddress=0x%p; Flags=0x%x; Irp=0x%p", Socket, LocalAddress, Flags, Irp); + UNREFERENCED_PARAMETER(Flags); - return VioWskCompleteIrp(Irp, STATUS_NOT_IMPLEMENTED, 0); + Status = VioWskIrpAcquire(pSocket, Irp); + if (!NT_SUCCESS(Status)) + { + pSocket = NULL; + goto Complete; + } + + Status = VioWskSocketBuildIOCTL(pSocket, IOCTL_SOCKET_BIND, LocalAddress, sizeof(SOCKADDR_VM), NULL, 0, &BindIrp); + if (!NT_SUCCESS(Status)) + goto Complete; + + CompContext = WskCompContextAlloc(wsksBind, pSocket, Irp, NULL); + if (!CompContext) + { + Status = STATUS_INSUFFICIENT_RESOURCES; + goto FreeBindirp; + } + + Status = WskCompContextSendIrp(CompContext, BindIrp); + WskCompContextDereference(CompContext); + if (NT_SUCCESS(Status)) + BindIrp = NULL; + + Irp = NULL; +FreeBindirp: + if (BindIrp) + VioWskIrpFree(BindIrp, NULL, FALSE); +Complete: + if (Irp) + VioWskIrpComplete(pSocket, Irp, Status, 0); + + DEBUG_EXIT_FUNCTION("0x%x", Status); + return Status; } NTSTATUS diff --git a/viosock/wsk/wsk-completion.c b/viosock/wsk/wsk-completion.c index 95ee456f3..ce68c1f58 100644 --- a/viosock/wsk/wsk-completion.c +++ b/viosock/wsk/wsk-completion.c @@ -88,6 +88,26 @@ WskGeneralIrpCompletion( memcpy(Irp->UserBuffer, Irp->AssociatedIrp.SystemBuffer, Irp->IoStatus.Information); opState = wsksFinished; break; + case wsksBind: + if (Ctx->Socket->Type == WSK_FLAG_LISTEN_SOCKET) + { + ULONG SendLog = 128; + + Irp->IoStatus.Status = VioWskSocketBuildIOCTL(Ctx->Socket, IOCTL_SOCKET_LISTEN, &SendLog, sizeof(SendLog), NULL, 0, &NextIrp); + if (!NT_SUCCESS(Irp->IoStatus.Status)) + break; + + Ctx->State = wsksListen; + NextIrpStatus = WskCompContextSendIrp(Ctx, NextIrp); + if (!NT_SUCCESS(NextIrpStatus)) + { + Irp->IoStatus.Status = NextIrpStatus; + Ctx->MasterIrp = NULL; + VioWskIrpFree(NextIrp, NULL, FALSE); + } + } + else opState = wsksFinished; + break; case wsksAcceptLocal: memcpy(Ctx->Specific.Accept.LocalAddress, Irp->AssociatedIrp.SystemBuffer, sizeof(SOCKADDR_VM)); if (Ctx->Specific.Accept.RemoteAddress) @@ -122,6 +142,7 @@ WskGeneralIrpCompletion( Ctx->IOSBInformation = (ULONG_PTR)Ctx->Specific.Accept.Socket; Ctx->UseIOSBInformation = 1; break; + case wsksListen: case wsksDisconnected: opState = wsksFinished; break; diff --git a/viosock/wsk/wsk-completion.h b/viosock/wsk/wsk-completion.h index 86a90f758..ae9c42046 100644 --- a/viosock/wsk/wsk-completion.h +++ b/viosock/wsk/wsk-completion.h @@ -45,6 +45,8 @@ typedef enum _EWSKState { wsksDisconnected, wsksAcceptLocal, wsksAcceptRemote, + wsksBind, + wsksListen, wsksFinished, } EWSKState, * PEWSKState; From a569f19127a847bcae8cb6e79749da3321791a1e Mon Sep 17 00:00:00 2001 From: Martin Drab Date: Mon, 25 Jul 2022 11:48:05 -0700 Subject: [PATCH 16/18] WSK: Let unsupported calls correctly report STATUS_NOT_IMPLEMENTED The following routines are not implemented: - WskControlClient - WskInspectComplete - WskConnectEx - WskSendEd - WskReceiveEx ref: [WskInspectComplete](https://learn.microsoft.com/en-us/windows-hardware/drivers/ddi/wsk/nc-wsk-pfn_wsk_inspect_complete) ref: [WskConnectEx](https://learn.microsoft.com/en-us/windows-hardware/drivers/ddi/wsk/nc-wsk-pfn_wsk_connect_ex) Signed-off-by: Martin Drab --- viosock/wsk/provider.c | 12 +++++- viosock/wsk/socket.c | 93 ++++++++++++++++++++++++++++++++++++++---- 2 files changed, 95 insertions(+), 10 deletions(-) diff --git a/viosock/wsk/provider.c b/viosock/wsk/provider.c index 486552756..ccc61c7e6 100644 --- a/viosock/wsk/provider.c +++ b/viosock/wsk/provider.c @@ -150,6 +150,9 @@ VioWskControlClient( _Inout_opt_ PIRP Irp ) { + NTSTATUS Status = STATUS_UNSUCCESSFUL; + DEBUG_ENTER_FUNCTION("Client=0x%p; ControlCode=0x%x; InputSize=%zu; InputBuffer=0x%p; OutputSize=%zu; OutputBuffer=0x%p; OutputSizeReturned=0x%p; Irp=0x%p", Client, ControlCode, InputSize, InputBuffer, OutputSize, OutputBuffer, OutputSizeReturned, Irp); + UNREFERENCED_PARAMETER(Client); UNREFERENCED_PARAMETER(ControlCode); UNREFERENCED_PARAMETER(InputSize); @@ -158,7 +161,14 @@ VioWskControlClient( UNREFERENCED_PARAMETER(OutputBuffer); UNREFERENCED_PARAMETER(OutputSizeReturned); - return VioWskCompleteIrp(Irp, STATUS_NOT_IMPLEMENTED, 0); + Status = STATUS_NOT_IMPLEMENTED; + if (Irp) + { + VioWskIrpComplete(NULL, Irp, Status, 0); + } + + DEBUG_EXIT_FUNCTION("0x%x", Status); + return Status; } _At_(*Result, __drv_allocatesMem(Mem)) diff --git a/viosock/wsk/socket.c b/viosock/wsk/socket.c index 1023ca317..828378d1d 100644 --- a/viosock/wsk/socket.c +++ b/viosock/wsk/socket.c @@ -626,11 +626,27 @@ VioWskInspectComplete( _Inout_ PIRP Irp ) { + NTSTATUS Status = STATUS_UNSUCCESSFUL; + PVIOWSK_SOCKET pSocket = CONTAINING_RECORD(ListenSocket, VIOWSK_SOCKET, WskSocket); + DEBUG_ENTER_FUNCTION("ListenSocket=0x%p; InspectID=0x%p; Action=%u; Irp=0x%p", ListenSocket, InspectID, Action, Irp); + UNREFERENCED_PARAMETER(ListenSocket); UNREFERENCED_PARAMETER(InspectID); UNREFERENCED_PARAMETER(Action); - return VioWskCompleteIrp(Irp, STATUS_NOT_IMPLEMENTED, 0); + Status = VioWskIrpAcquire(pSocket, Irp); + if (!NT_SUCCESS(Status)) + { + pSocket = NULL; + goto CompleteIrp; + } + + Status = STATUS_NOT_IMPLEMENTED; +CompleteIrp: + VioWskIrpComplete(pSocket, Irp, Status, 0); + + DEBUG_EXIT_FUNCTION("0x%x", Status); + return Status; } NTSTATUS @@ -894,12 +910,27 @@ VioWskConnectEx( _Inout_ PIRP Irp ) { - UNREFERENCED_PARAMETER(Socket); + NTSTATUS Status = STATUS_UNSUCCESSFUL; + PVIOWSK_SOCKET pSocket = CONTAINING_RECORD(Socket, VIOWSK_SOCKET, WskSocket); + DEBUG_ENTER_FUNCTION("Socket=0x%p; RemoteAddress=0x%p; Buffer=0x%p; Flags=0x%x; Irp=0x%p", Socket, RemoteAddress, Buffer, Flags, Irp); + UNREFERENCED_PARAMETER(RemoteAddress); UNREFERENCED_PARAMETER(Buffer); UNREFERENCED_PARAMETER(Flags); - return VioWskCompleteIrp(Irp, STATUS_NOT_IMPLEMENTED, 0); + Status = VioWskIrpAcquire(pSocket, Irp); + if (!NT_SUCCESS(Status)) + { + pSocket = NULL; + goto CompleteIrp; + } + + Status = STATUS_NOT_IMPLEMENTED; +CompleteIrp: + VioWskIrpComplete(pSocket, Irp, Status, 0); + + DEBUG_EXIT_FUNCTION("0x%x", Status); + return Status; } NTSTATUS @@ -913,13 +944,28 @@ VioWskSendEx( _Inout_ PIRP Irp ) { - UNREFERENCED_PARAMETER(Socket); + NTSTATUS Status = STATUS_UNSUCCESSFUL; + PVIOWSK_SOCKET pSocket = CONTAINING_RECORD(Socket, VIOWSK_SOCKET, WskSocket); + DEBUG_ENTER_FUNCTION("Socket=0x%p; Buffer=0x%p; Flags=0x%x; ControlInfoLength=%u; ControlInfo=0x%p; Irp=0x%p", Socket, Buffer, Flags, ControlInfoLength, ControlInfo, Irp); + UNREFERENCED_PARAMETER(Buffer); UNREFERENCED_PARAMETER(Flags); UNREFERENCED_PARAMETER(ControlInfoLength); UNREFERENCED_PARAMETER(ControlInfo); - return VioWskCompleteIrp(Irp, STATUS_NOT_IMPLEMENTED, 0); + Status = VioWskIrpAcquire(pSocket, Irp); + if (!NT_SUCCESS(Status)) + { + pSocket = NULL; + goto CompleteIrp; + } + + Status = STATUS_NOT_IMPLEMENTED; +CompleteIrp: + VioWskIrpComplete(pSocket, Irp, Status, 0); + + DEBUG_EXIT_FUNCTION("0x%x", Status); + return Status; } NTSTATUS @@ -934,14 +980,29 @@ VioWskReceiveEx( _Inout_ PIRP Irp ) { - UNREFERENCED_PARAMETER(Socket); + NTSTATUS Status = STATUS_UNSUCCESSFUL; + PVIOWSK_SOCKET pSocket = CONTAINING_RECORD(Socket, VIOWSK_SOCKET, WskSocket); + DEBUG_ENTER_FUNCTION("Socket=0x%p; Buffer=0x%p; Flags=0x%x; ControlInfoLength=0x%p; ControlInfo=0x%p; ControlFlags=0x%p; Irp=0x%p", Socket, Buffer, Flags, ControlInfoLength, ControlInfo, ControlFlags, Irp); + UNREFERENCED_PARAMETER(Buffer); UNREFERENCED_PARAMETER(Flags); UNREFERENCED_PARAMETER(ControlInfoLength); UNREFERENCED_PARAMETER(ControlInfo); UNREFERENCED_PARAMETER(ControlFlags); - return VioWskCompleteIrp(Irp, STATUS_NOT_IMPLEMENTED, 0); + Status = VioWskIrpAcquire(pSocket, Irp); + if (!NT_SUCCESS(Status)) + { + pSocket = NULL; + goto CompleteIrp; + } + + Status = STATUS_NOT_IMPLEMENTED; +CompleteIrp: + VioWskIrpComplete(pSocket, Irp, Status, 0); + + DEBUG_EXIT_FUNCTION("0x%x", Status); + return Status; } NTSTATUS @@ -951,7 +1012,21 @@ VioWskListen( _Inout_ PIRP Irp ) { - UNREFERENCED_PARAMETER(Socket); + NTSTATUS Status = STATUS_UNSUCCESSFUL; + PVIOWSK_SOCKET pSocket = CONTAINING_RECORD(Socket, VIOWSK_SOCKET, WskSocket); + DEBUG_ENTER_FUNCTION("Socket=0x%p; Irp=0x%p", Socket, Irp); + + Status = VioWskIrpAcquire(pSocket, Irp); + if (!NT_SUCCESS(Status)) + { + pSocket = NULL; + goto CompleteIrp; + } - return VioWskCompleteIrp(Irp, STATUS_NOT_IMPLEMENTED, 0); + Status = STATUS_NOT_IMPLEMENTED; +CompleteIrp: + VioWskIrpComplete(pSocket, Irp, Status, 0); + + DEBUG_EXIT_FUNCTION("0x%x", Status); + return Status; } From 493efc4c2a2597858198d0f98b693322828d90c2 Mon Sep 17 00:00:00 2001 From: Martin Drab Date: Wed, 20 Jul 2022 15:22:18 -0700 Subject: [PATCH 17/18] Implement a test driver demonstrating WSK capabilities The leverages WSK to create a set of client and server threads. The server thread listen for connection and exchange a message with connecting clients. Cryptographic hash function (SHA256) is used to verify that message data are not damaged during their transmission. Server threads are working in non-blocking mode. Such behavior is necessary in order to be able to cancel pending WskAccept calls since the viosock.sys driver does not permit their cancellation (unlike send/recv requests which ARE cancellable by design). Signed-off-by: Martin Drab --- buildAll.bat | 2 + viosock/buildAll.bat | 2 + viosock/cleanAll.bat | 4 + viosock/viosock-wsk-test/LICENSE | 28 + viosock/viosock-wsk-test/cleanAll.bat | 2 + viosock/viosock-wsk-test/test-messages.c | 235 +++++ viosock/viosock-wsk-test/test-messages.h | 78 ++ .../viosock-wsk-test/viosock-wsk-test.vcxproj | 111 +++ .../viosock-wsk-test.vcxproj.filters | 30 + viosock/viosock-wsk-test/viosockwsk-test.c | 870 ++++++++++++++++++ viosock/viosock.sln | 24 +- 11 files changed, 1384 insertions(+), 2 deletions(-) create mode 100644 viosock/viosock-wsk-test/LICENSE create mode 100644 viosock/viosock-wsk-test/cleanAll.bat create mode 100644 viosock/viosock-wsk-test/test-messages.c create mode 100644 viosock/viosock-wsk-test/test-messages.h create mode 100644 viosock/viosock-wsk-test/viosock-wsk-test.vcxproj create mode 100644 viosock/viosock-wsk-test/viosock-wsk-test.vcxproj.filters create mode 100644 viosock/viosock-wsk-test/viosockwsk-test.c diff --git a/buildAll.bat b/buildAll.bat index dd18e79ae..5450fc276 100755 --- a/buildAll.bat +++ b/buildAll.bat @@ -28,6 +28,8 @@ call tools\build.bat viosock\sys\viosock.vcxproj "Win11_SDV" %* if errorlevel 1 goto :fail call tools\build.bat viosock\wsk\wsk.vcxproj "Win11_SDV" %* if errorlevel 1 goto :fail +call tools\build.bat viosock\viosock-wsk-test\viosock-wsk-test.vcxproj "Win11_SDV" %* +if errorlevel 1 goto :fail call tools\build.bat viofs\pci\viofs.vcxproj "Win11_SDV" %* if errorlevel 1 goto :fail call tools\build.bat vioinput\hidpassthrough\hidpassthrough.vcxproj "Win11_SDV" %* diff --git a/viosock/buildAll.bat b/viosock/buildAll.bat index 26581edd2..05b7cc4f3 100644 --- a/viosock/buildAll.bat +++ b/viosock/buildAll.bat @@ -6,3 +6,5 @@ if errorlevel 1 goto :eof call ..\tools\build.bat sys\viosock.vcxproj "Win11_SDV" %* if errorlevel 1 goto :eof call ..\tools\build.bat wsk\wsk.vcxproj "Win11_SDV" %* +if errorlevel 1 goto :eof +call ..\tools\build.bat viosock-wsk-test\viosock-wsk-test.vcxproj "Win11_SDV" %* diff --git a/viosock/cleanAll.bat b/viosock/cleanAll.bat index 47f7edb45..2616cfa4a 100644 --- a/viosock/cleanAll.bat +++ b/viosock/cleanAll.bat @@ -17,6 +17,10 @@ pushd wsk call cleanAll.bat popd +pushd viosock-wsk-test +call cleanAll.bat +popd + pushd viosock-test call cleanAll.bat popd diff --git a/viosock/viosock-wsk-test/LICENSE b/viosock/viosock-wsk-test/LICENSE new file mode 100644 index 000000000..95c4900e4 --- /dev/null +++ b/viosock/viosock-wsk-test/LICENSE @@ -0,0 +1,28 @@ +Copyright 2019 Virtuozzo, Inc. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions +are met: + +Redistributions of source code must retain the above copyright +notice, this list of conditions and the following disclaimer. + +Redistributions in binary form must reproduce the above copyright +notice, this list of conditions and the following disclaimer in the +documentation and/or other materials provided with the distribution. + +Neither the name of the copyright holder nor the names of its +contributors may be used to endorse or promote products derived from +this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/viosock/viosock-wsk-test/cleanAll.bat b/viosock/viosock-wsk-test/cleanAll.bat new file mode 100644 index 000000000..de6f2cfe5 --- /dev/null +++ b/viosock/viosock-wsk-test/cleanAll.bat @@ -0,0 +1,2 @@ +@echo off +call ..\..\Tools\clean.bat diff --git a/viosock/viosock-wsk-test/test-messages.c b/viosock/viosock-wsk-test/test-messages.c new file mode 100644 index 000000000..b1c196810 --- /dev/null +++ b/viosock/viosock-wsk-test/test-messages.c @@ -0,0 +1,235 @@ +/* + * Copyright (c) 2023 Virtuozzo International GmbH + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions + * are met : + * 1. Redistributions of source code must retain the above copyright + * notice, this list of conditions and the following disclaimer. + * 2. Redistributions in binary form must reproduce the above copyright + * notice, this list of conditions and the following disclaimer in the + * documentation and / or other materials provided with the distribution. + * 3. Neither the names of the copyright holders nor the names of their contributors + * may be used to endorse or promote products derived from this software + * without specific prior written permission. + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS ``AS IS'' AND + * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + * ARE DISCLAIMED.IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS + * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) + * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT + * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY + * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF + * SUCH DAMAGE. + */ + +#include +#include +#include +#include "..\inc\debug-utils.h" +#include "test-messages.h" + + +static ULONG _randSeed; + + +NTSTATUS +VioWskMessageGenerate( + _In_opt_ BCRYPT_HASH_HANDLE SHA256Handle, + _Out_ PWSK_BUF WskBuffer, + _Out_ PVOID* FlatBuffer +) +{ + PMDL mdl = NULL; + PVOID buffer = NULL; + SIZE_T bufferSize = 0; + NTSTATUS Status = STATUS_UNSUCCESSFUL; + DEBUG_ENTER_FUNCTION("SHA256Handle=0x%p; WskBuffer=0x%p", SHA256Handle, WskBuffer); + + Status = STATUS_SUCCESS; + bufferSize = VIOWSK_MSG_SIZE; + buffer = ExAllocatePoolUninitialized(NonPagedPool, bufferSize, VIOWSK_TEST_MSG_TAG); + if (!buffer) { + Status = STATUS_INSUFFICIENT_RESOURCES; + goto Exit; + } + + if (SHA256Handle) { + for (SIZE_T i = 0; i < (VIOWSK_MSG_SIZE - 32) / sizeof(ULONG); ++i) + ((PULONG)buffer)[i] = RtlRandomEx(&_randSeed); + + Status = BCryptHashData(SHA256Handle, (PUCHAR)buffer, VIOWSK_MSG_SIZE - 32, 0); + if (!NT_SUCCESS(Status)) + goto FreeBuffer; + + Status = BCryptFinishHash(SHA256Handle, (PUCHAR)buffer + VIOWSK_MSG_SIZE - 32, 32, 0); + if (!NT_SUCCESS(Status)) + goto FreeBuffer; + } + + WskBuffer->Length = VIOWSK_MSG_SIZE; + WskBuffer->Offset = 0; + mdl = IoAllocateMdl(buffer, VIOWSK_MSG_SIZE, FALSE, FALSE, NULL); + if (!mdl) { + Status = STATUS_INSUFFICIENT_RESOURCES; + goto FreeBuffer; + } + + MmBuildMdlForNonPagedPool(mdl); + WskBuffer->Mdl = mdl; + *FlatBuffer = buffer; + buffer = NULL; +FreeBuffer: + if (buffer) + ExFreePoolWithTag(buffer, VIOWSK_TEST_MSG_TAG); +Exit: + DEBUG_EXIT_FUNCTION("0x%x, *FlatBuffer=0x%p", Status, *FlatBuffer); + return Status; +} + + +NTSTATUS +VIoWskMessageVerify( + _In_ BCRYPT_HASH_HANDLE SHA256Handle, + _In_ const WSK_BUF* WskBuf, + _Out_ PBOOLEAN Verified +) +{ + ULONG offset = 0; + SIZE_T remainingLength = 0; + PVOID buffer = NULL; + unsigned char digest[32]; + NTSTATUS Status = STATUS_UNSUCCESSFUL; + DEBUG_ENTER_FUNCTION("SHA256Handle=0x%p; WskBuffer=0x%p; Verified=0x%p", SHA256Handle, WskBuf, Verified); + + Status = STATUS_SUCCESS; + buffer = ExAllocatePoolUninitialized(NonPagedPool, WskBuf->Length, VIOWSK_TEST_MSG_TAG); + if (!buffer) { + Status = STATUS_INSUFFICIENT_RESOURCES; + goto Exit; + } + + offset = WskBuf->Offset; + remainingLength = WskBuf->Length; + for (PMDL mdl = WskBuf->Mdl; mdl != NULL; mdl = mdl->Next) { + SIZE_T mdlSize = 0; + PVOID mdlBuffer = NULL; + + mdlBuffer = MmGetSystemAddressForMdlSafe(mdl, NormalPagePriority); + if (!mdlBuffer) { + Status = STATUS_INSUFFICIENT_RESOURCES; + break; + } + + mdlSize = MmGetMdlByteCount(mdl) - offset; + if (mdlSize > remainingLength) + mdlSize = remainingLength; + + memcpy((unsigned char*)buffer + (WskBuf->Length - remainingLength), (unsigned char*)mdlBuffer + offset, mdlSize); + offset = 0; + remainingLength -= mdlSize; + } + + if (!NT_SUCCESS(Status)) + goto FreeBuffer; + + Status = BCryptHashData(SHA256Handle, (PUCHAR)buffer, (ULONG)(WskBuf->Length - sizeof(digest)), 0); + if (!NT_SUCCESS(Status)) + goto FreeBuffer; + + Status = BCryptFinishHash(SHA256Handle, digest, sizeof(digest), 0); + if (!NT_SUCCESS(Status)) + goto FreeBuffer; + + *Verified = (memcmp((unsigned char*)buffer + WskBuf->Length - sizeof(digest), digest, sizeof(digest)) == 0); +FreeBuffer: + ExFreePoolWithTag(buffer, VIOWSK_TEST_MSG_TAG); +Exit: + DEBUG_EXIT_FUNCTION("0x%x, *Verified=%u", Status, *Verified); + return Status; +} + + +NTSTATUS +VIoWskMessageVerifyBuffer( + _In_ BCRYPT_HASH_HANDLE SHA256Handle, + _In_ const void* Buffer, + _Out_ PBOOLEAN Verified +) +{ + unsigned char digest[32]; + NTSTATUS Status = STATUS_UNSUCCESSFUL; + DEBUG_ENTER_FUNCTION("SHA256Handle=0x%p; Buffer=0x%p; Verified=0x%p", SHA256Handle, Buffer, Verified); + + Status = BCryptHashData(SHA256Handle, (PUCHAR)Buffer, VIOWSK_MSG_SIZE - sizeof(digest), 0); + if (!NT_SUCCESS(Status)) + goto Exit; + + Status = BCryptFinishHash(SHA256Handle, digest, sizeof(digest), 0); + if (!NT_SUCCESS(Status)) + goto Exit; + + *Verified = (memcmp((unsigned char*)Buffer + VIOWSK_MSG_SIZE - sizeof(digest), digest, sizeof(digest)) == 0); + +Exit: + DEBUG_EXIT_FUNCTION("0x%x, *Verified=%u", Status, *Verified); + return Status; +} + + +void +VioWskMessageAdvance( + _Inout_ PWSK_BUF WskBuffer, + _In_ SIZE_T Length +) +{ + PMDL currentMdl = NULL; + ULONG advanceAmount = 0; + DEBUG_ENTER_FUNCTION("WskBuffer=0x%p; Length=%Iu", WskBuffer, Length); + + currentMdl = WskBuffer->Mdl; + while (WskBuffer->Length > 0 && Length > 0) { + advanceAmount = MmGetMdlByteCount(WskBuffer->Mdl) - WskBuffer->Offset; + if (advanceAmount > Length) + advanceAmount = (ULONG)Length; + + WskBuffer->Offset += advanceAmount; + if (WskBuffer->Offset == MmGetMdlByteCount(currentMdl)) { + WskBuffer->Mdl = WskBuffer->Mdl->Next; + WskBuffer->Offset = 0; + IoFreeMdl(currentMdl); + currentMdl = WskBuffer->Mdl; + } + + WskBuffer->Length -= advanceAmount; + Length -= advanceAmount; + } + + DEBUG_EXIT_FUNCTION_VOID(); + return; +} + +void +VioWskMessageFree( + _In_ PWSK_BUF WskBuffer, + _In_opt_ PVOID FlatBuffer +) +{ + PMDL mdl = NULL; + DEBUG_ENTER_FUNCTION("WskBuffer=0x%p; FlatBuffer=0x%p", WskBuffer, FlatBuffer); + + mdl = WskBuffer->Mdl; + while (mdl != NULL) { + mdl = mdl->Next; + IoFreeMdl(WskBuffer->Mdl); + WskBuffer->Mdl = mdl; + } + + if (FlatBuffer) + ExFreePoolWithTag(FlatBuffer, VIOWSK_TEST_MSG_TAG); + + DEBUG_EXIT_FUNCTION_VOID(); + return; +} diff --git a/viosock/viosock-wsk-test/test-messages.h b/viosock/viosock-wsk-test/test-messages.h new file mode 100644 index 000000000..4b668943c --- /dev/null +++ b/viosock/viosock-wsk-test/test-messages.h @@ -0,0 +1,78 @@ +/* + * Copyright (c) 2023 Virtuozzo International GmbH + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions + * are met : + * 1. Redistributions of source code must retain the above copyright + * notice, this list of conditions and the following disclaimer. + * 2. Redistributions in binary form must reproduce the above copyright + * notice, this list of conditions and the following disclaimer in the + * documentation and / or other materials provided with the distribution. + * 3. Neither the names of the copyright holders nor the names of their contributors + * may be used to endorse or promote products derived from this software + * without specific prior written permission. + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS ``AS IS'' AND + * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + * ARE DISCLAIMED.IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS + * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) + * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT + * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY + * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF + * SUCH DAMAGE. + */ + +#ifndef __VIOWSK_TEST_MESSAGES_H__ +#define __VIOWSK_TEST_MESSAGES_H__ + + +#include +#include +#include + + + +#define VIOWSK_TEST_MSG_TAG (ULONG)'MKSW' + +#define VIOWSK_MSG_SIZE 64 + +NTSTATUS +VioWskMessageGenerate( + _In_opt_ BCRYPT_HASH_HANDLE SHA256Handle, + _Out_ PWSK_BUF WskBuffer, + _Out_ PVOID* FlatBuffer +); + +NTSTATUS +VIoWskMessageVerify( + _In_ BCRYPT_HASH_HANDLE SHA256Handle, + _In_ const WSK_BUF* WskBuf, + _Out_ PBOOLEAN Verified +); + +NTSTATUS +VIoWskMessageVerifyBuffer( + _In_ BCRYPT_HASH_HANDLE SHA256Handle, + _In_ const void* Buffer, + _Out_ PBOOLEAN Verified +); + +void +VioWskMessageAdvance( + _Inout_ PWSK_BUF WskBuffer, + _In_ SIZE_T Length +); + +void +VioWskMessageFree( + _In_ PWSK_BUF WskBuffer, + _In_opt_ PVOID FlatBuffer +); + + + + +#endif diff --git a/viosock/viosock-wsk-test/viosock-wsk-test.vcxproj b/viosock/viosock-wsk-test/viosock-wsk-test.vcxproj new file mode 100644 index 000000000..4b35c9115 --- /dev/null +++ b/viosock/viosock-wsk-test/viosock-wsk-test.vcxproj @@ -0,0 +1,111 @@ + + + + + Win10 Release + Win32 + + + Win10 Release + x64 + + + Win10 Release + ARM64 + + + + {A2A7AFFC-BEFE-4066-A1DD-23242FAFAD70} + {dd38f7fc-d7bd-488b-9242-7d8754cde80d} + v4.5 + 12.0 + Debug + Win32 + viosock_wsk_test + viosockwsk-test + $(LatestTargetPlatformVersion) + + + + Windows10 + false + WindowsKernelModeDriver10.0 + Driver + WDM + Desktop + + + Windows10 + false + WindowsKernelModeDriver10.0 + Driver + WDM + Desktop + + + Windows10 + false + WindowsKernelModeDriver10.0 + Driver + WDM + Desktop + + + + + + + + + + + DbgengKernelDebugger + false + objfre_win10_x86\i386\ + objfre_win10_x86\i386\ + + + DbgengKernelDebugger + false + objfre_win10_amd64\amd64\ + objfre_win10_amd64\amd64\ + + + DbgengKernelDebugger + objfre_win10_arm64\arm64\ + objfre_win10_arm64\arm64\ + + + + cng.lib;%(AdditionalDependencies);$(KernelBufferOverflowLib);$(DDK_LIB_PATH)ntoskrnl.lib;$(DDK_LIB_PATH)hal.lib;$(DDK_LIB_PATH)wmilib.lib + + + + + cng.lib;%(AdditionalDependencies);$(KernelBufferOverflowLib);$(DDK_LIB_PATH)ntoskrnl.lib;$(DDK_LIB_PATH)hal.lib;$(DDK_LIB_PATH)wmilib.lib + + + + + cng.lib;%(AdditionalDependencies);$(KernelBufferOverflowLib);$(DDK_LIB_PATH)ntoskrnl.lib;$(DDK_LIB_PATH)hal.lib;$(DDK_LIB_PATH)wmilib.lib + + + + + + + + + + + + {96fdd976-0035-4e24-a61b-e93bed675101} + + + + + + + + + \ No newline at end of file diff --git a/viosock/viosock-wsk-test/viosock-wsk-test.vcxproj.filters b/viosock/viosock-wsk-test/viosock-wsk-test.vcxproj.filters new file mode 100644 index 000000000..c745a968c --- /dev/null +++ b/viosock/viosock-wsk-test/viosock-wsk-test.vcxproj.filters @@ -0,0 +1,30 @@ + + + + + {4FC737F1-C7A5-4376-A066-2A32D752A2FF} + cpp;c;cc;cxx;def;odl;idl;hpj;bat;asm;asmx + + + {93995380-89BD-4b04-88EB-625FBE52EBFB} + h;hpp;hxx;hm;inl;inc;xsd + + + {67DA6AB6-F800-4c08-8B7A-83BB121AAD01} + rc;ico;cur;bmp;dlg;rc2;rct;bin;rgs;gif;jpg;jpeg;jpe;resx;tiff;tif;png;wav;mfcribbon-ms + + + + + Source Files + + + Source Files + + + + + Header Files + + + \ No newline at end of file diff --git a/viosock/viosock-wsk-test/viosockwsk-test.c b/viosock/viosock-wsk-test/viosockwsk-test.c new file mode 100644 index 000000000..b123358ba --- /dev/null +++ b/viosock/viosock-wsk-test/viosockwsk-test.c @@ -0,0 +1,870 @@ +/* + * Copyright (c) 2023 Virtuozzo International GmbH + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions + * are met : + * 1. Redistributions of source code must retain the above copyright + * notice, this list of conditions and the following disclaimer. + * 2. Redistributions in binary form must reproduce the above copyright + * notice, this list of conditions and the following disclaimer in the + * documentation and / or other materials provided with the distribution. + * 3. Neither the names of the copyright holders nor the names of their contributors + * may be used to endorse or promote products derived from this software + * without specific prior written permission. + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS ``AS IS'' AND + * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + * ARE DISCLAIMED.IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS + * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) + * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT + * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY + * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF + * SUCH DAMAGE. + */ + +/** This driver shows how to use the WSK interface for the Virtio Socket Driver. + The driver creates four client and four server threads. Servers listen + on predefined ports, clients are randomly connecting to them and exchanging + serveral messages. Messages are 64 bytes and consist of 32 byte long sequence of + random bytes plus SHA256 hash of this sequence (to check that the message was transferred + intact). + + Things worth of noting: + - the WSK interface is implemented in the viosockwsk library that is linked to the test driver, + - the library needs to be initialized via VioWskModuleInit and finalized through VioWskModuleFinit, + - a device object can be optionally provided to VioWskModuleInit; the library uses it as parent for + internal workitems and other objects, + - server threads perform accept() (WskAccept) in non-blocking mode in order to be able to correctly + terminate themselves when the driver is being unloaded (the socket driver does not support accept cancellation), + - the driver runs its tests infinitely (results can be observer via WPP Tracing) until it gets unloaded. +*/ + +#include +#include +#include +#include +#include +#include "..\inc\debug-utils.h" +#include "..\inc\vio_wsk.h" +#include "..\inc\vio_sockets.h" +#include "..\sys\public.h" +#include "test-messages.h" + + + +#define LISTEN_PORT_MIN 1337 +#define LISTEN_PORT_MAX 1340 +#define SERVER_THREAD_COUNT ((LISTEN_PORT_MAX) - (LISTEN_PORT_MIN) + 1) +#define CLIENT_THREAD_COUNT 16 + +#define VIOWSK_TEST_TAG (ULONG)'TKSW' + +static volatile LONG _readyThreads; +static volatile LONG _terminate; +static PETHREAD _clientThreads[CLIENT_THREAD_COUNT]; +static PETHREAD _serverThreads[SERVER_THREAD_COUNT]; +static KEVENT _initEvent; +static PDEVICE_OBJECT _shutdownDeviceObject = NULL; +static WSK_REGISTRATION _vioWskRegistration; +static WSK_PROVIDER_NPI _vioWskProviderNPI; +static WSK_CLIENT_NPI _vioWskClientNPI = { + NULL, + NULL, +}; + + +static +NTSTATUS +_GeneralIrpComplete( + _In_ PDEVICE_OBJECT DeviceObject, + _In_ PIRP Irp, + _In_ PVOID Context +) +{ + NTSTATUS status = STATUS_UNSUCCESSFUL; + DEBUG_ENTER_FUNCTION("DeviceObject=0x%p; Irp=0x%p; Context=0x%p", DeviceObject, Irp, Context); + + UNREFERENCED_PARAMETER(DeviceObject); + UNREFERENCED_PARAMETER(Irp); + + KeSetEvent((PKEVENT)Context, IO_NO_INCREMENT, FALSE); + status = STATUS_MORE_PROCESSING_REQUIRED; + + DEBUG_EXIT_FUNCTION("0x%x", status); + return status; +} + + +#define WSK_SYNCHRONOUS_CALL(aIrp, aEvent, aCall, aIosb) do { \ + IoSetCompletionRoutine((aIrp), _GeneralIrpComplete, (aEvent), TRUE, TRUE, TRUE); \ + (aIosb)->Status = (aCall); \ + if ((aIosb)->Status == STATUS_PENDING) \ + { \ + NTSTATUS __waitResult = STATUS_TIMEOUT; \ + LARGE_INTEGER __timeout; \ + \ + __timeout.QuadPart = -10000000; \ + do \ + { \ + __waitResult = KeWaitForSingleObject((aEvent), Executive, KernelMode, FALSE, &__timeout); \ + if (__waitResult == STATUS_TIMEOUT && \ + InterlockedCompareExchange(&_terminate, 1, 1)) \ + { \ + IoCancelIrp((aIrp)); \ + break; \ + } \ + } while (__waitResult != STATUS_WAIT_0); \ + \ + KeWaitForSingleObject((aEvent), Executive, KernelMode, FALSE, NULL); \ + (aIosb)->Status = (aIrp)->IoStatus.Status; \ + } \ + \ + ASSERT((aIosb)->Status == (aIrp)->IoStatus.Status); \ + (aIosb)->Information = (aIrp)->IoStatus.Information; \ + KeResetEvent((aEvent)); \ + IoReuseIrp((aIrp), STATUS_UNSUCCESSFUL); \ +} while (FALSE) \ + + +static +NTSTATUS +_TestSocket( + _In_ PWSK_SOCKET Socket, + _In_ PIRP Irp, + _In_ BOOLEAN Server +) +{ + KEVENT event; + IO_STATUS_BLOCK iosb; + WSK_BUF msg; + PVOID msgFlat = NULL; + WSK_BUF recvMsg; + PVOID recvFlat = NULL; + ULONG hashObjectSize = 0; + ULONG returnedLength = 0; + void *hashObject = NULL; + BCRYPT_HASH_HANDLE hashHandle = NULL; + BOOLEAN verified = FALSE; + BCRYPT_ALG_HANDLE sha256Handle = NULL; + DEBUG_ENTER_FUNCTION("Socket=0x%p; Irp=0x%p; Server=%u", Socket, Irp, Server); + + KeInitializeEvent(&event, NotificationEvent, FALSE); + iosb.Status = BCryptOpenAlgorithmProvider(&sha256Handle, BCRYPT_SHA256_ALGORITHM, NULL, BCRYPT_PROV_DISPATCH); + if (!NT_SUCCESS(iosb.Status)) + { + DEBUG_ERROR("Unable to open SHA256 provider: 0x%x", iosb.Status); + goto Exit; + } + + iosb.Status = BCryptGetProperty(sha256Handle, BCRYPT_OBJECT_LENGTH, (PUCHAR)&hashObjectSize, sizeof(hashObjectSize), &returnedLength, 0); + if (!NT_SUCCESS(iosb.Status)) + { + DEBUG_ERROR("BCryptGetProperty: 0x%x", iosb.Status); + goto CloseProvider; + } + + hashObject = ExAllocatePoolUninitialized(NonPagedPool, hashObjectSize, VIOWSK_TEST_TAG); + if (!hashObject) + { + iosb.Status = STATUS_INSUFFICIENT_RESOURCES; + goto CloseProvider; + } + + for (size_t i = 0; i < 16; ++i) + { + if (!NT_SUCCESS(iosb.Status)) + break; + + iosb.Status = VioWskMessageGenerate(NULL, &recvMsg, &recvFlat); + if (!NT_SUCCESS(iosb.Status)) + { + DEBUG_ERROR("VioWskMessageGenerate: 0x%x", iosb.Status); + continue; + } + + iosb.Status = BCryptCreateHash(sha256Handle, &hashHandle, hashObject, hashObjectSize, NULL, 0, 0); + if (!NT_SUCCESS(iosb.Status)) + { + DEBUG_ERROR("BCryptCreateHash: 0x%x", iosb.Status); + goto FreeRecvMessage; + } + + iosb.Status = VioWskMessageGenerate(hashHandle, &msg, &msgFlat); + BCryptDestroyHash(hashObject); + if (!NT_SUCCESS(iosb.Status)) + { + DEBUG_ERROR("Unable to generate test message: 0x%x", iosb.Status); + goto FreeRecvMessage; + } + + iosb.Status = BCryptCreateHash(sha256Handle, &hashHandle, hashObject, hashObjectSize, NULL, 0, 0); + if (!NT_SUCCESS(iosb.Status)) + { + DEBUG_ERROR("BCryptCreateHash: 0x%x", iosb.Status); + goto FreeMessage; + } + + iosb.Status = VIoWskMessageVerify(hashHandle, &msg, &verified); + BCryptDestroyHash(hashObject); + if (!NT_SUCCESS(iosb.Status) || !verified) + { + if (!verified) + { + iosb.Status = STATUS_UNSUCCESSFUL; + DEBUG_ERROR("Generated test message is invalid: 0x%x", iosb.Status); + } else { + DEBUG_ERROR("Unable to verify test message: 0x%x", iosb.Status); + } + + goto FreeMessage; + } + + if (!Server) + { + while (NT_SUCCESS(iosb.Status) && msg.Length > 0) + { + WSK_SYNCHRONOUS_CALL(Irp, &event, ((PWSK_PROVIDER_CONNECTION_DISPATCH)Socket->Dispatch)->WskSend(Socket, &msg, 0, Irp), &iosb); + if (iosb.Status == STATUS_CANT_WAIT) + { + LARGE_INTEGER timeout; + + iosb.Status = STATUS_SUCCESS; + DEBUG_WARNING("Cannot wait for receive, sleeping\n"); + timeout.QuadPart = -10000000; + KeDelayExecutionThread(KernelMode, FALSE, &timeout); + continue; + } + + if (!NT_SUCCESS(iosb.Status) || iosb.Information != msg.Length) + { + DEBUG_ERROR("Unable to send the test message: 0x%x (%Iu sent, %Iu requested)\n", iosb.Status, iosb.Information, msg.Length); + goto FreeMessage; + } + + DEBUG_INFO("%Iu bytes sent (0x%x)\n", iosb.Information, iosb.Status); + VioWskMessageAdvance(&msg, iosb.Information); + } + } + + while (NT_SUCCESS(iosb.Status) && recvMsg.Length > 0) + { + WSK_SYNCHRONOUS_CALL(Irp, &event, ((PWSK_PROVIDER_CONNECTION_DISPATCH)Socket->Dispatch)->WskReceive(Socket, &recvMsg, 0, Irp), &iosb); + if (iosb.Status == STATUS_CANT_WAIT) + { + LARGE_INTEGER timeout; + + iosb.Status = STATUS_SUCCESS; + DEBUG_WARNING("Cannot wait for receive, sleeping\n"); + timeout.QuadPart = -10000000; + KeDelayExecutionThread(KernelMode, FALSE, &timeout); + continue; + } + + if (!NT_SUCCESS(iosb.Status) || recvMsg.Length != iosb.Information) + { + DEBUG_ERROR("Unable to receive the tes message: 0x%x (%Iu bytes length, %Iu bytes received)", iosb.Status, recvMsg.Length, iosb.Information); + goto FreeMessage; + } + + DEBUG_INFO("%Iu bytes received (0x%x)\n", iosb.Information, iosb.Status); + VioWskMessageAdvance(&recvMsg, iosb.Information); + } + + iosb.Status = BCryptCreateHash(sha256Handle, &hashHandle, hashObject, hashObjectSize, NULL, 0, 0); + if (!NT_SUCCESS(iosb.Status)) + { + DEBUG_ERROR("BCryptCreateHash: 0x%x", iosb.Status); + goto SendMessage; + } + + iosb.Status = VIoWskMessageVerifyBuffer(hashHandle, recvFlat, &verified); + BCryptDestroyHash(hashObject); + if (!NT_SUCCESS(iosb.Status) || !verified) + { + if (!verified) + iosb.Status = STATUS_UNSUCCESSFUL; + + DEBUG_ERROR("Unable to verify test message: 0x%x", iosb.Status); + goto SendMessage; + } + + SendMessage: + if (Server) + { + while (NT_SUCCESS(iosb.Status) && msg.Length > 0) + { + WSK_SYNCHRONOUS_CALL(Irp, &event, ((PWSK_PROVIDER_CONNECTION_DISPATCH)Socket->Dispatch)->WskSend(Socket, &msg, 0, Irp), &iosb); + if (iosb.Status == STATUS_CANT_WAIT) + { + LARGE_INTEGER timeout; + + iosb.Status = STATUS_SUCCESS; + DEBUG_WARNING("Cannot wait for receive, sleeping\n"); + timeout.QuadPart = -10000000; + KeDelayExecutionThread(KernelMode, FALSE, &timeout); + continue; + } + + if (!NT_SUCCESS(iosb.Status) || iosb.Information != msg.Length) + { + DEBUG_ERROR("Unable to send the test message: 0x%x (%Iu sent, %Iu requested)\n", iosb.Status, iosb.Information, msg.Length); + goto FreeMessage; + } + + DEBUG_INFO("%Iu bytes sent (0x%x)\n", iosb.Information, iosb.Status); + VioWskMessageAdvance(&msg, iosb.Information); + } + } + FreeMessage: + VioWskMessageFree(&msg, msgFlat); + FreeRecvMessage: + VioWskMessageFree(&recvMsg, recvFlat); + } + + if (Server) + { + DEBUG_INFO("Server thread test finished\n"); + } + + ExFreePoolWithTag(hashObject, VIOWSK_TEST_TAG); +CloseProvider: + BCryptCloseAlgorithmProvider(sha256Handle, 0); +Exit: + DEBUG_EXIT_FUNCTION("0x%x", iosb.Status); + return iosb.Status; +} + + +typedef struct _TEST_THREAD_CONTEXT { + LIST_ENTRY Entry; + PKSPIN_LOCK ListLock; + PETHREAD Thread; + PIRP Irp; + PWSK_SOCKET Socket; + volatile LONG Terminated; +} TEST_THREAD_CONTEXT, *PTEST_THREAD_CONTEXT; + +static +NTSTATUS +_SocketTestThreadCreate( + _In_ PWSK_SOCKET Socket, + _In_ PLIST_ENTRY ListHead, + _In_ PKSPIN_LOCK ListLock +); + +static +void +_TestThreadRoutine( + _In_ PVOID Context +) +{ + KIRQL irql; + KEVENT event; + IO_STATUS_BLOCK iosb; + BOOLEAN isActive = FALSE; + NTSTATUS Status = STATUS_UNSUCCESSFUL; + PTEST_THREAD_CONTEXT ctx = (PTEST_THREAD_CONTEXT)Context; + DEBUG_ENTER_FUNCTION("Context=0x%p", Context); + + Status = _TestSocket(ctx->Socket, ctx->Irp, TRUE); + KeAcquireSpinLock(ctx->ListLock, &irql); + if (!IsListEmpty(&ctx->Entry)) + { + RemoveEntryList(&ctx->Entry); + isActive = TRUE; + } + + KeReleaseSpinLock(ctx->ListLock, irql); + if (isActive) + { + memset(&iosb, 0, sizeof(iosb)); + KeInitializeEvent(&event, NotificationEvent, FALSE); + WSK_SYNCHRONOUS_CALL(ctx->Irp, &event, ((PWSK_PROVIDER_CONNECTION_DISPATCH)ctx->Socket->Dispatch)->WskDisconnect(ctx->Socket, NULL, 0, ctx->Irp), &iosb); + if (iosb.Status == STATUS_CONNECTION_INVALID) + iosb.Status = STATUS_SUCCESS; + + if (!NT_SUCCESS(iosb.Status)) + { + DEBUG_ERROR("Unable to disconnect the server socket: 0x%x", iosb.Status); + } + + WSK_SYNCHRONOUS_CALL(ctx->Irp, &event, ((PWSK_PROVIDER_BASIC_DISPATCH)ctx->Socket->Dispatch)->WskCloseSocket(ctx->Socket, ctx->Irp), &iosb); + if (!NT_SUCCESS(iosb.Status)) + { + DEBUG_ERROR("Unable to close the server socket: 0x%x", iosb.Status); + } + + ObDereferenceObject(ctx->Thread); + IoFreeIrp(ctx->Irp); + ExFreePoolWithTag(ctx, VIOWSK_TEST_TAG); + } + + DEBUG_EXIT_FUNCTION("0x%x", Status); + return; +} + + +static +void +_ServerThreadRoutine( + _In_opt_ PVOID Context +) +{ + KIRQL irql; + KEVENT event; + ULONG nonBlocking = 0; + PIRP irp = NULL; + SOCKADDR_VM listenAddress; + DECLARE_UNICODE_STRING_SIZE(listenHost, 16); + DECLARE_UNICODE_STRING_SIZE(listenPort, 16); + PADDRINFOEXW addrInfo = NULL; + PWSK_SOCKET serverSocket = NULL; + IO_STATUS_BLOCK iosb; + KSPIN_LOCK threadListLock; + LIST_ENTRY threadListHead; + DEBUG_ENTER_FUNCTION("Context=0x%p", Context); + + if (InterlockedIncrement(&_readyThreads) == (CLIENT_THREAD_COUNT + SERVER_THREAD_COUNT)) + KeSetEvent(&_initEvent, IO_NO_INCREMENT, FALSE); + + irp = IoAllocateIrp(1, FALSE); + if (!irp) { + iosb.Status = STATUS_INSUFFICIENT_RESOURCES; + DEBUG_ERROR("Unable to allocate IRP: 0x%x", iosb.Status); + goto Exit; + } + + KeInitializeEvent(&event, NotificationEvent, FALSE); + WSK_SYNCHRONOUS_CALL(irp, &event, _vioWskProviderNPI.Dispatch->WskSocket(_vioWskProviderNPI.Client, AF_VSOCK, SOCK_STREAM, 0, WSK_FLAG_LISTEN_SOCKET, NULL, NULL, NULL, NULL, NULL, irp), &iosb); + if (!NT_SUCCESS(iosb.Status)) { + DEBUG_ERROR("Unable to create server socket: 0x%x", iosb.Status); + goto FreeIrp; + } + + serverSocket = (PWSK_SOCKET)iosb.Information; + iosb.Status = RtlUnicodeStringPrintf(&listenHost, L"%u", (ULONG)VMADDR_CID_ANY); + if (!NT_SUCCESS(iosb.Status)) { + DEBUG_ERROR("Unable to prepare server address hos: 0x%x", iosb.Status); + goto CloseSocket; + } + + iosb.Status = RtlUnicodeStringPrintf(&listenPort, L"%Iu", LISTEN_PORT_MIN + (SIZE_T)Context); + if (!NT_SUCCESS(iosb.Status)) { + DEBUG_ERROR("Unable to prepare server address port: 0x%x", iosb.Status); + goto CloseSocket; + } + + WSK_SYNCHRONOUS_CALL(irp, &event, _vioWskProviderNPI.Dispatch->WskGetAddressInfo(_vioWskProviderNPI.Client, &listenHost, &listenPort, 0, NULL, NULL, &addrInfo, NULL, NULL, irp), &iosb); + if (!NT_SUCCESS(iosb.Status)) { + DEBUG_ERROR("Unable to translate the listen address: 0x%x", iosb.Status); + goto CloseSocket; + } + + listenAddress = *(PSOCKADDR_VM)addrInfo->ai_addr; + _vioWskProviderNPI.Dispatch->WskFreeAddressInfo(_vioWskProviderNPI.Client, addrInfo); + + WSK_SYNCHRONOUS_CALL(irp, &event, ((PWSK_PROVIDER_LISTEN_DISPATCH)serverSocket->Dispatch)->WskBind(serverSocket, (PSOCKADDR)&listenAddress, 0, irp), &iosb); + if (!NT_SUCCESS(iosb.Status)) { + DEBUG_ERROR("Unable to bind: 0x%x", iosb.Status); + goto CloseSocket; + } + + nonBlocking = TRUE; + WSK_SYNCHRONOUS_CALL(irp, &event, ((PWSK_PROVIDER_LISTEN_DISPATCH)serverSocket->Dispatch)->WskControlSocket(serverSocket, WskIoctl, FIONBIO, 0, sizeof(nonBlocking), &nonBlocking, 0, NULL, NULL, irp), &iosb); + if (!NT_SUCCESS(iosb.Status)) { + DEBUG_ERROR("Unable to be nonblocking: 0x%x", iosb.Status); + goto CloseSocket; + } + + nonBlocking = FALSE; + InitializeListHead(&threadListHead); + KeInitializeSpinLock(&threadListLock); + while (!InterlockedCompareExchange(&_terminate, 1, 1)) { + PWSK_SOCKET clientSocket = NULL; + SOCKADDR_VM localAddr; + SOCKADDR_VM remoteAddr; + LARGE_INTEGER timeout; + + timeout.QuadPart = -10000000; + WSK_SYNCHRONOUS_CALL(irp, &event, ((PWSK_PROVIDER_LISTEN_DISPATCH)serverSocket->Dispatch)->WskAccept(serverSocket, WSK_FLAG_CONNECTION_SOCKET, NULL, NULL, (PSOCKADDR)&localAddr, (PSOCKADDR)&remoteAddr, irp), &iosb); + if (iosb.Status == STATUS_CANT_WAIT) { + KeDelayExecutionThread(KernelMode, FALSE, &timeout); + continue; + } + + if (!NT_SUCCESS(iosb.Status)) { + DEBUG_ERROR("Unable to accept: 0x%x", iosb.Status); + break; + } + + clientSocket = (PWSK_SOCKET)iosb.Information; + WSK_SYNCHRONOUS_CALL(irp, &event, ((PWSK_PROVIDER_LISTEN_DISPATCH)clientSocket->Dispatch)->WskControlSocket(clientSocket, WskIoctl, FIONBIO, 0, sizeof(nonBlocking), &nonBlocking, 0, NULL, NULL, irp), &iosb); + if (!NT_SUCCESS(iosb.Status)) { + DEBUG_ERROR("Unable to be blocking: 0x%x", iosb.Status); + } + + if (NT_SUCCESS(iosb.Status)) { + WSK_SYNCHRONOUS_CALL(irp, &event, _vioWskProviderNPI.Dispatch->WskGetNameInfo(_vioWskProviderNPI.Client, (PSOCKADDR)&localAddr, sizeof(localAddr), &listenHost, &listenPort, 0, NULL, NULL, irp), &iosb); + if (NT_SUCCESS(iosb.Status)) { + DEBUG_INFO("Accepted connection:"); + DEBUG_INFO(" Local: %wZ:%wZ", &listenHost, &listenPort); + WSK_SYNCHRONOUS_CALL(irp, &event, _vioWskProviderNPI.Dispatch->WskGetNameInfo(_vioWskProviderNPI.Client, (PSOCKADDR)&remoteAddr, sizeof(remoteAddr), &listenHost, &listenPort, 0, NULL, NULL, irp), &iosb); + if (!NT_SUCCESS(iosb.Status)) { + DEBUG_ERROR("Unable to get client address strings: 0x%x", iosb.Status); + } + } + } + + if (NT_SUCCESS(iosb.Status)) { + DEBUG_INFO(" Remote: %wZ:%wZ", &listenHost, &listenPort); + iosb.Status = _SocketTestThreadCreate(clientSocket, &threadListHead, &threadListLock); + if (!NT_SUCCESS(iosb.Status)) { + DEBUG_ERROR("Socket test thread failed to create: 0x%x", iosb.Status); + } + } + + if (!NT_SUCCESS(iosb.Status)) { + WSK_SYNCHRONOUS_CALL(irp, &event, ((PWSK_PROVIDER_BASIC_DISPATCH)clientSocket->Dispatch)->WskCloseSocket(clientSocket, irp), &iosb); + break; + } + } + + KeAcquireSpinLock(&threadListLock, &irql); + while (!IsListEmpty(&threadListHead)) + { + PTEST_THREAD_CONTEXT ctx = CONTAINING_RECORD(threadListHead.Flink, TEST_THREAD_CONTEXT, Entry); + + RemoveEntryList(&ctx->Entry); + InitializeListHead(&ctx->Entry); + KeReleaseSpinLock(&threadListLock, irql); + InterlockedExchange(&ctx->Terminated, 1); + KeWaitForSingleObject(ctx->Thread, Executive, KernelMode, FALSE, NULL); + ObDereferenceObject(ctx->Thread); + WSK_SYNCHRONOUS_CALL(irp, &event, ((PWSK_PROVIDER_BASIC_DISPATCH)ctx->Socket->Dispatch)->WskCloseSocket(ctx->Socket, irp), &iosb); + IoFreeIrp(ctx->Irp); + ExFreePoolWithTag(ctx, VIOWSK_TEST_TAG); + KeAcquireSpinLock(&threadListLock, &irql); + } + + KeReleaseSpinLock(&threadListLock, irql); +CloseSocket: + WSK_SYNCHRONOUS_CALL(irp, &event, ((PWSK_PROVIDER_BASIC_DISPATCH)serverSocket->Dispatch)->WskCloseSocket(serverSocket, irp), &iosb); + if (!NT_SUCCESS(iosb.Status)) { + DEBUG_ERROR("Unable to close the server socket: 0x%x", iosb.Status); + } +FreeIrp: + IoFreeIrp(irp); +Exit: + InterlockedExchange(&_terminate, 1); + + DEBUG_EXIT_FUNCTION("0x%x", iosb.Status); + return; +} + + +static +void +_ClientThreadRoutine( + _In_opt_ PVOID Context +) +{ + KEVENT event; + PIRP irp = NULL; + LARGE_INTEGER timeSeed; + LARGE_INTEGER timeout; + IO_STATUS_BLOCK iosb; + SOCKADDR_VM localAddr; + SOCKADDR_VM remoteAddr; + PWSK_SOCKET socket = NULL; + DEBUG_ENTER_FUNCTION("Context=0x%p", Context); + + UNREFERENCED_PARAMETER(Context); + + iosb.Status = STATUS_SUCCESS; + if (InterlockedIncrement(&_readyThreads) == (CLIENT_THREAD_COUNT + SERVER_THREAD_COUNT)) + KeSetEvent(&_initEvent, IO_NO_INCREMENT, FALSE); + + irp = IoAllocateIrp(1, FALSE); + if (!irp) { + iosb.Status = STATUS_INSUFFICIENT_RESOURCES; + DEBUG_ERROR("Failed to allocate IRP: 0x%x", iosb.Status); + goto Exit; + } + + timeout.QuadPart = -10000000; + KeQuerySystemTime(&timeSeed); + KeInitializeEvent(&event, NotificationEvent, FALSE); + while (!InterlockedCompareExchange(&_terminate, 1, 1)) { + memset(&localAddr, 0, sizeof(localAddr)); + localAddr.svm_family = AF_VSOCK; + localAddr.svm_cid = (UINT)VMADDR_CID_ANY; + localAddr.svm_port = (UINT)VMADDR_PORT_ANY; + memset(&remoteAddr, 0, sizeof(remoteAddr)); + remoteAddr.svm_family = AF_VSOCK; + remoteAddr.svm_cid = (UINT)VMADDR_CID_ANY; + remoteAddr.svm_port = LISTEN_PORT_MIN + (RtlRandomEx(&timeSeed.LowPart) % (SERVER_THREAD_COUNT)); + WSK_SYNCHRONOUS_CALL(irp, &event, _vioWskProviderNPI.Dispatch->WskSocket(_vioWskProviderNPI.Client, AF_VSOCK, SOCK_STREAM, 0, WSK_FLAG_CONNECTION_SOCKET, NULL, NULL, NULL, NULL, NULL, irp), &iosb); + if (!NT_SUCCESS(iosb.Status)) { + DEBUG_ERROR("Unable to create client socket: 0x%x", iosb.Status); + break; + } + + socket = (PWSK_SOCKET)iosb.Information; + WSK_SYNCHRONOUS_CALL(irp, &event, ((PWSK_PROVIDER_CONNECTION_DISPATCH)socket->Dispatch)->WskBind(socket, (PSOCKADDR)&localAddr, 0, irp), &iosb); + if (NT_SUCCESS(iosb.Status)) + { + WSK_SYNCHRONOUS_CALL(irp, &event, ((PWSK_PROVIDER_CONNECTION_DISPATCH)socket->Dispatch)->WskConnect(socket, (PSOCKADDR)&remoteAddr, 0, irp), &iosb); + } + + if (NT_SUCCESS(iosb.Status)) { + iosb.Status = _TestSocket(socket, irp, FALSE); + if (!NT_SUCCESS(iosb.Status)) { + DEBUG_ERROR("Client socket test failed: 0x%x", iosb.Status); + } + + WSK_SYNCHRONOUS_CALL(irp, &event, ((PWSK_PROVIDER_CONNECTION_DISPATCH)socket->Dispatch)->WskDisconnect(socket, NULL, 0, irp), &iosb); + if (iosb.Status == STATUS_CONNECTION_INVALID) + iosb.Status = STATUS_SUCCESS; + + if (!NT_SUCCESS(iosb.Status)) { + DEBUG_ERROR("Client socket disconnect failed: 0x%x", iosb.Status); + } + } else { + DEBUG_ERROR("Unable to connect to the server: 0x%x", iosb.Status); + } + + WSK_SYNCHRONOUS_CALL(irp, &event, ((PWSK_PROVIDER_CONNECTION_DISPATCH)socket->Dispatch)->WskCloseSocket(socket, irp), &iosb); + if (!NT_SUCCESS(iosb.Status)) { + DEBUG_ERROR("Unable to close the client socket: 0x%x", iosb.Status); + } + + DEBUG_INFO("Client thread test finished\n"); + KeDelayExecutionThread(KernelMode, FALSE, &timeout); + } + + IoFreeIrp(irp); +Exit: + InterlockedExchange(&_terminate, 1); + + DEBUG_EXIT_FUNCTION("0x%x", iosb.Status); + return; +} + + +static +void +_DestroyThreadGroup( + _In_ PETHREAD* ObjectArray, + _In_ SIZE_T Count +) +{ + DEBUG_ENTER_FUNCTION("ObjectArray=0x%p; Count=%Iu", ObjectArray, Count); + + InterlockedExchange(&_terminate, 1); + for (SIZE_T i = 0; i < Count; ++i) { + KeWaitForSingleObject(ObjectArray[i], Executive, KernelMode, FALSE, NULL); + ObDereferenceObject(ObjectArray[i]); + } + + DEBUG_EXIT_FUNCTION_VOID(); + return; +} + + +static +NTSTATUS +_CreateThreadGroup( + _In_ SIZE_T Count, + _In_ PKSTART_ROUTINE Routine, + _Out_ PETHREAD *ObjectArray +) +{ + CLIENT_ID clientId; + OBJECT_ATTRIBUTES oa; + HANDLE hThread = NULL; + NTSTATUS Status = STATUS_UNSUCCESSFUL; +#pragma warning(push) +#pragma warning (disable : 4152) + DEBUG_ENTER_FUNCTION("Count=%Iu; Routine=0x%p; ObjectArray=0x%p", Count, Routine, ObjectArray); +#pragma warning (pop) + + InitializeObjectAttributes(&oa, NULL, OBJ_KERNEL_HANDLE, NULL, NULL); + for (SIZE_T i = 0; i < Count; ++i) { + Status = PsCreateSystemThread(&hThread, SYNCHRONIZE, &oa, NULL, &clientId, Routine, (PVOID)i); + if (NT_SUCCESS(Status)) { + Status = ObReferenceObjectByHandle(hThread, SYNCHRONIZE, *PsThreadType, KernelMode, ObjectArray + i, NULL); + ZwClose(hThread); + } + + if (!NT_SUCCESS(Status)) { + _DestroyThreadGroup(ObjectArray, i); + break; + } + } + + DEBUG_EXIT_FUNCTION("0x%x", Status); + return Status; +} + + +static +NTSTATUS +_SocketTestThreadCreate( + _In_ PWSK_SOCKET Socket, + _In_ PLIST_ENTRY ListHead, + _In_ PKSPIN_LOCK ListLock +) +{ + KIRQL irql; + CLIENT_ID clientId; + OBJECT_ATTRIBUTES oa; + HANDLE hThread = NULL; + PTEST_THREAD_CONTEXT ctx = NULL; + NTSTATUS Status = STATUS_UNSUCCESSFUL; + DEBUG_ENTER_FUNCTION("Socket=0x%p; ListHead=0x%p; ListLock=0x%p", Socket, ListHead, ListLock); + + ctx = ExAllocatePoolUninitialized(NonPagedPool, sizeof(TEST_THREAD_CONTEXT), VIOWSK_TEST_TAG); + if (!ctx) { + Status = STATUS_INSUFFICIENT_RESOURCES; + goto Exit; + } + + InterlockedExchange(&ctx->Terminated, 0); + InitializeListHead(&ctx->Entry); + ctx->ListLock = ListLock; + ctx->Socket = Socket; + ctx->Irp = IoAllocateIrp(1, FALSE); + if (!ctx->Irp) { + Status = STATUS_INSUFFICIENT_RESOURCES; + goto FreeCtx; + } + + KeAcquireSpinLock(ListLock, &irql); + InsertTailList(ListHead, &ctx->Entry); + KeReleaseSpinLock(ListLock, irql); + InitializeObjectAttributes(&oa, NULL, OBJ_KERNEL_HANDLE, NULL, NULL); + Status = PsCreateSystemThread(&hThread, SYNCHRONIZE, &oa, NULL, &clientId, _TestThreadRoutine, ctx); + if (!NT_SUCCESS(Status)) + goto FreeIrp;; + + Status = ObReferenceObjectByHandle(hThread, SYNCHRONIZE, *PsThreadType, KernelMode, &ctx->Thread, NULL); + if (!NT_SUCCESS(Status)) { + InterlockedExchange(&ctx->Terminated, 1); + ZwWaitForSingleObject(hThread, FALSE, NULL); + goto CloseThread; + } + + ctx = NULL; +CloseThread: + ZwClose(hThread); + hThread = NULL; +FreeIrp: + if (hThread) + { + KeAcquireSpinLock(ListLock, &irql); + RemoveEntryList(&ctx->Entry); + InitializeListHead(&ctx->Entry); + KeReleaseSpinLock(ListLock, irql); + } + + if (ctx && ctx->Irp) + IoFreeIrp(ctx->Irp); +FreeCtx: + if (ctx) + ExFreePoolWithTag(ctx, VIOWSK_TEST_TAG); +Exit: + DEBUG_EXIT_FUNCTION("0x%x", Status); + return Status; +} + + +static +void +DriverUnload( + _In_ PDRIVER_OBJECT DriverObject +) +{ + DEBUG_ENTER_FUNCTION("DriverObject=0x%p", DriverObject); + + UNREFERENCED_PARAMETER(DriverObject); + + _DestroyThreadGroup(_clientThreads, sizeof(_clientThreads) / sizeof(_clientThreads[0])); + _DestroyThreadGroup(_serverThreads, sizeof(_serverThreads) / sizeof(_serverThreads[0])); + VioWskReleaseProviderNPI(&_vioWskRegistration); + VioWskDeregister(&_vioWskRegistration); + VioWskModuleFinit(); + IoDeleteDevice(_shutdownDeviceObject); + + DEBUG_EXIT_FUNCTION_VOID(); + return; +} + + +NTSTATUS +DriverEntry( + _In_ PDRIVER_OBJECT DriverObject, + _In_ PUNICODE_STRING RegistryPath +) +{ + PDEVICE_OBJECT Device = NULL; + NTSTATUS Status = STATUS_UNSUCCESSFUL; + DEBUG_ENTER_FUNCTION("DriverObject=0x%p; RegistryPath=\"%wZ\"", DriverObject, RegistryPath); + + Status = IoCreateDevice(DriverObject, 0, NULL, FILE_DEVICE_UNKNOWN, 0, FALSE, &Device); + if (!NT_SUCCESS(Status)) + goto Exit; + + Status = VioWskModuleInit(DriverObject, RegistryPath, Device); + if (!NT_SUCCESS(Status)) + goto DeleteDevice; + + Status = VioWskRegister(&_vioWskClientNPI, &_vioWskRegistration); + if (!NT_SUCCESS(Status)) + goto VioWskFinit; + + Status = VioWskCaptureProviderNPI(&_vioWskRegistration, WSK_INFINITE_WAIT, &_vioWskProviderNPI); + if (!NT_SUCCESS(Status)) + goto VioWskDeregister; + + KeInitializeEvent(&_initEvent, NotificationEvent, FALSE); + Status = _CreateThreadGroup(sizeof(_serverThreads) / sizeof(_serverThreads[0]), _ServerThreadRoutine, _serverThreads); + if (!NT_SUCCESS(Status)) + goto VIoWskReleaseNPI; + + Status = _CreateThreadGroup(sizeof(_clientThreads) / sizeof(_clientThreads[0]), _ClientThreadRoutine, _clientThreads); + if (!NT_SUCCESS(Status)) + goto DestroyServers; + + KeWaitForSingleObject(&_initEvent, Executive, KernelMode, FALSE, NULL); + if (_terminate) { + Status = STATUS_UNSUCCESSFUL; + goto DestroyClients; + } + + DriverObject->DriverUnload = DriverUnload; + _shutdownDeviceObject = (PDEVICE_OBJECT)InterlockedExchangePointer(&Device, NULL); +DestroyClients: + if (!NT_SUCCESS(Status)) + _DestroyThreadGroup(_clientThreads, sizeof(_clientThreads) / sizeof(_clientThreads[0])); +DestroyServers: + if (!NT_SUCCESS(Status)) + _DestroyThreadGroup(_serverThreads, sizeof(_serverThreads) / sizeof(_serverThreads[0])); +VIoWskReleaseNPI: + if (!NT_SUCCESS(Status)) + VioWskReleaseProviderNPI(&_vioWskRegistration); +VioWskDeregister: + if (!NT_SUCCESS(Status)) + VioWskDeregister(&_vioWskRegistration); +VioWskFinit: + if (!NT_SUCCESS(Status)) + VioWskModuleFinit(); +DeleteDevice: + if (Device) + IoDeleteDevice(Device); +Exit: + DEBUG_EXIT_FUNCTION("0x%x", Status); + return Status; +} \ No newline at end of file diff --git a/viosock/viosock.sln b/viosock/viosock.sln index d27453e5a..80fd45307 100644 --- a/viosock/viosock.sln +++ b/viosock/viosock.sln @@ -1,6 +1,6 @@ Microsoft Visual Studio Solution File, Format Version 12.00 -# Visual Studio Version 16 -VisualStudioVersion = 16.0.31702.278 +# Visual Studio Version 17 +VisualStudioVersion = 17.8.34316.72 MinimumVisualStudioVersion = 10.0.40219.1 Project("{8BC9CEB8-8B4A-11D0-8D11-00A0C91BC942}") = "VirtioLib", "..\VirtIO\VirtioLib.vcxproj", "{01D87C47-437A-4A16-8FD9-33FA5C99339E}" EndProject @@ -44,6 +44,8 @@ Project("{8BC9CEB8-8B4A-11D0-8D11-00A0C91BC942}") = "ViosockPackage", "ViosockPa EndProject Project("{8BC9CEB8-8B4A-11D0-8D11-00A0C91BC942}") = "viosockwsk", "wsk\wsk.vcxproj", "{96FDD976-0035-4E24-A61B-E93BED675101}" EndProject +Project("{8BC9CEB8-8B4A-11D0-8D11-00A0C91BC942}") = "viosockwsk-test", "viosock-wsk-test\viosock-wsk-test.vcxproj", "{A2A7AFFC-BEFE-4066-A1DD-23242FAFAD70}" +EndProject Global GlobalSection(SolutionConfigurationPlatforms) = preSolution Win10 Release|ARM64 = Win10 Release|ARM64 @@ -152,6 +154,24 @@ Global {96FDD976-0035-4E24-A61B-E93BED675101}.Win11 Release|x64.Build.0 = Win11 Release|x64 {96FDD976-0035-4E24-A61B-E93BED675101}.Win11 Release|x64.Deploy.0 = Win11 Release|x64 {96FDD976-0035-4E24-A61B-E93BED675101}.Win11 Release|x86.ActiveCfg = Win11 Release|Win32 + {A2A7AFFC-BEFE-4066-A1DD-23242FAFAD70}.Win10 Release|ARM64.ActiveCfg = Win10 Release|ARM64 + {A2A7AFFC-BEFE-4066-A1DD-23242FAFAD70}.Win10 Release|ARM64.Build.0 = Win10 Release|ARM64 + {A2A7AFFC-BEFE-4066-A1DD-23242FAFAD70}.Win10 Release|ARM64.Deploy.0 = Win10 Release|ARM64 + {A2A7AFFC-BEFE-4066-A1DD-23242FAFAD70}.Win10 Release|x64.ActiveCfg = Win10 Release|x64 + {A2A7AFFC-BEFE-4066-A1DD-23242FAFAD70}.Win10 Release|x64.Build.0 = Win10 Release|x64 + {A2A7AFFC-BEFE-4066-A1DD-23242FAFAD70}.Win10 Release|x64.Deploy.0 = Win10 Release|x64 + {A2A7AFFC-BEFE-4066-A1DD-23242FAFAD70}.Win10 Release|x86.ActiveCfg = Win10 Release|Win32 + {A2A7AFFC-BEFE-4066-A1DD-23242FAFAD70}.Win10 Release|x86.Build.0 = Win10 Release|Win32 + {A2A7AFFC-BEFE-4066-A1DD-23242FAFAD70}.Win10 Release|x86.Deploy.0 = Win10 Release|Win32 + {A2A7AFFC-BEFE-4066-A1DD-23242FAFAD70}.Win11 Release|ARM64.ActiveCfg = Win11 Release|ARM64 + {A2A7AFFC-BEFE-4066-A1DD-23242FAFAD70}.Win11 Release|ARM64.Build.0 = Win11 Release|ARM64 + {A2A7AFFC-BEFE-4066-A1DD-23242FAFAD70}.Win11 Release|ARM64.Deploy.0 = Win11 Release|ARM64 + {A2A7AFFC-BEFE-4066-A1DD-23242FAFAD70}.Win11 Release|x64.ActiveCfg = Win11 Release|x64 + {A2A7AFFC-BEFE-4066-A1DD-23242FAFAD70}.Win11 Release|x64.Build.0 = Win11 Release|x64 + {A2A7AFFC-BEFE-4066-A1DD-23242FAFAD70}.Win11 Release|x64.Deploy.0 = Win11 Release|x64 + {A2A7AFFC-BEFE-4066-A1DD-23242FAFAD70}.Win11 Release|x86.ActiveCfg = Win11 Release|Win32 + {A2A7AFFC-BEFE-4066-A1DD-23242FAFAD70}.Win11 Release|x86.Build.0 = Win11 Release|Win32 + {A2A7AFFC-BEFE-4066-A1DD-23242FAFAD70}.Win11 Release|x86.Deploy.0 = Win11 Release|Win32 EndGlobalSection GlobalSection(SolutionProperties) = preSolution HideSolutionNode = FALSE From a831488c871369efac99c67e72799cb8451a9235 Mon Sep 17 00:00:00 2001 From: Martin Drab Date: Thu, 18 May 2023 06:26:56 -0700 Subject: [PATCH 18/18] Enable WPP Event Tracing for both the WSK library and the test driver The WPP Event Tracing can be disabled which means return to the good old DbgPrintEx routine. To disable WPP, the following steps need to be done: - commenting the line `#define EVENT_TRACING` in project's trace.h file, - setting the "Run Wpp Tracing" option in project settings to No. Signed-off-by: Martin Drab --- viosock/inc/debug-utils.h | 8 ++ viosock/viosock-wsk-test/test-messages.c | 3 + .../viosock-wsk-test/viosock-wsk-test.vcxproj | 123 +++++++++++++++++- .../viosock-wsk-test.vcxproj.filters | 3 + viosock/viosock-wsk-test/viosockwsk-test.c | 11 +- viosock/viosock-wsk-test/wpp-trace.h | 117 +++++++++++++++++ viosock/wsk/provider.c | 8 +- viosock/wsk/socket-internal.c | 5 + viosock/wsk/socket.c | 7 +- viosock/wsk/viowsk.c | 4 + viosock/wsk/wpp-trace.h | 113 ++++++++++++++++ viosock/wsk/wsk-completion.c | 6 +- viosock/wsk/wsk-utils.c | 6 +- viosock/wsk/wsk-workitem.c | 3 + viosock/wsk/wsk.vcxproj | 29 ++++- viosock/wsk/wsk.vcxproj.filters | 3 + 16 files changed, 437 insertions(+), 12 deletions(-) create mode 100644 viosock/viosock-wsk-test/wpp-trace.h create mode 100644 viosock/wsk/wpp-trace.h diff --git a/viosock/inc/debug-utils.h b/viosock/inc/debug-utils.h index ce8c319c1..75e79fa8a 100644 --- a/viosock/inc/debug-utils.h +++ b/viosock/inc/debug-utils.h @@ -31,6 +31,12 @@ #define __DEBUG_UTILS_H__ +#include +#include "wpp-trace.h" +#include "..\..\virtio\kdebugprint.h" + +#ifndef EVENT_TRACING + #if defined(DBG) || defined(_DEBUG) @@ -81,3 +87,5 @@ #endif + +#endif diff --git a/viosock/viosock-wsk-test/test-messages.c b/viosock/viosock-wsk-test/test-messages.c index b1c196810..ccf3398e7 100644 --- a/viosock/viosock-wsk-test/test-messages.c +++ b/viosock/viosock-wsk-test/test-messages.c @@ -30,6 +30,9 @@ #include #include "..\inc\debug-utils.h" #include "test-messages.h" +#ifdef EVENT_TRACING +#include "test-messages.tmh" +#endif static ULONG _randSeed; diff --git a/viosock/viosock-wsk-test/viosock-wsk-test.vcxproj b/viosock/viosock-wsk-test/viosock-wsk-test.vcxproj index 4b35c9115..6ce9f8232 100644 --- a/viosock/viosock-wsk-test/viosock-wsk-test.vcxproj +++ b/viosock/viosock-wsk-test/viosock-wsk-test.vcxproj @@ -13,6 +13,18 @@ Win10 Release ARM64 + + Win11 Release + ARM64 + + + Win11 Release + Win32 + + + Win11 Release + x64 + {A2A7AFFC-BEFE-4066-A1DD-23242FAFAD70} @@ -34,6 +46,14 @@ WDM Desktop + + Windows10 + false + WindowsKernelModeDriver10.0 + Driver + WDM + Desktop + Windows10 false @@ -42,6 +62,14 @@ WDM Desktop + + Windows10 + false + WindowsKernelModeDriver10.0 + Driver + WDM + Desktop + Windows10 false @@ -50,6 +78,14 @@ WDM Desktop + + Windows10 + false + WindowsKernelModeDriver10.0 + Driver + WDM + Desktop + @@ -64,31 +100,111 @@ objfre_win10_x86\i386\ objfre_win10_x86\i386\ + + DbgengKernelDebugger + false + objfre_win11_x86\i386\ + objfre_win11_x86\i386\ + DbgengKernelDebugger false objfre_win10_amd64\amd64\ objfre_win10_amd64\amd64\ + + DbgengKernelDebugger + false + objfre_win11_amd64\amd64\ + objfre_win11_amd64\amd64\ + DbgengKernelDebugger objfre_win10_arm64\arm64\ objfre_win10_arm64\arm64\ + + DbgengKernelDebugger + objfre_win11_arm64\arm64\ + objfre_win11_arm64\arm64\ + - cng.lib;%(AdditionalDependencies);$(KernelBufferOverflowLib);$(DDK_LIB_PATH)ntoskrnl.lib;$(DDK_LIB_PATH)hal.lib;$(DDK_LIB_PATH)wmilib.lib + Ntstrsafe.lib;cng.lib;%(AdditionalDependencies);$(KernelBufferOverflowLib);$(DDK_LIB_PATH)ntoskrnl.lib;$(DDK_LIB_PATH)hal.lib;$(DDK_LIB_PATH)wmilib.lib + + + true + wpp-trace.h + TRACE_PROJECT_NAME=VioWskTest;_X86_=1;i386=1;STD_CALL;%(PreprocessorDefinitions) + + + SHA256 + + + + + Ntstrsafe.lib;cng.lib;%(AdditionalDependencies);$(KernelBufferOverflowLib);$(DDK_LIB_PATH)ntoskrnl.lib;$(DDK_LIB_PATH)hal.lib;$(DDK_LIB_PATH)wmilib.lib + + true + wpp-trace.h + TRACE_PROJECT_NAME=VioWskTest;_X86_=1;i386=1;STD_CALL;%(PreprocessorDefinitions) + + + SHA256 + - cng.lib;%(AdditionalDependencies);$(KernelBufferOverflowLib);$(DDK_LIB_PATH)ntoskrnl.lib;$(DDK_LIB_PATH)hal.lib;$(DDK_LIB_PATH)wmilib.lib + Ntstrsafe.lib;cng.lib;%(AdditionalDependencies);$(KernelBufferOverflowLib);$(DDK_LIB_PATH)ntoskrnl.lib;$(DDK_LIB_PATH)hal.lib;$(DDK_LIB_PATH)wmilib.lib + + true + wpp-trace.h + TRACE_PROJECT_NAME=VioWskTest;_WIN64;_AMD64_;AMD64;%(PreprocessorDefinitions) + + + SHA256 + + + + + Ntstrsafe.lib;cng.lib;%(AdditionalDependencies);$(KernelBufferOverflowLib);$(DDK_LIB_PATH)ntoskrnl.lib;$(DDK_LIB_PATH)hal.lib;$(DDK_LIB_PATH)wmilib.lib + + + true + wpp-trace.h + TRACE_PROJECT_NAME=VioWskTest;_WIN64;_AMD64_;AMD64;%(PreprocessorDefinitions) + + + SHA256 + - cng.lib;%(AdditionalDependencies);$(KernelBufferOverflowLib);$(DDK_LIB_PATH)ntoskrnl.lib;$(DDK_LIB_PATH)hal.lib;$(DDK_LIB_PATH)wmilib.lib + Ntstrsafe.lib;cng.lib;%(AdditionalDependencies);$(KernelBufferOverflowLib);$(DDK_LIB_PATH)ntoskrnl.lib;$(DDK_LIB_PATH)hal.lib;$(DDK_LIB_PATH)wmilib.lib + + + true + wpp-trace.h + TRACE_PROJECT_NAME=VioWskTest;_ARM64_;ARM64;_USE_DECLSPECS_FOR_SAL=1;STD_CALL;%(PreprocessorDefinitions) + + + SHA256 + + + + + Ntstrsafe.lib;cng.lib;%(AdditionalDependencies);$(KernelBufferOverflowLib);$(DDK_LIB_PATH)ntoskrnl.lib;$(DDK_LIB_PATH)hal.lib;$(DDK_LIB_PATH)wmilib.lib + + true + wpp-trace.h + TRACE_PROJECT_NAME=VioWskTest;_ARM64_;ARM64;_USE_DECLSPECS_FOR_SAL=1;STD_CALL;%(PreprocessorDefinitions) + + + SHA256 + @@ -104,6 +220,7 @@ + diff --git a/viosock/viosock-wsk-test/viosock-wsk-test.vcxproj.filters b/viosock/viosock-wsk-test/viosock-wsk-test.vcxproj.filters index c745a968c..9a2656595 100644 --- a/viosock/viosock-wsk-test/viosock-wsk-test.vcxproj.filters +++ b/viosock/viosock-wsk-test/viosock-wsk-test.vcxproj.filters @@ -26,5 +26,8 @@ Header Files + + Header Files + \ No newline at end of file diff --git a/viosock/viosock-wsk-test/viosockwsk-test.c b/viosock/viosock-wsk-test/viosockwsk-test.c index b123358ba..ac84c3d6c 100644 --- a/viosock/viosock-wsk-test/viosockwsk-test.c +++ b/viosock/viosock-wsk-test/viosockwsk-test.c @@ -52,7 +52,9 @@ #include "..\inc\vio_sockets.h" #include "..\sys\public.h" #include "test-messages.h" - +#ifdef EVENT_TRACING +#include "viosockwsk-test.tmh" +#endif #define LISTEN_PORT_MIN 1337 @@ -797,6 +799,7 @@ DriverUnload( VioWskDeregister(&_vioWskRegistration); VioWskModuleFinit(); IoDeleteDevice(_shutdownDeviceObject); + WPP_CLEANUP(DriverObject); DEBUG_EXIT_FUNCTION_VOID(); return; @@ -813,6 +816,7 @@ DriverEntry( NTSTATUS Status = STATUS_UNSUCCESSFUL; DEBUG_ENTER_FUNCTION("DriverObject=0x%p; RegistryPath=\"%wZ\"", DriverObject, RegistryPath); + WPP_INIT_TRACING(DriverObject, RegistryPath); Status = IoCreateDevice(DriverObject, 0, NULL, FILE_DEVICE_UNKNOWN, 0, FALSE, &Device); if (!NT_SUCCESS(Status)) goto Exit; @@ -865,6 +869,11 @@ DriverEntry( if (Device) IoDeleteDevice(Device); Exit: + if (!NT_SUCCESS(Status)) + { + WPP_CLEANUP(DriverObject); + } + DEBUG_EXIT_FUNCTION("0x%x", Status); return Status; } \ No newline at end of file diff --git a/viosock/viosock-wsk-test/wpp-trace.h b/viosock/viosock-wsk-test/wpp-trace.h new file mode 100644 index 000000000..8673c69c5 --- /dev/null +++ b/viosock/viosock-wsk-test/wpp-trace.h @@ -0,0 +1,117 @@ +/* + * Copyright (c) 2023 Virtuozzo International GmbH + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions + * are met : + * 1. Redistributions of source code must retain the above copyright + * notice, this list of conditions and the following disclaimer. + * 2. Redistributions in binary form must reproduce the above copyright + * notice, this list of conditions and the following disclaimer in the + * documentation and / or other materials provided with the distribution. + * 3. Neither the names of the copyright holders nor the names of their contributors + * may be used to endorse or promote products derived from this software + * without specific prior written permission. + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS ``AS IS'' AND + * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + * ARE DISCLAIMED.IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS + * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) + * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT + * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY + * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF + * SUCH DAMAGE. + */ + + + +#define EVENT_TRACING + +#if !defined(EVENT_TRACING) + +#if !defined(TRACE_LEVEL_NONE) + #define TRACE_LEVEL_NONE 0 + #define TRACE_LEVEL_CRITICAL 1 + #define TRACE_LEVEL_FATAL 1 + #define TRACE_LEVEL_ERROR 2 + #define TRACE_LEVEL_WARNING 3 + #define TRACE_LEVEL_INFORMATION 4 + #define TRACE_LEVEL_VERBOSE 5 + #define TRACE_LEVEL_RESERVED6 6 + #define TRACE_LEVEL_RESERVED7 7 + #define TRACE_LEVEL_RESERVED8 8 + #define TRACE_LEVEL_RESERVED9 9 +#endif + + +// +// Define Debug Flags +// +#define DBG_VIOWSK 0x00000001 + +#define DBG_TEST 0x00000001 + +#define WPP_INIT_TRACING(a,b) +#define WPP_CLEANUP(DriverObject) + +#else +#define WPP_CHECK_FOR_NULL_STRING + +// +// Define the tracing flags. +// +// Tracing GUID - C2D7F82F-CE5F-4408-8A37-8B9FE2B3D52E +// + +// {13b9cfb4-b962-4b43-b59d-92242fab52e3} +// {46e3298a-70b1-49c6-b9fd-8691980b7adf} +#define WPP_CONTROL_GUIDS \ + WPP_DEFINE_CONTROL_GUID(WskTraceGuid,(13b9cfb4,b962,4b43,b59d,92242fab52e3), \ + WPP_DEFINE_BIT(DBG_VIOWSK) /* bit 0 = 0x00000001 */ \ + ) \ + WPP_DEFINE_CONTROL_GUID(WskTestTraceGuid,(46e3298a,70b1,49c6,b9fd,8691980b7adf), \ + WPP_DEFINE_BIT(DBG_TEST) /* bit 0 = 0x00000001 */ \ + ) + + +#define WPP_FLAG_LEVEL_LOGGER(flag, level) \ + WPP_LEVEL_LOGGER(flag) + +#define WPP_FLAG_LEVEL_ENABLED(flag, level) \ + (WPP_LEVEL_ENABLED(flag) && WPP_CONTROL(WPP_BIT_ ## flag).Level >= level) + +#define WPP_LEVEL_FLAGS_LOGGER(lvl,flags) \ + WPP_LEVEL_LOGGER(flags) + +#define WPP_LEVEL_FLAGS_ENABLED(lvl, flags) \ + (WPP_LEVEL_ENABLED(flags) && WPP_CONTROL(WPP_BIT_ ## flags).Level >= lvl) + + // + // This comment block is scanned by the trace preprocessor to define our + // Trace function. + // + // begin_wpp config + // + // USEPREFIX(DEBUG_ENTER_FUNCTION, "%!STDPREFIX! %!FUNC!("); + // USESUFFIX(DEBUG_ENTER_FUNCTION, ")"); + // USEPREFIX(DEBUG_ENTER_FUNCTION_NO_ARGS, "%!STDPREFIX! %!FUNC!()"); + // USEPREFIX(DEBUG_EXIT_FUNCTION, "%!STDPREFIX! %!FUNC!(-)"); + // USEPREFIX(DEBUG_EXIT_FUNCTION_VOID, "%!STDPREFIX! %!FUNC!"); + // USESUFFIX(DEBUG_EXIT_FUNCTION_VOID, "(-)"); + // + // FUNC TraceEvents(LEVEL, FLAGS, MSG, ...); + // FUNC DEBUG_ENTER_FUNCTION{LEVEL=TRACE_LEVEL_VERBOSE, FLAGS=DBG_TEST}(MSG, ...); + // FUNC DEBUG_ENTER_FUNCTION_NO_ARGS{LEVEL=TRACE_LEVEL_VERBOSE, FLAGS=DBG_TEST}(); + // FUNC DEBUG_EXIT_FUNCTION{LEVEL=TRACE_LEVEL_VERBOSE, FLAGS=DBG_TEST}(MSG, ...); + // FUNC DEBUG_EXIT_FUNCTION_VOID{LEVEL=TRACE_LEVEL_VERBOSE, FLAGS=DBG_TEST}(); + // FUNC DEBUG_ERROR{LEVEL=TRACE_LEVEL_ERROR, FLAGS=DBG_TEST}(MSG, ...); + // FUNC DEBUG_WARNING{LEVEL=TRACE_LEVEL_WARNING, FLAGS=DBG_TEST}(MSG, ...); + // FUNC DEBUG_TRACE{LEVEL=TRACE_LEVEL_VERBOSE, FLAGS=DBG_TEST}(MSG, ...); + // FUNC DEBUG_INFO{LEVEL=TRACE_LEVEL_INFORMATION, FLAGS=DBG_TEST}(MSG, ...); + // + // end_wpp + // + +#endif diff --git a/viosock/wsk/provider.c b/viosock/wsk/provider.c index ccc61c7e6..b2ab13e19 100644 --- a/viosock/wsk/provider.c +++ b/viosock/wsk/provider.c @@ -34,6 +34,10 @@ #include "viowsk-internal.h" #include "wsk-workitem.h" #include "..\inc\vio_wsk.h" +#ifdef EVENT_TRACING +#include "provider.tmh" +#endif + #ifdef ALLOC_PRAGMA #pragma alloc_text (PAGE, VioWskGetAddressInfo) @@ -61,7 +65,7 @@ VioWskSocket( PWSK_WORKITEM WorkItem = NULL; PVIOWSK_SOCKET pSocket = NULL; NTSTATUS Status = STATUS_UNSUCCESSFUL; - DEBUG_ENTER_FUNCTION("Client=0x%p; AddressFamily=%u; SocketType=%u; Protocol=%u; Flags=0x%x; SocketContext=0x%p; Dispatch=0x%p; OwningProcess=0x%p; OwningThread=0x%p; SecurityDescriptor=0x%p; Irp=0x%p", Client, AddressFamily, SocketType, Protocol, Flags, SocketContext, Dispatch, OwningProcess, OwningThread, SecurityDescriptor, Irp); + DEBUG_ENTER_FUNCTION("Client=0x%p; AddressFamily=%u; SocketType=%u; Protocol=%u; Flags=0x%x; SocketContext=0x%p; Dispatch=0x%p; OwningProcess=0x%p; OwningThread=0x%p; SecurityDescriptor=0x%p; Irp=0x%p", Client, AddressFamily, SocketType, Protocol, Flags, SocketContext, Dispatch, OwningProcess, OwningThread, SecurityDescriptor, Irp); _At_((void*)Irp->IoStatus.Information, __drv_allocatesMem(Mem)) @@ -151,7 +155,7 @@ VioWskControlClient( ) { NTSTATUS Status = STATUS_UNSUCCESSFUL; - DEBUG_ENTER_FUNCTION("Client=0x%p; ControlCode=0x%x; InputSize=%zu; InputBuffer=0x%p; OutputSize=%zu; OutputBuffer=0x%p; OutputSizeReturned=0x%p; Irp=0x%p", Client, ControlCode, InputSize, InputBuffer, OutputSize, OutputBuffer, OutputSizeReturned, Irp); + DEBUG_ENTER_FUNCTION("Client=0x%p; ControlCode=0x%x; InputSize=%Iu; InputBuffer=0x%p; OutputSize=%Iu; OutputBuffer=0x%p; OutputSizeReturned=0x%p; Irp=0x%p", Client, ControlCode, InputSize, InputBuffer, OutputSize, OutputBuffer, OutputSizeReturned, Irp); UNREFERENCED_PARAMETER(Client); UNREFERENCED_PARAMETER(ControlCode); diff --git a/viosock/wsk/socket-internal.c b/viosock/wsk/socket-internal.c index 737bf8e29..874f87580 100644 --- a/viosock/wsk/socket-internal.c +++ b/viosock/wsk/socket-internal.c @@ -32,6 +32,9 @@ #include "..\inc\debug-utils.h" #include "wsk-utils.h" #include "viowsk-internal.h" +#ifdef EVENT_TRACING +#include "socket-internal.tmh" +#endif #ifdef ALLOC_PRAGMA #pragma alloc_text (PAGE, VioWskSocketInternal) @@ -39,6 +42,8 @@ #endif + + _Must_inspect_result_ static NTSTATUS diff --git a/viosock/wsk/socket.c b/viosock/wsk/socket.c index 828378d1d..669e2ba4d 100644 --- a/viosock/wsk/socket.c +++ b/viosock/wsk/socket.c @@ -35,6 +35,11 @@ #include "wsk-completion.h" #include "wsk-workitem.h" #include "..\inc\vio_wsk.h" +#ifdef EVENT_TRACING +#include "socket.tmh" +#endif + + NTSTATUS WSKAPI @@ -293,7 +298,7 @@ VioWskControlSocket( PVIOSOCKET_COMPLETION_CONTEXT CompContext = NULL; NTSTATUS Status = STATUS_UNSUCCESSFUL; PVIOWSK_SOCKET pSocket = CONTAINING_RECORD(Socket, VIOWSK_SOCKET, WskSocket); - DEBUG_ENTER_FUNCTION("Socket=0x%p; RequestType=%u; ControlCode=0x%x; Level=%u; InputSize=%zu; InputBuffer=0x%p; OutputSize=%zu; OutputBuffer=0x%p; OutputSizeReturned=0x%p; Irp=0x%p", Socket, RequestType, ControlCode, Level, InputSize, InputBuffer, OutputSize, OutputBuffer, OutputSizeReturned, Irp); + DEBUG_ENTER_FUNCTION("Socket=0x%p; RequestType=%u; ControlCode=0x%x; Level=%u; InputSize=%Iu; InputBuffer=0x%p; OutputSize=%Iu; OutputBuffer=0x%p; OutputSizeReturned=0x%p; Irp=0x%p", Socket, RequestType, ControlCode, Level, InputSize, InputBuffer, OutputSize, OutputBuffer, OutputSizeReturned, Irp); UNREFERENCED_PARAMETER(OutputSizeReturned); diff --git a/viosock/wsk/viowsk.c b/viosock/wsk/viowsk.c index 012d67499..4e6c8f10b 100644 --- a/viosock/wsk/viowsk.c +++ b/viosock/wsk/viowsk.c @@ -32,6 +32,9 @@ #include "viowsk.h" #include "..\inc\vio_wsk.h" #include "viowsk-internal.h" +#ifdef EVENT_TRACING +#include "viowsk.tmh" +#endif #ifdef ALLOC_PRAGMA #pragma alloc_text (PAGE, VioWskRegister) @@ -39,6 +42,7 @@ #endif + typedef struct _VIOSOCK_WAIT_CONTEXT { UNICODE_STRING SymbolicLinkName; KEVENT Event; diff --git a/viosock/wsk/wpp-trace.h b/viosock/wsk/wpp-trace.h new file mode 100644 index 000000000..54fe6cb3e --- /dev/null +++ b/viosock/wsk/wpp-trace.h @@ -0,0 +1,113 @@ +/* + * This file contains trace and debugging related definitions + * + * Copyright (c) 2019 Virtuozzo International GmbH + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions + * are met : + * 1. Redistributions of source code must retain the above copyright + * notice, this list of conditions and the following disclaimer. + * 2. Redistributions in binary form must reproduce the above copyright + * notice, this list of conditions and the following disclaimer in the + * documentation and / or other materials provided with the distribution. + * 3. Neither the names of the copyright holders nor the names of their contributors + * may be used to endorse or promote products derived from this software + * without specific prior written permission. + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS ``AS IS'' AND + * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + * ARE DISCLAIMED.IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS + * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) + * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT + * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY + * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF + * SUCH DAMAGE. + */ + + #pragma once + +#define EVENT_TRACING + +#if !defined(EVENT_TRACING) + +#if !defined(TRACE_LEVEL_NONE) +#define TRACE_LEVEL_NONE 0 +#define TRACE_LEVEL_CRITICAL 1 +#define TRACE_LEVEL_FATAL 1 +#define TRACE_LEVEL_ERROR 2 +#define TRACE_LEVEL_WARNING 3 +#define TRACE_LEVEL_INFORMATION 4 +#define TRACE_LEVEL_VERBOSE 5 +#define TRACE_LEVEL_RESERVED6 6 +#define TRACE_LEVEL_RESERVED7 7 +#define TRACE_LEVEL_RESERVED8 8 +#define TRACE_LEVEL_RESERVED9 9 +#endif + + +// +// Define Debug Flags +// +#define DBG_VIOWSK 0x00000001 + +#define WPP_INIT_TRACING(a,b) +#define WPP_CLEANUP(DriverObject) + +#else + +#define WPP_CHECK_FOR_NULL_STRING + +// +// Define the tracing flags. +// +// Tracing GUID - C2D7F82F-CE5F-4408-8A37-8B9FE2B3D52E +// + +// {13b9cfb4-b962-4b43-b59d-92242fab52e3} +#define WPP_CONTROL_GUIDS \ + WPP_DEFINE_CONTROL_GUID(WskTraceGuid,(13b9cfb4,b962,4b43,b59d,92242fab52e3), \ + WPP_DEFINE_BIT(DBG_VIOWSK) /* bit 0 = 0x00000001 */ \ + ) + +#define WPP_FLAG_LEVEL_LOGGER(flag, level) \ + WPP_LEVEL_LOGGER(flag) + +#define WPP_FLAG_LEVEL_ENABLED(flag, level) \ + (WPP_LEVEL_ENABLED(flag) && WPP_CONTROL(WPP_BIT_ ## flag).Level >= level) + +#define WPP_LEVEL_FLAGS_LOGGER(lvl,flags) \ + WPP_LEVEL_LOGGER(flags) + +#define WPP_LEVEL_FLAGS_ENABLED(lvl, flags) \ + (WPP_LEVEL_ENABLED(flags) && WPP_CONTROL(WPP_BIT_ ## flags).Level >= lvl) + + // + // This comment block is scanned by the trace preprocessor to define our + // Trace function. + // + // begin_wpp config + // + // USEPREFIX(DEBUG_ENTER_FUNCTION, "%!STDPREFIX! %!FUNC!("); + // USESUFFIX(DEBUG_ENTER_FUNCTION, ")"); + // USEPREFIX(DEBUG_ENTER_FUNCTION_NO_ARGS, "%!STDPREFIX! %!FUNC!()"); + // USEPREFIX(DEBUG_EXIT_FUNCTION, "%!STDPREFIX! %!FUNC!(-)"); + // USEPREFIX(DEBUG_EXIT_FUNCTION_VOID, "%!STDPREFIX! %!FUNC!"); + // USESUFFIX(DEBUG_EXIT_FUNCTION_VOID, "(-)"); + // + // FUNC TraceEvents(LEVEL, FLAGS, MSG, ...); + // FUNC DEBUG_ENTER_FUNCTION{LEVEL=TRACE_LEVEL_VERBOSE, FLAGS=DBG_VIOWSK}(MSG, ...); + // FUNC DEBUG_ENTER_FUNCTION_NO_ARGS{LEVEL=TRACE_LEVEL_VERBOSE, FLAGS=DBG_VIOWSK}(); + // FUNC DEBUG_EXIT_FUNCTION{LEVEL=TRACE_LEVEL_VERBOSE, FLAGS=DBG_VIOWSK}(MSG, ...); + // FUNC DEBUG_EXIT_FUNCTION_VOID{LEVEL=TRACE_LEVEL_VERBOSE, FLAGS=DBG_VIOWSK}(); + // FUNC DEBUG_ERROR{LEVEL=TRACE_LEVEL_ERROR, FLAGS=DBG_VIOWSK}(MSG, ...); + // FUNC DEBUG_WARNING{LEVEL=TRACE_LEVEL_WARNING, FLAGS=DBG_VIOWSK}(MSG, ...); + // FUNC DEBUG_TRACE{LEVEL=TRACE_LEVEL_VERBOSE, FLAGS=DBG_VIOWSK}(MSG, ...); + // FUNC DEBUG_INFO{LEVEL=TRACE_LEVEL_INFORMATION, FLAGS=DBG_VIOWSK}(MSG, ...); + // + // end_wpp + // + +#endif diff --git a/viosock/wsk/wsk-completion.c b/viosock/wsk/wsk-completion.c index ce68c1f58..2477c06e3 100644 --- a/viosock/wsk/wsk-completion.c +++ b/viosock/wsk/wsk-completion.c @@ -32,7 +32,9 @@ #include "wsk-utils.h" #include "viowsk-internal.h" #include "wsk-completion.h" - +#ifdef EVENT_TRACING +#include "wsk-completion.tmh" +#endif @@ -230,7 +232,7 @@ WskGeneralIrpCompletion( WskCompContextDereference(Ctx); - DEBUG_EXIT_FUNCTION("0x%x", STATUS_MORE_PROCESSING_REQUIRED); + DEBUG_EXIT_FUNCTION("0x%ix", STATUS_MORE_PROCESSING_REQUIRED); return STATUS_MORE_PROCESSING_REQUIRED; } diff --git a/viosock/wsk/wsk-utils.c b/viosock/wsk/wsk-utils.c index 7d42deff2..64911e012 100644 --- a/viosock/wsk/wsk-utils.c +++ b/viosock/wsk/wsk-utils.c @@ -34,6 +34,9 @@ #include "viowsk-internal.h" #include "wsk-completion.h" #include "wsk-utils.h" +#ifdef EVENT_TRACING +#include "wsk-utils.tmh" +#endif #ifdef ALLOC_PRAGMA #endif @@ -89,6 +92,7 @@ VioWskIrpRelease( } + NTSTATUS VioWskIrpComplete( _Inout_opt_ PVIOWSK_SOCKET Socket, @@ -97,7 +101,7 @@ VioWskIrpComplete( _In_ ULONG_PTR Information ) { - DEBUG_ENTER_FUNCTION("Socket=0x%p; Irp=0x%p; Status=0x%x; Information=%zu", Socket, Irp, Status, Information); + DEBUG_ENTER_FUNCTION("Socket=0x%p; Irp=0x%p; Status=0x%x; Information=%Iu", Socket, Irp, Status, (ULONG64)Information); IoSetCancelRoutine(Irp, NULL); if (Socket) diff --git a/viosock/wsk/wsk-workitem.c b/viosock/wsk/wsk-workitem.c index 70223530a..caa9dfabc 100644 --- a/viosock/wsk/wsk-workitem.c +++ b/viosock/wsk/wsk-workitem.c @@ -34,6 +34,9 @@ #include "..\inc\vio_wsk.h" #include "viowsk-internal.h" #include "wsk-workitem.h" +#ifdef EVENT_TRACING +#include "wsk-workitem.tmh" +#endif #pragma warning(disable : 4996) diff --git a/viosock/wsk/wsk.vcxproj b/viosock/wsk/wsk.vcxproj index 463fc1814..381ff01ee 100644 --- a/viosock/wsk/wsk.vcxproj +++ b/viosock/wsk/wsk.vcxproj @@ -29,6 +29,7 @@ + @@ -125,12 +126,36 @@ - WINAPI_FAMILY=WINAPI_FAMILY_DESKTOP_APP;WINAPI_PARTITION_DESKTOP=1;WINAPI_PARTITION_SYSTEM=1;WINAPI_PARTITION_APP=1;WINAPI_PARTITION_PC_APP=1;%(PreprocessorDefinitions) + TRACE_PROJECT_NAME=VioWsk;WINAPI_FAMILY=WINAPI_FAMILY_DESKTOP_APP;WINAPI_PARTITION_DESKTOP=1;WINAPI_PARTITION_SYSTEM=1;WINAPI_PARTITION_APP=1;WINAPI_PARTITION_PC_APP=1;%(PreprocessorDefinitions) + true + true + true + true + true + wpp-trace.h + wpp-trace.h + trace.h + wpp-trace.h + trace.h + true + true + true + wpp-trace.h + wpp-trace.h + wpp-trace.h - _DEBUG;%(PreprocessorDefinitions) + TRACE_PROJECT_NAME=VioWsk;_DEBUG;%(PreprocessorDefinitions) + true + true + true + true + trace.h + trace.h + trace.h + trace.h diff --git a/viosock/wsk/wsk.vcxproj.filters b/viosock/wsk/wsk.vcxproj.filters index a2e1b70c1..1e60147a1 100644 --- a/viosock/wsk/wsk.vcxproj.filters +++ b/viosock/wsk/wsk.vcxproj.filters @@ -36,6 +36,9 @@ Header Files + + Header Files +