diff --git a/UsbDk/WdfRequest.h b/UsbDk/WdfRequest.h index 06e1261..4083d3c 100644 --- a/UsbDk/WdfRequest.h +++ b/UsbDk/WdfRequest.h @@ -29,6 +29,18 @@ public: return WdfRequestRetrieveOutputBuffer(m_Request, sizeof(*objectPtr), (PVOID *)&objectPtr, Length); } + template + NTSTATUS FetchInputMemory(WDFMEMORY *Memory) + { + return WdfRequestRetrieveInputMemory(m_Request, Memory); + } + + template + NTSTATUS FetchOutputMemory(WDFMEMORY *Memory) + { + return WdfRequestRetrieveOutputMemory(m_Request, Memory); + } + template NTSTATUS FetchInputArray(TObject &arrayPtr, size_t &numElements) { @@ -155,3 +167,34 @@ static void UsbDkHandleRequestWithOutput(CWdfRequest &Request, Request.SetOutputDataLen(outputLength); Request.SetStatus(status); } + +template +static void UsbDkHandleRequestWithIOMemory(CWdfRequest &Request, + THandler Handler) +{ + WDFMEMORY output; + + auto status = Request.FetchOutputMemory(&output); + if (!NT_SUCCESS(status)) + { + Request.SetOutputDataLen(0); + Request.SetStatus(status); + return; + } + + WDFMEMORY input; + + status = Request.FetchInputMemory(&input); + if (!NT_SUCCESS(status)) + { + Request.SetOutputDataLen(0); + Request.SetStatus(status); + return; + } + + size_t outputLength; + status = Handler(input, output, outputLength); + + Request.SetOutputDataLen(outputLength); + Request.SetStatus(status); +}