//
// Copyright (c) Microsoft Corporation.  All rights reserved.
//
//
// Use of this source code is subject to the terms of the Microsoft end-user
// license agreement (EULA) under which you licensed this SOFTWARE PRODUCT.
// If you did not accept the terms of the EULA, you are not authorized to use
// this source code. For a copy of the EULA, please see the LICENSE.RTF on your
// install media.
//

//
//	Registry Utility Functions
//

#include <windows.h>

BOOL
GetRegistryValue(
                IN  HKEY     hKey,
                IN  TCHAR   *tszSubKeyName, OPTIONAL
                IN  TCHAR   *tszValueName,
                OUT LPBYTE   pValue,
                IN  DWORD    cbValue,
                IN  DWORD    dwRequiredType)
//
//	Read in a registry value.
//
//	Return TRUE if successful, FALSE otherwise.
//
{
    LONG    lResult;
    DWORD   dwType;

    //
    //	If a subkeyname specified, open that key and it becomes the new hKey.
    //
    if (tszSubKeyName) {
        lResult = RegOpenKeyEx(hKey, tszSubKeyName, 0, KEY_READ, &hKey);
        if (lResult != ERROR_SUCCESS) {
            return FALSE;
        }
    }

    lResult = RegQueryValueEx(hKey, tszValueName, NULL, &dwType, pValue, &cbValue);
    if (lResult == ERROR_SUCCESS) {
        if (dwType != dwRequiredType) {
            lResult = !ERROR_SUCCESS;
        }
    }

    if (tszSubKeyName)
        RegCloseKey(hKey);

    return lResult == ERROR_SUCCESS;
}

BOOL
GetRegistryStringValue(
                      IN  HKEY     hKey,
                      IN  TCHAR   *tszSubKeyName, OPTIONAL
                      IN  TCHAR   *tszValueName,
                      OUT TCHAR   *tszValue,
                      IN  DWORD    cbValue)
//
//	Read in a registry value that should have a string type
//
{
    return GetRegistryValue(hKey, tszSubKeyName, tszValueName, (LPBYTE)tszValue, cbValue, REG_SZ);
}

BOOL
GetRegistryMultiSzValue(
                       IN  HKEY     hKey,
                       IN  TCHAR   *tszSubKeyName, OPTIONAL
                       IN  TCHAR   *tszValueName,
                       OUT TCHAR   *tszValue,
                       IN  DWORD    cbValue)
//
//	Read in a registry value that should have a multi string type
//
{
    return GetRegistryValue(hKey, tszSubKeyName, tszValueName, (LPBYTE)tszValue, cbValue, REG_MULTI_SZ);
}

BOOL
GetRegistryDwordValue(
                     IN  HKEY     hKey,
                     IN  TCHAR   *tszSubKeyName, OPTIONAL
                     IN  TCHAR   *tszValueName,
                     OUT DWORD   *pdwValue)
//
//	Read in a registry value that should have a dword
//
{
    return GetRegistryValue(hKey, tszSubKeyName, tszValueName, (LPBYTE)pdwValue, sizeof(*pdwValue), REG_DWORD);
}

BOOL
SetRegistryDwordValue(
                     IN  HKEY     hKey,
                     IN  TCHAR   *tszSubKeyName, OPTIONAL
                     IN  TCHAR   *tszValueName,
                     IN  DWORD    dwValue)
//
//	Set a registry value to a dword.
//
//	Return TRUE if successful, FALSE otherwise.
//
{
    LONG    lResult;

    //
    //	If a subkeyname specified, open or create that key and it becomes the new hKey.
    //
    if (tszSubKeyName) {
        lResult = RegCreateKeyEx(hKey, tszSubKeyName, 0, NULL, 0, KEY_ALL_ACCESS, NULL, &hKey, NULL);
        if (lResult != ERROR_SUCCESS) {
            return FALSE;
        }
    }

    lResult = RegSetValueEx(hKey, tszValueName, 0, REG_DWORD, (LPBYTE)&dwValue, sizeof(DWORD));

    if (tszSubKeyName)
        RegCloseKey(hKey);

    return lResult == ERROR_SUCCESS;
}

//
// dwFlags values:
//
#define REG_ALLOC_MEMORY	(1<<0)

DWORD
RegReadValues(
             IN  HKEY     hKey,
             IN  TCHAR   *tszSubKeyName, OPTIONAL
             IN  PVOID   (*pfMemAlloc)(DWORD), OPTIONAL
             IN  VOID    (*pfMemFree)(PVOID, DWORD), OPTIONAL
             ...)
