#include "stdafx.h"

#define STRSAFE_NO_DEPRECATE
#include <Strsafe.h>

// Used by CopyRegValue specifically for copying
// REG_MULTI_SZ values, because those in particular are
// really funky.
static BOOL
CopyMultiValue(HKEY hKey,
			   PVOID pVal,
			   LPCTSTR szName,
			   DWORD dwSize)
{
	int iStrCount = 0;
	LPTSTR szMulti = NULL,
		ptr;
	BOOL bWasNull = FALSE;
	LONG lRes;
	MultiString *msData = (MultiString *)pVal;

	HANDLE hHeap = GetProcessHeap();

	szMulti = HeapAlloc(hHeap,
		HEAP_ZERO_MEMORY,
		dwSize + sizeof(TCHAR));
	if (!szMulti) return FALSE;

	// Get the multi-string value
	if (!(lRes = RegQueryValueEx(hKey,
		szName,
		0,
		NULL,
		szMulti,
		&dwSize)))
	{
		HeapFree(hHeap, 0, szMulti);
		SetLastError(lRes);
		return FALSE;
	}

	for (ptr = szMulti; (ptr - szMulti) < (LONG)dwSize; ptr++) {
		if (*ptr == 0) {
			if (bWasNull) {
				// Two NULLs in a row, we're done!
				break;
			} else {
				// End of string, add to the list
				bWasNull = TRUE;
				iStrCount++;
			}
		}
		bWasNull = FALSE;
	}

	if (!bWasNull) {
		// Not properly NULL terminated, sigh...
		// Make sure to count the last string.
		iStrCount++;
	}

	if (!MultiString_Init(msData, iStrCount)) {
		HeapFree(hHeap, 0, szMulti);
		return FALSE;
	} else {
		int i;
		for (i = 0, ptr = szMulti; i < iStrCount; i++, ptr += strlen(ptr) + 1) {
			if (!MultiString_Append(msData, ptr)) {
				DWORD err = GetLastError();

				HeapFree(hHeap, 0, szMulti);
				MultiString_Dispose(msData);

				SetLastError(err);
				return FALSE;
			}
		}
	}

	HeapFree(hHeap, 0, szMulti);

	return TRUE;
}

BOOL
CopyRegValue(HKEY hKey,
			 PVOID pVal,
			 LPCTSTR szName,
			 DWORD dwType,
			 DWORD dwSize)
{
	HANDLE hHeap = GetProcessHeap();
	PVOID pData = pVal;
	LONG lRes;

	if (pVal == NULL ||
		szName == NULL ||
		hKey == NULL) {
		return FALSE;
	}

	if (dwType == REG_BINARY ||
		dwType == REG_SZ) {
		pData = HeapAlloc(hHeap, HEAP_ZERO_MEMORY, dwSize);
		if (!pData) return FALSE;
		*(PVOID *)pVal = pData;
	} else if (dwType == REG_MULTI_SZ) {
		// REG_MULTI_SZ is complex enough to warrant its own
		// function
		return CopyMultiValue(hKey, pVal, szName, dwSize);
	}

	if ((lRes = RegQueryValueEx(hKey,
		szName,
		NULL,
		NULL,
		pData,
		&dwSize)) != ERROR_SUCCESS)
	{
		SetLastError(lRes);
		return FALSE;
	}

	return TRUE;
}

// Enumerates across all values in a key, running the given
// RegEnumCallback. Use `pData` to pass application-defined
// data to your enumerator callback.
BOOL RegGetAllValues(HKEY hKey,
					 RegEnumCallback cb,
					 PVOID pData)
{
	LONG lRes;
	TCHAR szValName[512];
	DWORD dwNameLen = 512,
		dwIdx,
		dwValType,
		dwValSize;
	if (!hKey || !cb) return FALSE;

	for (dwIdx = 0;	
		dwNameLen = 512,	// reset dwNameLen first
		!(lRes = RegEnumValue(hKey,
			dwIdx,
			szValName,
			&dwNameLen,
			0,
			&dwValType,
			NULL,
			&dwValSize));
		dwIdx++) {
		// Enforce null termination
		szValName[min(511, dwNameLen)] = 0;
		// pass to callback
		if (!cb(hKey,
				szValName,
				dwValSize,
				dwValType,
				pData)) {
			// Callback returned failure, bail out
			return FALSE;
		}
	}

	if (lRes != ERROR_NO_MORE_ITEMS) {
		// Some enumeration error happened, store it
		// and report the failure
		SetLastError(lRes);
		return FALSE;
	}

	return TRUE;
}

BOOL
RegStoreMultiString(HKEY hKey,
					LPTSTR szValName,
					MultiString *msData)
{
	LPTSTR szMulti;
	DWORD dwMultiSize;
	LONG lRes;
	HANDLE hHeap = GetProcessHeap();

	if (!msData->TotalLength) return FALSE;

	// Total size + second NULL terminator
	dwMultiSize = (DWORD)msData->TotalLength + 1;
	szMulti = HeapAlloc(hHeap,
		HEAP_ZERO_MEMORY,
		dwMultiSize * sizeof(TCHAR));
	if (!szMulti) return FALSE;

	// Copy all the strings into one buffer, as per REG_MULTI_SZ
	{
		DWORD i, dwCurLen;
		for (i = 0, dwCurLen = 0;
			i < msData->Count && dwCurLen < dwMultiSize;
			dwCurLen += (DWORD)_tcslen(msData->Values[i++]) + 1) {
			HRESULT hrRes = StringCchCopy(&szMulti[dwCurLen],
				_tcslen(msData->Values[i]) + 1,
				msData->Values[i]);
			if (FAILED(hrRes)) {
				HeapFree(hHeap, 0, szMulti);
				return FALSE;
			}
		}
	}

	lRes = RegSetValueEx(hKey,
		szValName,
		0,
		REG_MULTI_SZ,
		szMulti,
		dwMultiSize);

	HeapFree(hHeap, 0, szMulti);

	if (lRes != ERROR_SUCCESS) {
		SetLastError(lRes);
		return FALSE;
	}

	return TRUE;
}

// Stores a value in the registry, making assumptions based on the
// key type.
BOOL
RegStoreValue(HKEY hKey,
			  LPTSTR szValName,
			  PVOID pData,
			  DWORD dwValType)
{
	DWORD dwStoreSize;
	LONG lRes;
	if (!hKey || !szValName || !pData) return FALSE;

	switch (dwValType) {
	case REG_SZ:
	case REG_EXPAND_SZ:
		dwStoreSize = (DWORD)_tcslen((LPTSTR)pData) + sizeof(TCHAR);
		break;
	case REG_DWORD:
		dwStoreSize = sizeof(DWORD);
		break;
#ifdef QWORD
	case REG_QWORD:
		dwStoreSize = sizeof(QWORD);
		break;
#endif
	case REG_BINARY:
		// TODO: handle REG_BINARY? will we ever use it?
		OutputDebugString(_T("REG_BINARY attempted to be stored! Unsupported"));
		SetLastError(ERROR_INVALID_PARAMETER);
		return FALSE;
	case REG_MULTI_SZ:
		return RegStoreMultiString(hKey,
			szValName,
			(MultiString *)pData);
	default:
		// ignore
		SetLastError(ERROR_INVALID_PARAMETER);
		return FALSE;
	}

	lRes = RegSetValueEx(hKey,
		szValName,
		0,
		dwValType,
		pData,
		dwStoreSize);
	if (lRes != ERROR_SUCCESS) {
		SetLastError(lRes);
		return FALSE;
	}

	return TRUE;
}