Files
Katayama Hirofumi MZ 29ca914bd4 [SHLWAPI] Use <new> std::nothrow (#9150)
Use preferred memory handling on memory
shortage.
JIRA issue: N/A
- Link cppstl instead of cpprt.
- Use <new> new(std::throw) for memory
  allocation.
- Return E_OUTOFMEMORY on allocation
  failure.
- Don't use <stdlib.h> min macro but use
  <cstdlib> __min.
2026-06-15 08:50:15 +09:00

519 lines
13 KiB
C++

/*
* PROJECT: ReactOS Shell
* LICENSE: LGPL-2.0-or-later (https://spdx.org/licenses/LGPL-2.0-or-later)
* PURPOSE: Implement QuerySourceCreateFromKey
* COPYRIGHT: Copyright 2026 Katayama Hirofumi MZ <[email protected]>
*/
#include "precomp.h"
#include <shlwapi.h>
#include <shlwapi_undoc.h>
#include <shlobj_undoc.h>
#include <shlguid_undoc.h>
#include <new>
WINE_DEFAULT_DEBUG_CHANNEL(shell);
static HRESULT SHAllocBlob(ULONG cbData, const BYTE *pbData, FLAGGED_BYTE_BLOB** ppBlob)
{
FLAGGED_BYTE_BLOB* pBlob =
(FLAGGED_BYTE_BLOB*)CoTaskMemAlloc(sizeof(FLAGGED_BYTE_BLOB) + cbData);
if (!pBlob)
return E_OUTOFMEMORY;
pBlob->clSize = cbData;
if (pbData)
CopyMemory(pBlob->abData, pbData, cbData);
*ppBlob = pBlob;
return S_OK;
}
class CRegistrySource;
/******************************************************************************
* CRegistryEnumBase
*/
class CRegistryEnumBase : public IEnumString
{
protected:
LONG m_cRefs = 1;
DWORD m_dwIndex = 0;
CRegistrySource *m_pSource = NULL;
HKEY m_hKey = NULL;
PWSTR m_pszName = NULL;
WCHAR m_szBuf[64] = {};
DWORD m_cchNameMax = 0;
BOOL _Next(PWSTR *ppwsz);
virtual BOOL _RegNext(DWORD dwIndex) = 0;
virtual DWORD _MaxLen() = 0;
public:
virtual ~CRegistryEnumBase();
HRESULT Init(HKEY hKey, CRegistrySource *pSource);
// IUnknown methods
STDMETHODIMP QueryInterface(REFIID riid, void **ppvObject) override;
STDMETHODIMP_(ULONG) AddRef() override;
STDMETHODIMP_(ULONG) Release() override;
// IEnumString methods
STDMETHODIMP Next(ULONG celt, LPWSTR* rgelt, ULONG* pceltFetched) override;
STDMETHODIMP Skip(ULONG celt) override;
STDMETHODIMP Reset() override;
STDMETHODIMP Clone(IEnumString ** ppenum) override;
};
/******************************************************************************
* CRegistryEnumValues
*/
class CRegistryEnumValues : public CRegistryEnumBase
{
public:
BOOL _RegNext(DWORD dwIndex) override;
DWORD _MaxLen() override;
};
/******************************************************************************
* CRegistryEnumKeys
*/
class CRegistryEnumKeys : public CRegistryEnumBase
{
public:
BOOL _RegNext(DWORD dwIndex) override;
DWORD _MaxLen() override;
};
/******************************************************************************
* CRegistrySource
*/
class CRegistrySource
: public IQuerySourceOld
, public IObjectWithRegistryKeyOld
{
LONG m_cRefs = 1;
HKEY m_hKey = NULL;
public:
virtual ~CRegistrySource();
HRESULT Init(HKEY hKey, PCWSTR pszSubKey, BOOL bCreate);
// IUnknown methods
STDMETHODIMP QueryInterface(REFIID riid, void **ppvObject) override;
STDMETHODIMP_(ULONG) AddRef() override;
STDMETHODIMP_(ULONG) Release() override;
// IQuerySourceOld methods
STDMETHODIMP EnumValues(IEnumString **ppEnum) override;
STDMETHODIMP EnumSources(IEnumString **ppEnum) override;
STDMETHODIMP QueryValueString(PCWSTR keyName, PCWSTR valueName, PWSTR *ppszValue) override;
STDMETHODIMP QueryValueDword(PCWSTR keyName, PCWSTR valueName, DWORD *pdwValue) override;
STDMETHODIMP QueryValueExists(PCWSTR keyName, PCWSTR valueName) override;
STDMETHODIMP QueryValueDirect(PCWSTR keyName, PCWSTR valueName,
FLAGGED_BYTE_BLOB **ppBlob) override;
STDMETHODIMP OpenSource(PCWSTR keyName, BOOL bCreate, IQuerySourceOld **ppSource) override;
STDMETHODIMP SetValueDirect(PCWSTR keyName, PCWSTR valueName, DWORD dwType, DWORD cbData,
LPCVOID pvData) override;
// IObjectWithRegistryKeyOld methods
STDMETHODIMP SetKey(HKEY hKey) override;
STDMETHODIMP GetKey(HKEY *phKey) override;
};
/******************************************************************************/
CRegistryEnumBase::~CRegistryEnumBase()
{
if (m_pszName && m_pszName != m_szBuf)
LocalFree(m_pszName);
if (m_pSource)
m_pSource->Release();
}
HRESULT CRegistryEnumBase::Init(HKEY hKey, CRegistrySource *pSource)
{
m_hKey = hKey;
m_pSource = pSource;
m_pSource->AddRef();
m_cchNameMax = _MaxLen();
if (m_cchNameMax > _countof(m_szBuf))
{
m_pszName = (PWSTR)LocalAlloc(LPTR, m_cchNameMax * sizeof(WCHAR));
}
else
{
m_cchNameMax = _countof(m_szBuf);
m_pszName = m_szBuf;
}
return m_pszName ? S_OK : E_OUTOFMEMORY;
}
BOOL CRegistryEnumBase::_Next(PWSTR *ppwsz)
{
return _RegNext(m_dwIndex) && SUCCEEDED(SHStrDupW(m_pszName, ppwsz));
}
HRESULT CRegistryEnumBase::QueryInterface(REFIID riid, PVOID* ppv)
{
if (!ppv)
return E_POINTER;
if (riid == IID_IEnumString)
{
*ppv = static_cast<IEnumString*>(this);
AddRef();
return S_OK;
}
return E_NOINTERFACE;
}
STDMETHODIMP_(ULONG) CRegistryEnumBase::AddRef()
{
return ::InterlockedIncrement(&m_cRefs);
}
STDMETHODIMP_(ULONG) CRegistryEnumBase::Release()
{
LONG refs = ::InterlockedDecrement(&m_cRefs);
if (!refs)
delete this;
return refs;
}
STDMETHODIMP CRegistryEnumBase::Next(ULONG celt, LPWSTR* rgelt, ULONG* pceltFetched)
{
if (!rgelt || (celt > 1 && !pceltFetched))
return E_INVALIDARG;
ULONG cFetched = 0;
while (cFetched < celt)
{
if (!_Next(&rgelt[cFetched]))
break;
++m_dwIndex;
++cFetched;
}
if (pceltFetched)
*pceltFetched = cFetched;
return (cFetched == celt) ? S_OK : S_FALSE;
}
STDMETHODIMP CRegistryEnumBase::Skip(ULONG celt)
{
return E_NOTIMPL;
}
STDMETHODIMP CRegistryEnumBase::Reset()
{
m_dwIndex = 0;
return S_OK;
}
STDMETHODIMP CRegistryEnumBase::Clone(IEnumString ** ppenum)
{
return E_NOTIMPL;
}
/******************************************************************************/
BOOL CRegistryEnumKeys::_RegNext(DWORD dwIndex)
{
return RegEnumKeyW(m_hKey, dwIndex, m_pszName, m_cchNameMax) == ERROR_SUCCESS;
}
DWORD CRegistryEnumKeys::_MaxLen()
{
DWORD cchKeyNameMax = 0;
RegQueryInfoKeyW(m_hKey, NULL, NULL, NULL, NULL, &cchKeyNameMax, NULL, NULL, NULL, NULL,
NULL, NULL);
return cchKeyNameMax + 1; // Including NUL
}
/******************************************************************************/
BOOL CRegistryEnumValues::_RegNext(DWORD dwIndex)
{
DWORD cchValueNameMax = m_cchNameMax;
return RegEnumValueW(m_hKey, dwIndex, m_pszName, &cchValueNameMax, NULL,
NULL, NULL, NULL) == ERROR_SUCCESS;
}
DWORD CRegistryEnumValues::_MaxLen()
{
DWORD cchValueNameMax = 0;
RegQueryInfoKeyW(m_hKey, NULL, NULL, NULL, NULL, NULL, NULL, NULL, &cchValueNameMax,
NULL, NULL, NULL);
return cchValueNameMax + 1; // Including NUL
}
/******************************************************************************/
CRegistrySource::~CRegistrySource()
{
if (m_hKey)
RegCloseKey(m_hKey);
}
HRESULT CRegistrySource::Init(HKEY hKey, PCWSTR pszSubKey, BOOL bCreate)
{
LSTATUS error;
if (bCreate)
error = RegCreateKeyExW(hKey, pszSubKey, 0, NULL, 0, MAXIMUM_ALLOWED, NULL, &m_hKey, NULL);
else
error = RegOpenKeyExW(hKey, pszSubKey, 0, MAXIMUM_ALLOWED, &m_hKey);
return error ? HRESULT_FROM_WIN32(error) : S_OK;
}
STDMETHODIMP CRegistrySource::QueryInterface(REFIID riid, void **ppvObject)
{
if (!ppvObject)
return E_POINTER;
if (riid == IID_IQuerySourceOld)
{
*ppvObject = static_cast<IQuerySourceOld*>(this);
AddRef();
return S_OK;
}
if (riid == IID_IObjectWithRegistryKeyOld)
{
*ppvObject = static_cast<IObjectWithRegistryKeyOld*>(this);
AddRef();
return S_OK;
}
*ppvObject = NULL;
return E_NOINTERFACE;
}
STDMETHODIMP_(ULONG) CRegistrySource::AddRef()
{
return ::InterlockedIncrement(&m_cRefs);
}
STDMETHODIMP_(ULONG) CRegistrySource::Release()
{
LONG refs = ::InterlockedDecrement(&m_cRefs);
if (!refs)
delete this;
return refs;
}
STDMETHODIMP CRegistrySource::EnumValues(IEnumString **ppEnum)
{
CRegistryEnumValues* pEnum = new(std::nothrow) CRegistryEnumValues();
if (!pEnum)
return E_OUTOFMEMORY;
HRESULT hr = pEnum->Init(m_hKey, this);
if (FAILED(hr))
{
pEnum->Release();
pEnum = NULL;
}
*ppEnum = pEnum;
return pEnum ? S_OK : E_OUTOFMEMORY;
}
STDMETHODIMP CRegistrySource::EnumSources(IEnumString **ppEnum)
{
CRegistryEnumKeys* pEnum = new(std::nothrow) CRegistryEnumKeys();
if (!pEnum)
return E_OUTOFMEMORY;
HRESULT hr = pEnum->Init(m_hKey, this);
if (FAILED(hr))
{
pEnum->Release();
pEnum = NULL;
}
*ppEnum = pEnum;
return pEnum ? S_OK : E_OUTOFMEMORY;
}
STDMETHODIMP CRegistrySource::QueryValueString(
PCWSTR keyName,
PCWSTR valueName,
PWSTR *ppszValue)
{
*ppszValue = NULL;
WCHAR szData[128];
DWORD dwType = REG_NONE, cbData = sizeof(szData);
LSTATUS error = SHGetValueW(m_hKey, keyName, valueName, &dwType, szData, &cbData);
if (error == ERROR_SUCCESS)
{
if (dwType != REG_SZ)
return E_DATATYPE_MISMATCH;
if (!valueName && !szData[0])
return HRESULT_FROM_WIN32(ERROR_FILE_NOT_FOUND);
return SHStrDupW(szData, ppszValue);
}
if (error != ERROR_MORE_DATA)
return error ? HRESULT_FROM_WIN32(error) : S_OK;
*ppszValue = (PWSTR)CoTaskMemAlloc(cbData);
if (!*ppszValue)
return E_OUTOFMEMORY;
HRESULT hr = S_OK;
error = SHGetValueW(m_hKey, keyName, valueName, &dwType, *ppszValue, &cbData);
if (error)
{
CoTaskMemFree(*ppszValue);
*ppszValue = NULL;
hr = HRESULT_FROM_WIN32(error);
}
if (SUCCEEDED(hr) && dwType != REG_SZ)
{
CoTaskMemFree(*ppszValue);
*ppszValue = NULL;
hr = E_DATATYPE_MISMATCH;
}
return hr;
}
STDMETHODIMP CRegistrySource::QueryValueDword(PCWSTR keyName, PCWSTR valueName, DWORD *pdwValue)
{
DWORD dwType, cbValue = sizeof(*pdwValue);
LSTATUS error = SHGetValueW(m_hKey, keyName, valueName, &dwType, pdwValue, &cbValue);
if (error)
return HRESULT_FROM_WIN32(error);
if (dwType != REG_DWORD)
return E_DATATYPE_MISMATCH;
return S_OK;
}
STDMETHODIMP CRegistrySource::QueryValueExists(PCWSTR keyName, PCWSTR valueName)
{
LSTATUS error = SHGetValueW(m_hKey, keyName, valueName, NULL, NULL, NULL);
if (error)
return HRESULT_FROM_WIN32(error);
return S_OK;
}
STDMETHODIMP CRegistrySource::QueryValueDirect(
PCWSTR keyName,
PCWSTR valueName,
FLAGGED_BYTE_BLOB **ppBlob)
{
HRESULT hr = E_FAIL;
HKEY hKey = m_hKey;
DWORD dwType = REG_NONE;
DWORD cbData = 256;
BYTE abData[256];
*ppBlob = NULL;
if (keyName && *keyName &&
RegOpenKeyExW(m_hKey, keyName, 0, KEY_READ, &hKey) != ERROR_SUCCESS)
{
return E_FAIL;
}
LSTATUS error = RegQueryValueExW(hKey, valueName, NULL, &dwType, abData, &cbData);
if (error == ERROR_SUCCESS)
{
hr = SHAllocBlob(cbData, abData, ppBlob);
}
else if (error == ERROR_MORE_DATA)
{
hr = SHAllocBlob(cbData, NULL, ppBlob);
if (SUCCEEDED(hr))
{
error = RegQueryValueExW(hKey, valueName, NULL, &dwType, (*ppBlob)->abData, &cbData);
if (error != ERROR_SUCCESS)
{
CoTaskMemFree(*ppBlob);
*ppBlob = NULL;
hr = error ? HRESULT_FROM_WIN32(error) : S_OK;
}
}
}
else
{
hr = HRESULT_FROM_WIN32(error);
}
if (hKey != m_hKey)
RegCloseKey(hKey);
if (SUCCEEDED(hr))
(*ppBlob)->fFlags = dwType;
return hr;
}
STDMETHODIMP CRegistrySource::OpenSource(PCWSTR keyName, BOOL bCreate, IQuerySourceOld **ppSource)
{
return QuerySourceCreateFromKey(m_hKey, keyName, bCreate, IID_IQuerySourceOld, (PVOID*)ppSource);
}
STDMETHODIMP CRegistrySource::SetValueDirect(
PCWSTR keyName,
PCWSTR valueName,
DWORD dwType,
DWORD cbData,
LPCVOID pvData)
{
LSTATUS error = SHSetValueW(m_hKey, keyName, valueName, dwType, pvData, cbData);
if (error)
return HRESULT_FROM_WIN32(error);
return S_OK;
}
STDMETHODIMP CRegistrySource::SetKey(HKEY hKey)
{
if (m_hKey)
return E_UNEXPECTED;
m_hKey = SHRegDuplicateHKey(hKey);
return m_hKey ? S_OK : E_UNEXPECTED;
}
STDMETHODIMP CRegistrySource::GetKey(HKEY *phKey)
{
if (!m_hKey)
return E_UNEXPECTED;
*phKey = SHRegDuplicateHKey(m_hKey);
if (!*phKey)
return E_UNEXPECTED;
return S_OK;
}
/**************************************************************************
* QuerySourceCreateFromKey (SHLWAPI.544)
*
* @see https://www.geoffchappell.com/studies/windows/shell/shlwapi/api/regsrc/createfromkey.htm
*/
EXTERN_C
HRESULT WINAPI
QuerySourceCreateFromKey(
_In_ HKEY hKey,
_In_opt_ PCWSTR lpSubKey,
_In_ BOOL bCreate,
_In_ REFIID riid,
_Outptr_ PVOID *ppv)
{
*ppv = NULL;
CRegistrySource* pRS = new(std::nothrow) CRegistrySource();
if (!pRS)
return E_OUTOFMEMORY;
HRESULT hr = pRS->Init(hKey, lpSubKey, bCreate);
if (SUCCEEDED(hr))
hr = pRS->QueryInterface(riid, ppv);
pRS->Release();
return hr;
}