//
//	Read in a number of registry values.
//
//	For each value to be read in, the following parameters must be passed:
//
//		TCHAR  *tszValueName
//		DWORD	valueType
//		DWORD	dwFlags
//		void   *pValue
//		DWORD	cbValue		// size of space pointed to by pValue
//
// The list is terminated with a NULL tszValueName
//
// Returns the count of the number of values that were read in successfully.
//
{
    LONG    lResult;
    DWORD   dwNumRead = 0,
    dwRegistryType,
    cbRegistryValue,
    dwFlags,
    cbValue,
    *pcbValue,
    dwRequiredType;
    void    *pValue,
    **ppValue;
    va_list argptr;
    PTCHAR  tszValueName;

    //
    // If a subkeyname specified, open that key and it becomes the new hKey.
    //
    if (tszSubKeyName) {
        lResult = RegOpenKeyEx(hKey, tszSubKeyName, 0, KEY_READ, &hKey);
        if (lResult != ERROR_SUCCESS) {
            return dwNumRead;
        }
    }

    va_start(argptr, pfMemFree);
    while (TRUE) {
        tszValueName = va_arg(argptr, PTCHAR);
        if (tszValueName == NULL)
            break;

        dwRequiredType  = va_arg(argptr, DWORD);
        dwFlags         = va_arg(argptr, DWORD);

        pcbValue = NULL;
        if (dwFlags & REG_ALLOC_MEMORY) {
            //
            //	We won't know the size of memory to allocate and read on the first
            //	try, so the first call to RegQueryValueEx will be just to find out
            //	the size.
            //
            pValue          = NULL;
            cbValue         = 0;
            ppValue         = va_arg(argptr, void **);
            pcbValue        = va_arg(argptr, PDWORD);
        } else {
            pValue          = va_arg(argptr, void *);

            if (dwRequiredType == REG_BINARY) {
                pcbValue = va_arg(argptr, PDWORD);
                cbValue = *pcbValue;
            } else
                cbValue         = va_arg(argptr, DWORD);
        }

        //
        //	2 passes used to read stuff from registry:
        //		1st pass: Find out type and size info
        //		2nd pass: Read it in if ok
        //	Problem: type and size could change between the 2 passess...
        //
        //	Find out the type and size of the registry value
        //
        lResult = RegQueryValueEx(hKey, tszValueName, NULL, &dwRegistryType, NULL, &cbRegistryValue);
        if (lResult == ERROR_SUCCESS || (lResult == ERROR_MORE_DATA)) {
            if (dwRegistryType == dwRequiredType) {
                if (dwFlags & REG_ALLOC_MEMORY) {
                    //
                    //	Allocate the memory to hold the object now that we know the size
                    //
                    pValue = pfMemAlloc(cbRegistryValue);
                    if (pValue == NULL) {
                        SetLastError(ERROR_OUTOFMEMORY);
                    }
                } else {
                    //
                    // Make sure that the size of the registry object is not greater than
                    // the space provided.
                    //
                    if (cbRegistryValue > cbValue || !pValue) {
                        if (dwRequiredType == REG_BINARY)
                            *pcbValue = cbRegistryValue;
                        pValue = NULL;
                        SetLastError(ERROR_INSUFFICIENT_BUFFER);
                    }
                }

                if (pValue) {
                    //
                    //	Read in the value data now that we have the memory for it
                    //
                    // Theoretically we should have an iteration, because the data size
                    // could grow between the two calls to ReqQueryValueEx.  At this time
                    // we take the simple approach and just give it one chance to read
                    // successfully.
                    //
                    lResult = RegQueryValueEx(hKey, tszValueName, NULL, NULL, pValue, &cbValue);
                    if (lResult == ERROR_SUCCESS) {
                        dwNumRead++;
                        if (dwRequiredType == REG_BINARY)
                            *pcbValue = cbValue;
                    } else {

                        if (dwFlags & REG_ALLOC_MEMORY) {
                            pfMemFree(pValue, cbRegistryValue);
                            pValue = NULL;
                            cbValue = 0;
                        }
                    }
                }
            } else {
                // Wrong data type
            }
        }

        if (dwFlags & REG_ALLOC_MEMORY) {
            *ppValue = pValue;
            if (pcbValue )
                *pcbValue = cbValue;
        }

    }
    va_end(argptr);

    if (tszSubKeyName)
        RegCloseKey(hKey);

    return dwNumRead;
}