diff --git a/UsbDk/Registry.cpp b/UsbDk/Registry.cpp new file mode 100644 index 0000000..0e344c3 --- /dev/null +++ b/UsbDk/Registry.cpp @@ -0,0 +1,107 @@ +/********************************************************************** +* Copyright (c) 2013-2014 Red Hat, Inc. +* +* Developed by Daynix Computing LTD. +* +* Authors: +* Dmitry Fleytman +* Kirill Moizik +* +* Licensed under the Apache License, Version 2.0 (the "License"); +* you may not use this file except in compliance with the License. +* You may obtain a copy of the License at +* +* http://www.apache.org/licenses/LICENSE-2.0 +* +* Unless required by applicable law or agreed to in writing, software +* distributed under the License is distributed on an "AS IS" BASIS, +* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +* See the License for the specific language governing permissions and +* limitations under the License. +* +**********************************************************************/ + +#include "stdafx.h" +#include "Alloc.h" +#include "MemoryBuffer.h" +#include "Trace.h" +#include "Registry.h" +#include "Registry.tmh" + +NTSTATUS CRegKey::Open(const CRegKey &ParentKey, const UNICODE_STRING &RegPath) +{ + OBJECT_ATTRIBUTES KeyAttributes; + + InitializeObjectAttributes(&KeyAttributes, + const_cast(&RegPath), + OBJ_KERNEL_HANDLE, ParentKey.m_Key, nullptr); + + auto status = ZwOpenKey(&m_Key, KEY_READ, &KeyAttributes); + if (!NT_SUCCESS(status)) + { + TraceEvents(TRACE_LEVEL_ERROR, TRACE_REGISTRY, + "%!FUNC! Failed to open registry key %wZ, parent 0x%p (status: %!STATUS!)", &RegPath, ParentKey.m_Key, status); + } + + return status; +} + +NTSTATUS CRegKey::QuerySubkeyInfo(ULONG Index, + KEY_INFORMATION_CLASS InfoClass, + CWdmMemoryBuffer &Buffer) +{ + ULONG BytesNeeded = 0; + NTSTATUS status = STATUS_BUFFER_TOO_SMALL; + + while ((status == STATUS_BUFFER_OVERFLOW) || (status == STATUS_BUFFER_TOO_SMALL)) + { + status = Buffer.Recreate(BytesNeeded, PagedPool); + if (NT_SUCCESS(status)) + { + status = ZwEnumerateKey(m_Key, + Index, + InfoClass, + Buffer.Ptr(), + static_cast(Buffer.Size()), + &BytesNeeded); + } + } + + if (!NT_SUCCESS(status) && (status != STATUS_NO_MORE_ENTRIES)) + { + TraceEvents(TRACE_LEVEL_ERROR, TRACE_REGISTRY, + "%!FUNC! Failed to query subkey #%d info, info class %d (status: %!STATUS!)", Index, InfoClass, status); + } + + return status; +} + +NTSTATUS CRegKey::QueryValueInfo(const UNICODE_STRING &ValueName, + KEY_VALUE_INFORMATION_CLASS InfoClass, + CWdmMemoryBuffer &Buffer) +{ + ULONG BytesNeeded = 0; + NTSTATUS status = STATUS_BUFFER_TOO_SMALL; + + while ((status == STATUS_BUFFER_OVERFLOW) || (status == STATUS_BUFFER_TOO_SMALL)) + { + status = Buffer.Recreate(BytesNeeded, PagedPool); + if (NT_SUCCESS(status)) + { + status = ZwQueryValueKey(m_Key, + const_cast(&ValueName), + InfoClass, + Buffer.Ptr(), + static_cast(Buffer.Size()), + &BytesNeeded); + } + } + + if (status != STATUS_SUCCESS) + { + TraceEvents(TRACE_LEVEL_ERROR, TRACE_REGISTRY, + "%!FUNC! Failed to query value %wZ, info class %d (status: %!STATUS!)", &ValueName, InfoClass, status); + } + + return status; +} diff --git a/UsbDk/Registry.h b/UsbDk/Registry.h new file mode 100644 index 0000000..cd57804 --- /dev/null +++ b/UsbDk/Registry.h @@ -0,0 +1,77 @@ +/********************************************************************** +* Copyright (c) 2013-2014 Red Hat, Inc. +* +* Developed by Daynix Computing LTD. +* +* Authors: +* Dmitry Fleytman +* Kirill Moizik +* +* Licensed under the Apache License, Version 2.0 (the "License"); +* you may not use this file except in compliance with the License. +* You may obtain a copy of the License at +* +* http://www.apache.org/licenses/LICENSE-2.0 +* +* Unless required by applicable law or agreed to in writing, software +* distributed under the License is distributed on an "AS IS" BASIS, +* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +* See the License for the specific language governing permissions and +* limitations under the License. +* +**********************************************************************/ + +#pragma once + +class CRegKey +{ +public: + NTSTATUS Open(const CRegKey &ParentKey, const UNICODE_STRING &RegPath); + NTSTATUS Open(const UNICODE_STRING &RegPath) + { return Open(CRegKey(), RegPath); } + + template + NTSTATUS ForEachSubKey(TFunctor Functor) + { + auto status = STATUS_SUCCESS; + for (ULONG Index = 0; status == STATUS_SUCCESS; Index++) + { + CWdmMemoryBuffer InfoBuffer; + status = QuerySubkeyInfo(Index, KeyBasicInformation, InfoBuffer); + if (NT_SUCCESS(status)) + { + auto Info = reinterpret_cast(InfoBuffer.Ptr()); + CStringHolder Name; + Name.Attach(Info->Name, static_cast(Info->NameLength)); + + Functor(Name); + } + } + + if (status == STATUS_NO_MORE_ENTRIES) + { + return STATUS_SUCCESS; + } + + return status; + } + + ~CRegKey() + { + if (m_Key != nullptr) + { + ZwClose(m_Key); + } + } + +protected: + NTSTATUS QuerySubkeyInfo(ULONG Index, + KEY_INFORMATION_CLASS InfoClass, + CWdmMemoryBuffer &Buffer); + + NTSTATUS QueryValueInfo(const UNICODE_STRING &ValueName, + KEY_VALUE_INFORMATION_CLASS InfoClass, + CWdmMemoryBuffer &Buffer); + + HANDLE m_Key = nullptr; +}; diff --git a/UsbDk/Trace.h b/UsbDk/Trace.h index 7270ec4..d25d517 100644 --- a/UsbDk/Trace.h +++ b/UsbDk/Trace.h @@ -46,6 +46,7 @@ WPP_DEFINE_BIT(TRACE_USBTARGET) \ WPP_DEFINE_BIT(TRACE_FILTERSTRATEGY) \ WPP_DEFINE_BIT(TRACE_URB) \ + WPP_DEFINE_BIT(TRACE_REGISTRY) \ ) #define WPP_FLAG_LEVEL_LOGGER(flag, level) \ diff --git a/UsbDk/UsbDk.vcxproj b/UsbDk/UsbDk.vcxproj index 05b2cee..ddb2aba 100644 --- a/UsbDk/UsbDk.vcxproj +++ b/UsbDk/UsbDk.vcxproj @@ -76,6 +76,7 @@ + @@ -97,6 +98,7 @@ + diff --git a/UsbDk/UsbDk.vcxproj.filters b/UsbDk/UsbDk.vcxproj.filters index 3de99ea..b5395f5 100644 --- a/UsbDk/UsbDk.vcxproj.filters +++ b/UsbDk/UsbDk.vcxproj.filters @@ -94,6 +94,9 @@ Header Files + + Header Files + @@ -147,6 +150,9 @@ Source Files + + Source Files +