diff --git a/inc/usersim/ex.h b/inc/usersim/ex.h index cb11664..a78d1a5 100644 --- a/inc/usersim/ex.h +++ b/inc/usersim/ex.h @@ -184,6 +184,77 @@ extern "C" USERSIM_API void ExRaiseDatatypeMisalignment(); + /** + * @brief Allocate memory. + * @param[in] pool_type Pool type to use. + * @param[in] size Size of memory to allocate. + * @param[in] tag Pool tag to use. + * @param[in] initialize False to return "uninitialized" memory. + * @returns Pointer to memory block allocated, or null on failure. + */ + __drv_allocatesMem(Mem) _Must_inspect_result_ _Ret_writes_maybenull_(size) void* usersim_allocate_with_tag( + _In_ __drv_strictTypeMatch(__drv_typeExpr) POOL_TYPE pool_type, size_t size, uint32_t tag, bool initialize); + + /** + * @brief Allocate memory. + * @param[in] size Size of memory to allocate. + * @returns Pointer to memory block allocated, or null on failure. + */ + __drv_allocatesMem(Mem) _Must_inspect_result_ _Ret_writes_maybenull_(size) void* usersim_allocate(size_t size); + + /** + * @brief Reallocate memory. + * @param[in] memory Allocation to be reallocated. + * @param[in] old_size Old size of memory to reallocate. + * @param[in] new_size New size of memory to reallocate. + * @returns Pointer to memory block allocated, or null on failure. + */ + __drv_allocatesMem(Mem) _Must_inspect_result_ _Ret_writes_maybenull_(new_size) void* usersim_reallocate( + _In_ _Post_invalid_ void* memory, size_t old_size, size_t new_size); + + /** + * @brief Reallocate memory with tag. + * @param[in] memory Allocation to be reallocated. + * @param[in] old_size Old size of memory to reallocate. + * @param[in] new_size New size of memory to reallocate. + * @param[in] tag Pool tag to use. + * @returns Pointer to memory block allocated, or null on failure. + */ + __drv_allocatesMem(Mem) _Must_inspect_result_ _Ret_writes_maybenull_(new_size) void* usersim_reallocate_with_tag( + _In_ _Post_invalid_ void* memory, size_t old_size, size_t new_size, uint32_t tag); + + /** + * @brief Free memory. + * @param[in] memory Allocation to be freed. + */ + void + usersim_free(_Frees_ptr_opt_ void* memory); + + /** + * @brief Allocate memory that has a starting address that is cache aligned. + * @param[in] size Size of memory to allocate + * @returns Pointer to memory block allocated, or null on failure. + */ + USERSIM_API + __drv_allocatesMem(Mem) _Must_inspect_result_ + _Ret_writes_maybenull_(size) void* usersim_allocate_cache_aligned(size_t size); + + /** + * @brief Allocate memory that has a starting address that is cache aligned with tag. + * @param[in] size Size of memory to allocate + * @param[in] tag Pool tag to use. + * @returns Pointer to memory block allocated, or null on failure. + */ + __drv_allocatesMem(Mem) _Must_inspect_result_ + _Ret_writes_maybenull_(size) void* usersim_allocate_cache_aligned_with_tag(size_t size, uint32_t tag); + + /** + * @brief Free memory that has a starting address that is cache aligned. + * @param[in] memory Allocation to be freed. + */ + void + usersim_free_cache_aligned(_Frees_ptr_opt_ void* memory); + #if defined(__cplusplus) } @@ -198,80 +269,6 @@ USERSIM_API _Ret_maybenull_ void* ExAllocatePoolWithTagCPP( _In_ __drv_strictTypeMatch(__drv_typeExpr) POOL_TYPE pool_type, SIZE_T number_of_bytes, ULONG tag); -/** - * @brief Allocate memory. - * @param[in] pool_type Pool type to use. - * @param[in] size Size of memory to allocate. - * @param[in] tag Pool tag to use. - * @param[in] initialize False to return "uninitialized" memory. - * @returns Pointer to memory block allocated, or null on failure. - */ -__drv_allocatesMem(Mem) _Must_inspect_result_ _Ret_writes_maybenull_(size) void* usersim_allocate_with_tag( - _In_ __drv_strictTypeMatch(__drv_typeExpr) POOL_TYPE pool_type, - size_t size, - uint32_t tag, - bool initialize); - - /** - * @brief Allocate memory. - * @param[in] size Size of memory to allocate. - * @returns Pointer to memory block allocated, or null on failure. - */ -__drv_allocatesMem(Mem) _Must_inspect_result_ _Ret_writes_maybenull_(size) void* usersim_allocate(size_t size); - -/** - * @brief Reallocate memory. - * @param[in] memory Allocation to be reallocated. - * @param[in] old_size Old size of memory to reallocate. - * @param[in] new_size New size of memory to reallocate. - * @returns Pointer to memory block allocated, or null on failure. - */ -__drv_allocatesMem(Mem) _Must_inspect_result_ _Ret_writes_maybenull_(new_size) void* usersim_reallocate( - _In_ _Post_invalid_ void* memory, size_t old_size, size_t new_size); - -/** - * @brief Reallocate memory with tag. - * @param[in] memory Allocation to be reallocated. - * @param[in] old_size Old size of memory to reallocate. - * @param[in] new_size New size of memory to reallocate. - * @param[in] tag Pool tag to use. - * @returns Pointer to memory block allocated, or null on failure. - */ -__drv_allocatesMem(Mem) _Must_inspect_result_ _Ret_writes_maybenull_(new_size) void* usersim_reallocate_with_tag( - _In_ _Post_invalid_ void* memory, size_t old_size, size_t new_size, uint32_t tag); - -/** - * @brief Free memory. - * @param[in] memory Allocation to be freed. - */ -void -usersim_free(_Frees_ptr_opt_ void* memory); - -/** - * @brief Allocate memory that has a starting address that is cache aligned. - * @param[in] size Size of memory to allocate - * @returns Pointer to memory block allocated, or null on failure. - */ -USERSIM_API -__drv_allocatesMem(Mem) _Must_inspect_result_ - _Ret_writes_maybenull_(size) void* usersim_allocate_cache_aligned(size_t size); - -/** - * @brief Allocate memory that has a starting address that is cache aligned with tag. - * @param[in] size Size of memory to allocate - * @param[in] tag Pool tag to use. - * @returns Pointer to memory block allocated, or null on failure. - */ -__drv_allocatesMem(Mem) _Must_inspect_result_ - _Ret_writes_maybenull_(size) void* usersim_allocate_cache_aligned_with_tag(size_t size, uint32_t tag); - -/** - * @brief Free memory that has a starting address that is cache aligned. - * @param[in] memory Allocation to be freed. - */ -void -usersim_free_cache_aligned(_Frees_ptr_opt_ void* memory); - USERSIM_API _Ret_maybenull_ void* ExAllocatePoolUninitializedCPP(_In_ POOL_TYPE pool_type, _In_ size_t number_of_bytes, _In_ unsigned long tag); diff --git a/inc/usersim/fwp_test.h b/inc/usersim/fwp_test.h index 02b87f6..b062609 100644 --- a/inc/usersim/fwp_test.h +++ b/inc/usersim/fwp_test.h @@ -7,6 +7,11 @@ #include #include +#if defined(__cplusplus) +extern "C" +{ +#endif + typedef struct _fwp_classify_parameters { ADDRESS_FAMILY family; @@ -51,4 +56,8 @@ usersim_fwp_sock_ops_v6(_In_ fwp_classify_parameters_t* parameters); USERSIM_API void usersim_fwp_set_sublayer_guids( - _In_ const GUID& default_sublayer, _In_ const GUID& connect_v4_sublayer, _In_ const GUID& connect_v6_sublayer); \ No newline at end of file + _In_ const GUID& default_sublayer, _In_ const GUID& connect_v4_sublayer, _In_ const GUID& connect_v6_sublayer); + +#if defined(__cplusplus) +} +#endif diff --git a/inc/usersim/wdf.h b/inc/usersim/wdf.h index 1335744..e53624f 100644 --- a/inc/usersim/wdf.h +++ b/inc/usersim/wdf.h @@ -76,9 +76,32 @@ extern "C" void WDF_DRIVER_CONFIG_INIT(_Out_ PWDF_DRIVER_CONFIG config, _In_opt_ PFN_WDF_DRIVER_DEVICE_ADD evt_driver_device_add); + typedef NTSTATUS(WdfDriverCreate_t)( + _In_ WDF_DRIVER_GLOBALS* driver_globals, + _In_ PDRIVER_OBJECT driver_object, + _In_ PCUNICODE_STRING registry_path, + _In_opt_ PWDF_OBJECT_ATTRIBUTES driver_attributes, + _In_ PWDF_DRIVER_CONFIG driver_config, + _Out_opt_ WDFDRIVER* driver); + + typedef NTSTATUS(WdfDeviceCreate_t)( + _In_ WDF_DRIVER_GLOBALS* driver_globals, + _Inout_ PWDFDEVICE_INIT* device_init, + _In_opt_ PWDF_OBJECT_ATTRIBUTES device_attributes, + _Out_ WDFDEVICE* device); + + typedef _IRQL_requires_max_(DISPATCH_LEVEL) + VOID(WdfDeviceInitFree_t)(_In_ PWDF_DRIVER_GLOBALS driver_globals, _In_ PWDFDEVICE_INIT device_init); + + typedef _Must_inspect_result_ _IRQL_requires_max_(PASSIVE_LEVEL) PWDFDEVICE_INIT(WdfControlDeviceInitAllocate_t)( + _In_ PWDF_DRIVER_GLOBALS driver_globals, + _In_ WDFDRIVER driver, + _In_ CONST UNICODE_STRING* sddl_string); + typedef enum _WDFFUNCENUM { WdfControlDeviceInitAllocateTableIndex = 25, + WdfDeviceInitFreeTableIndex = 54, WdfDeviceInitSetDeviceTypeTableIndex = 66, WdfDeviceInitAssignNameTableIndex = 67, WdfDeviceInitSetCharacteristicsTableIndex = 70, diff --git a/src/ex.cpp b/src/ex.cpp index 797f07b..3c95244 100644 --- a/src/ex.cpp +++ b/src/ex.cpp @@ -354,6 +354,7 @@ usersim_allocate_with_tag( // The pointer we return has to be cache aligned so we allocate // enough extra space to fill a cache line, and put the // usersim_allocation_header_t at the end of that space. + // TODO: move logic into usersim_allocate_cache_aligned_with_tag(). size_t full_size = USERSIM_CACHE_LINE_SIZE + size; uint8_t* pointer = (uint8_t*)_aligned_malloc(full_size, USERSIM_CACHE_LINE_SIZE); if (pointer == nullptr) { diff --git a/src/ndis.h b/src/ndis.h index c27e1aa..2d4f4fd 100644 --- a/src/ndis.h +++ b/src/ndis.h @@ -9,6 +9,11 @@ #include #include +#if defined(__cplusplus) +extern "C" +{ +#endif + #define NET_BUFFER_FIRST_MDL(_NB) ((_NB)->MdlChain) #define NDIS_STATUS_SUCCESS ((NDIS_STATUS)STATUS_SUCCESS) #define NET_BUFFER_LIST_FIRST_NB(_NBL) ((_NBL)->FirstNetBuffer) @@ -43,41 +48,41 @@ typedef struct _NET_BUFFER_LIST_CONTEXT NET_BUFFER_LIST_CONTEXT, *PNET_BUFFER_LI typedef struct _NDIS_GENERIC_OBJECT NDIS_GENERIC_OBJECT, *PNDIS_GENERIC_OBJECT; -__declspec(dllexport) PNDIS_GENERIC_OBJECT +USERSIM_API PNDIS_GENERIC_OBJECT NdisAllocateGenericObject(_In_opt_ DRIVER_OBJECT* driver_object, _In_ unsigned long tag, _In_ uint16_t size); -__declspec(dllexport) NDIS_HANDLE +USERSIM_API NDIS_HANDLE NdisAllocateNetBufferListPool(_In_opt_ NDIS_HANDLE ndis_handle, _In_ NET_BUFFER_LIST_POOL_PARAMETERS const* parameters); -__declspec(dllexport) NET_BUFFER_LIST* +USERSIM_API NET_BUFFER_LIST* NdisAllocateCloneNetBufferList( _In_ NET_BUFFER_LIST* original_net_buffer_list, _In_ NDIS_HANDLE net_buffer_list_pool_handle, _In_ NDIS_HANDLE net_buffer_pool_handle, ULONG allocate_clone_flags); -__declspec(dllexport) void +USERSIM_API void NdisFreeCloneNetBufferList(_In_ NET_BUFFER_LIST* clone_net_buffer_list, ULONG free_clone_flags); -__declspec(dllexport) PNET_BUFFER_LIST +USERSIM_API PNET_BUFFER_LIST NdisAllocateNetBufferList(_In_ NDIS_HANDLE nbl_pool_handle, _In_ USHORT context_size, _In_ USHORT context_backfill); -__declspec(dllexport) _Must_inspect_result_ __drv_allocatesMem(mem) NET_BUFFER* NdisAllocateNetBuffer( +USERSIM_API _Must_inspect_result_ __drv_allocatesMem(mem) NET_BUFFER* NdisAllocateNetBuffer( _In_ NDIS_HANDLE pool_handle, _In_opt_ MDL* mdl_chain, _In_ unsigned long data_offset, _In_ SIZE_T data_length); -__declspec(dllexport) VOID +USERSIM_API VOID NdisFreeNetBuffer(_In_ __drv_freesMem(mem) NET_BUFFER* net_buffer); -__declspec(dllexport) VOID +USERSIM_API VOID NdisFreeNetBufferList(_In_ __drv_freesMem(mem) NET_BUFFER_LIST* net_buffer_list); -__declspec(dllexport) void +USERSIM_API void NdisFreeNetBufferListPool(_In_ __drv_freesMem(mem) NDIS_HANDLE pool_handle); -__declspec(dllexport) void +USERSIM_API void NdisFreeGenericObject(_In_ PNDIS_GENERIC_OBJECT ndis_object); -__declspec(dllexport) void* +USERSIM_API void* NdisGetDataBuffer( _In_ NET_BUFFER* net_buffer, _In_ unsigned long bytes_needed, @@ -85,16 +90,20 @@ NdisGetDataBuffer( _In_ unsigned long align_multiple, _In_ unsigned long align_offset); -__declspec(dllexport) NDIS_STATUS +USERSIM_API NDIS_STATUS NdisRetreatNetBufferDataStart( _In_ NET_BUFFER* net_buffer, _In_ unsigned long data_offset_delta, _In_ unsigned long data_back_fill, _In_opt_ void* allocate_mdl_handler); -__declspec(dllexport) void +USERSIM_API void NdisAdvanceNetBufferDataStart( _In_ NET_BUFFER* net_buffer, _In_ unsigned long data_offset_delta, _In_ BOOLEAN free_mdl, _In_opt_ void* free_mdl_handler); + +#if defined(__cplusplus) +} +#endif \ No newline at end of file diff --git a/src/wdf.cpp b/src/wdf.cpp index 1eb609b..d1ecf32 100644 --- a/src/wdf.cpp +++ b/src/wdf.cpp @@ -9,6 +9,11 @@ WDF_DRIVER_GLOBALS g_UsersimWdfDriverGlobals = {0}; +static WdfDriverCreate_t _WdfDriverCreate; +static WdfDeviceCreate_t _WdfDeviceCreate; +static WdfControlDeviceInitAllocate_t _WdfControlDeviceInitAllocate; +static WdfDeviceInitFree_t _WdfDeviceInitFree; + static NTSTATUS _WdfDriverCreate( _In_ WDF_DRIVER_GLOBALS* driver_globals, @@ -42,7 +47,6 @@ _WdfDeviceCreate( _In_opt_ PWDF_OBJECT_ATTRIBUTES device_attributes, _Out_ WDFDEVICE* device) { - UNREFERENCED_PARAMETER(device_init); UNREFERENCED_PARAMETER(device_attributes); if (driver_globals != &g_UsersimWdfDriverGlobals) { @@ -52,7 +56,7 @@ _WdfDeviceCreate( return STATUS_INSUFFICIENT_RESOURCES; } - *device = nullptr; + *device = *device_init; return STATUS_SUCCESS; } @@ -89,6 +93,13 @@ _WdfControlDeviceInitAllocate( return device_init; } +static _IRQL_requires_max_(DISPATCH_LEVEL) VOID +_WdfDeviceInitFree(_In_ PWDF_DRIVER_GLOBALS driver_globals, _In_ PWDFDEVICE_INIT device_init) +{ + UNREFERENCED_PARAMETER(driver_globals); + usersim_free(device_init); +} + static _IRQL_requires_max_(DISPATCH_LEVEL) VOID _WdfDeviceInitSetDeviceType( @@ -247,6 +258,7 @@ void usersim_initialize_wdf() { g_UsersimWdfFunctions[WdfControlDeviceInitAllocateTableIndex] = (WDFFUNC)_WdfControlDeviceInitAllocate; + g_UsersimWdfFunctions[WdfDeviceInitFreeTableIndex] = (WDFFUNC)_WdfDeviceInitFree; g_UsersimWdfFunctions[WdfDeviceInitSetDeviceTypeTableIndex] = (WDFFUNC)_WdfDeviceInitSetDeviceType; g_UsersimWdfFunctions[WdfDeviceInitAssignNameTableIndex] = (WDFFUNC)_WdfDeviceInitAssignName; g_UsersimWdfFunctions[WdfDeviceInitSetCharacteristicsTableIndex] = (WDFFUNC)_WdfDeviceInitSetCharacteristics; diff --git a/tests/wdf_test.cpp b/tests/wdf_test.cpp index 093a62c..8ec0731 100644 --- a/tests/wdf_test.cpp +++ b/tests/wdf_test.cpp @@ -29,17 +29,43 @@ TEST_CASE("WdfDriverCreate", "[wdf]") WDF_DRIVER_CONFIG config; WDF_DRIVER_CONFIG_INIT(&config, nullptr); - typedef NTSTATUS(FN_WdfDriverCreate)( - _In_ WDF_DRIVER_GLOBALS* driver_globals, - _In_ PDRIVER_OBJECT driver_object, - _In_ PCUNICODE_STRING registry_path, - _In_opt_ PWDF_OBJECT_ATTRIBUTES driver_attributes, - _In_ PWDF_DRIVER_CONFIG driver_config, - _Out_opt_ WDFDRIVER* driver); - FN_WdfDriverCreate* WdfDriverCreate = (FN_WdfDriverCreate*)UsersimWdfFunctions[WdfDriverCreateTableIndex]; - WDFDRIVER driver; + WdfDriverCreate_t* WdfDriverCreate = (WdfDriverCreate_t*)UsersimWdfFunctions[WdfDriverCreateTableIndex]; + WDFDRIVER driver = nullptr; NTSTATUS status = WdfDriverCreate(UsersimWdfDriverGlobals, &driver_object, nullptr, nullptr, &config, &driver); REQUIRE(status == STATUS_SUCCESS); - REQUIRE(driver == &driver_object); + REQUIRE(driver != nullptr); + + // Verify that a duplicate create fails. + status = WdfDriverCreate(UsersimWdfDriverGlobals, &driver_object, nullptr, nullptr, &config, &driver); + REQUIRE(status == STATUS_DRIVER_INTERNAL_ERROR); + + UsersimWdfDriverGlobals->Driver = nullptr; +} + +TEST_CASE("WdfDeviceCreate", "[wdf]") +{ + DRIVER_OBJECT driver_object = {0}; + WDF_DRIVER_CONFIG config; + WDF_DRIVER_CONFIG_INIT(&config, nullptr); + WdfDriverCreate_t* WdfDriverCreate = (WdfDriverCreate_t*)UsersimWdfFunctions[WdfDriverCreateTableIndex]; + WDFDRIVER driver = nullptr; + NTSTATUS status = WdfDriverCreate(UsersimWdfDriverGlobals, &driver_object, nullptr, nullptr, &config, &driver); + REQUIRE(status == STATUS_SUCCESS); + + WdfDeviceCreate_t* WdfDeviceCreate = (WdfDeviceCreate_t*)UsersimWdfFunctions[WdfDeviceCreateTableIndex]; + WdfControlDeviceInitAllocate_t* WdfControlDeviceInitAllocate = + (WdfControlDeviceInitAllocate_t*)UsersimWdfFunctions[WdfControlDeviceInitAllocateTableIndex]; + DECLARE_CONST_UNICODE_STRING(security_descriptor, L""); + PWDFDEVICE_INIT init = WdfControlDeviceInitAllocate(UsersimWdfDriverGlobals, driver, &security_descriptor); + REQUIRE(init != nullptr); + + WDFDEVICE device = nullptr; + status = WdfDeviceCreate(UsersimWdfDriverGlobals, &init, nullptr, &device); + REQUIRE(status == STATUS_SUCCESS); + REQUIRE(device != nullptr); + + WdfDeviceInitFree_t* WdfDeviceInitFree = (WdfDeviceInitFree_t*)UsersimWdfFunctions[WdfDeviceInitFreeTableIndex]; + WdfDeviceInitFree(UsersimWdfDriverGlobals, init); + UsersimWdfDriverGlobals->Driver = nullptr; }