/*
    Tests all Com registered CLSID for common mistakes on QueryInterface
*/

#include <windows.h>
#include <shlwapi.h>
#include <objbase.h>
#include <sstream>
#include <string>

void CheckPtr(IUnknown* pUnk, IUnknown* pTest)
{
	if(pTest)
	{
		if(pTest != (IUnknown*)0xFFFFFFFF)
		{
			std::wcout << L"\tChanged ptr to 0x" << pTest << L" which is " <<
				(pTest == pUnk ? L" the same as " : L" different to ") << L"pUnk" << std::endl;
		}
		else
		{
			std::wcout << L"\tDidn't even return a NULL pointer" << std::endl;
		}
	}
}

void DoIt(REFIID actualClsid, LPCWSTR keyName)
{
	HRESULT hr = S_OK;
	PEXCEPTION_POINTERS pEx = NULL;
	__try
	{
		IUnknown* pUnk = NULL;
		hr = CoCreateInstance(actualClsid, NULL, CLSCTX_ALL, IID_IUnknown, (LPVOID*)&pUnk);
		if(SUCCEEDED(hr))
		{
			IUnknown* pUnk2 = NULL;
			hr = pUnk->QueryInterface(&pUnk2);
			if(hr == S_OK)
			{
				pUnk2->Release();
			}
			else
			{
				std::wcout << L"FAILED: IUnknown query for IUnknown on CLSID " << keyName << ". HRESULT " << std::hex << hr << std::endl;
			}
			IUnknown* pTest = (IUnknown*)0xFFFFFFFF;
			if(SUCCEEDED(hr = pUnk->QueryInterface(IID_NULL, (LPVOID*)&pTest)))
			{
				std::wcout << L"FAILED: QueryInterface for IID_NULL succeeded on CLSID " << keyName << std::endl;
				CheckPtr(pUnk, pTest);
				IID iDummy;
				memset(&iDummy, 0xFF, sizeof(iDummy));
				if(SUCCEEDED(pUnk->QueryInterface(iDummy, (LPVOID*)&pTest)))
				{
					std::wcout << L"\tSucceeded again with IID_Dummy" << std::endl;
					CheckPtr(pUnk, pTest);
				}
			}
			else
			{
				if(pTest != NULL)
				{
					std::wcout << L"FAILED: QueryInterface didn't NULL the output on failure, CLSID: " << keyName << std::endl;
				}
			}
			pUnk->Release();
		}
		else
		{
			std::wcout << L"FAILED: CoCreateInstance failed with error 0x" << std::hex << hr << L" - CLSID " << keyName << std::endl;
		}
	}
	__except(pEx = GetExceptionInformation(), EXCEPTION_EXECUTE_HANDLER)
	{
		PVOID where = pEx->ExceptionRecord->ExceptionAddress;
		std::wcout << L"\tCaught Exception: 0x" << std::hex << pEx->ExceptionRecord->ExceptionCode << L" from 0x" << pEx->ExceptionRecord->ExceptionAddress;
		HMODULE hMod = NULL;
		if(GetModuleHandleEx(GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT | GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS, (LPCWSTR)where, &hMod))
		{
			WCHAR modName[MAX_PATH] = L"";
			GetModuleFileName(hMod, modName, MAX_PATH);
			std::wcout << L"(" << PathFindFileName(modName) << L')';
			std::wcout << L"CLSID - " << keyName;
		}
		std::wcout << std::endl;
	}
}

int __cdecl wmain(int argc, wchar_t** argv)
//int WINAPI WinMain(HINSTANCE, HINSTANCE, LPSTR, int)
{
	DWORD index = 0;
	if(argc >= 2)
	{
		index = _wtoi(argv[1]);
	}
	CoInitialize(NULL);
	HKEY hCLSIDKey = NULL;
	if(RegOpenKeyEx(HKEY_CLASSES_ROOT, L"CLSID", 0, KEY_READ, &hCLSIDKey) == 0)
	{
		WCHAR keyName[100] = {0};
		DWORD nameSize = 100;
		DWORD error = 0;
		if((error = RegEnumKeyEx(hCLSIDKey, index, keyName, &nameSize, NULL, NULL, NULL, NULL)) == ERROR_SUCCESS)
		{
			//std::wcout << L"Processing " << keyName << ':' << std::endl;
			CLSID actualClsid = {0};
			if(SUCCEEDED(IIDFromString(keyName, &actualClsid)))
			{
				DoIt(actualClsid, keyName);
			}
			STARTUPINFO si2 = {sizeof(si2), 0};
			si2.dwFlags = STARTF_USESTDHANDLES;
			si2.hStdError = GetStdHandle(STD_ERROR_HANDLE);
			si2.hStdOutput = GetStdHandle(STD_OUTPUT_HANDLE);
			PROCESS_INFORMATION pi2 = {0};
			std::wostringstream str;
			str << L"\"" << _wpgmptr << L"\" " << ++index;
			std::wstring args = str.str();
			args.push_back(0);
			if(CreateProcess(NULL, &args[0], NULL, NULL, TRUE, 0, NULL, NULL, &si2, &pi2))
			{
				CloseHandle(pi2.hThread);
				CloseHandle(pi2.hProcess);
			}
			else
			{
				std::wcout << "Failed to create process: " << args << L" Error: " << GetLastError();
			}
		}
		RegCloseKey(hCLSIDKey);
	}
	CoUninitialize();
	return 0;
}