The Windows Access Control Model Part 4

/**
*	This code is by OShah. all code will have the following licence.
*	Copyright Shexec32. All code bears the following licence:
**/

/**
*	Copyright Shexec32 2004-2005. All rights reserved.
*
*	THIS CODE AND INFORMATION IS PROVIDED "AS IS" WITHOUT WARRANTY OF
*	ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING BUT NOT LIMITED
*	TO THE IMPLIED WARRANTIES OF MERCHANTABILITY AND/OR FITNESS FOR A
*	PARTICULAR PURPOSE.
**/

#include "SecurityCore.h"
#include "helper.h"
#include "resource2.h"



void CProgressDialog::IncProgressBar(const std::basic_string<TCHAR> pwzString)
{/** This CProgressObj accessor sets the progress bar notification. It's needed because we may not have a progress object.
*	Returns FALSE if the user cancelled.
**/
	this->ProgressCount ++;
	/* Increment the counter even if we don't have a UI */
	if(this->CProgressImpl != NULL)
	{
		if(this->ProgressCount == 1)
		{/* Make sure the lines say the correct things. */
			this->CProgressImpl->SetLine(1, _T("Setting Attributes for file:"), FALSE, NULL);
		}
		this->SetLines(2, pwzString.c_str());
		this->CProgressImpl->SetProgress64(this->ProgressCount, this->Count);
		if(this->CProgressImpl->HasUserCancelled())
		{
			this->Cancelled = TRUE;
		}
	}
	/* If the user has cancelled, inform the caller of this function. */
}



CProgressDialog::CProgressDialog (HWND TheirhWnd /* = NULL */, size_t CountIn /* = 1 */)
	: hWndOwner(TheirhWnd), Count(CountIn), ProgressCount(0), Cancelled(FALSE)
{
	::InitializeCriticalSection(&this->AutoCritSect);
	this->CProgressImpl.GetInterfacePtr() = NULL;
}


CProgressDialog &CProgressDialog::operator=(const CProgressDialog &OldClass)
{
	if(this == &OldClass) return *this;
	hWndOwner = OldClass.hWndOwner;
	Count = OldClass.Count;
	ProgressCount = OldClass.ProgressCount;
	Cancelled = OldClass.Cancelled;
	/* The default equals operator should be adequate for our purposes. It has reference counting semantics. */
	/* CProgressImpl->AddRef(); */
	CProgressImpl = OldClass.CProgressImpl;
	return *this;
}



CProgressDialog::~CProgressDialog()
{/* CProgressImpl.Release(); Unnecessary for auto _com_ptrs */
	::DeleteCriticalSection(&this->AutoCritSect);
}



void CProgressDialog::set_Count(const size_t CountIn)
{
	::EnterCriticalSection(&this->AutoCritSect);
	{
		/* This function is thread safe. */
		this->Count = CountIn;
	}
	::LeaveCriticalSection(&this->AutoCritSect);
}



size_t CProgressDialog::get_Count(void)
{
	/* This function is thread safe. */
	size_t Result = 0;
	::EnterCriticalSection(&this->AutoCritSect);
	{
		Result = this->Count;
	}
	::LeaveCriticalSection(&this->AutoCritSect);
	return Result;
}



BOOL CProgressDialog::get_Cancelled(void)
{
	BOOL Result = FALSE;
	::EnterCriticalSection(&this->AutoCritSect);
	{
		Result = this->Cancelled;
	}
	::LeaveCriticalSection(&this->AutoCritSect);
	return Result;
}




void CProgressDialog::InitProgressDialog(HWND TheirhWnd)
{/** For the progress bar, rshx32 has pointed me to a little gem from Shell32. There is a COM object called
*	IProgressDialog, which does almost all the work of producing a progress bar window for us!
*	It takes multithreading issues, worker thread marshalling, UI responsiveness, animations,
*	all without me not needing to reinvent the wheel!
*
*	This is the Explorer Progress Dialog box that appears during long shell operations. The interface is
*	fully documented in the Shell SDK. Including steps on how to program it.
**/
	/* This dialog is performed in the context of the second thread. */
	this->CProgressImpl.CreateInstance(CLSID_ProgressDialog);
	this->hWndOwner = TheirhWnd;
	/* No interface eh? Then just show nothing (The progress dialog is just an extra nicety, that's all). */
	if(this->CProgressImpl != NULL)
	{/* Initialize the Dialog */
		this->CProgressImpl->SetTitle(L"Setting Attributes for Objects");
		/* Import the video from our DLL. */
		this->CProgressImpl->SetAnimation(::g_hInst, IDD_AVICLIP1);
		this->CProgressImpl->SetCancelMsg(L"Aborting process...", NULL);
		this->CProgressImpl->SetLine(1, L"Beginning Processing...", TRUE, NULL);

		/* And Start the progress dialog box. */
		this->CProgressImpl->StartProgressDialog(TheirhWnd, NULL, PROGDLG_NORMAL | PROGDLG_MODAL | PROGDLG_AUTOTIME, NULL);
	}
}



void CProgressDialog::SetLines(DWORD dwLineNum, const std::basic_string<TCHAR> pwzString)
{/* sets the indicated text. */
	if(this->CProgressImpl != NULL)
        this->CProgressImpl->SetLine(dwLineNum, pwzString.c_str(), TRUE, NULL);
}



void CProgressDialog::Finish(void)
{/* Closes the Progress Dialog. */
	// IProgressDialog *IProgressImpl;
	if(this->CProgressImpl != NULL)
		this->CProgressImpl->StopProgressDialog();
}




BOOL CObjSecurity::GetListOfFiles(std::list< const std::basic_string<TCHAR> > &DirectoryCollection, BOOL fRecurse)
{/** Traverses a directory tree and returns all the files in that directory as a list list.
*	Input: DirectoryCollection contains a number of elements. All elements except the last
*	one will be ignored. The last element should specify the directory to be traversed.
*
*	This method will perform a "DIRCMD /s" on that directory, and will fill up the list with
*	the list of files.
*
*	Returns TRUE for success, FALSE for failure.
*	DirectoryCollection will contain the recursed directory tree (or the best it can come up with).
*	This is a recursive walk.
**/
	DWORD ListErr = 0;
	WIN32_FIND_DATA DirCmdStruct = {0};
	HANDLE hDirCmd = NULL;
	std::basic_string<TCHAR> sRootIter = DirectoryCollection.back(), sRootDir = sRootIter;
	sRootDir.append(_T("\\*"));


	/** Open up a search handle to the directory. **/
	hDirCmd = ::FindFirstFile(sRootDir.c_str(), &DirCmdStruct);
	if(hDirCmd == INVALID_HANDLE_VALUE || hDirCmd == NULL)
	{/* If we cannot enumerate the directory, Leave the rest of the directory alone. */
		return FALSE;
	}

	try {/* protect the loop with an exception handler (fail all in this directory if an exception is thrown.) */
		do {/* Do this at least once */
			std::basic_string<TCHAR> sTmpFName = DirCmdStruct.cFileName;
			if(sTmpFName.compare(_T(".")) != 0 && sTmpFName.compare(_T("..")) != 0 )
			{/* Do not do the '.'s; Then convert all paths to full filenames. */
				sTmpFName = sRootIter + _T("\\") + sTmpFName;
				DWORD FileAttrib = ::GetFileAttributes(sTmpFName.c_str());
				/* Armed with the full pathnames: get its current attributes. */

				if(FileAttrib == INVALID_FILE_ATTRIBUTES)/* E_INVAL: skip the file. */
					continue;

				/* put this entry into the DirectoryCollection. */
				DirectoryCollection.push_back(sTmpFName);

				/* 1.09: Do not walk into junctions, treat them as separate files. */
				if((FileAttrib & FILE_ATTRIBUTE_DIRECTORY) && !(FileAttrib & FILE_ATTRIBUTE_REPARSE_POINT))
				{/* It's a directory: Recall ourselves. And recurse. For failed recursion, continue silently. */
					if(fRecurse)
					{
						this->GetListOfFiles(DirectoryCollection, fRecurse);
					}
					/** This directory is the last entry in the collection. When we recall ourselves, the child
					*	recursion will fill out the list of filenames into DirectoryCollection. When it returns,
					*	The DirectoryCollection will have changed.
					**/
				}
			}

			::SetLastError(ERROR_SUCCESS);
			ListErr = ::FindNextFile(hDirCmd, &DirCmdStruct);
			if(!ListErr)
			{/* Get the next file (we've already done this file). */
				ListErr = ::GetLastError();
				break;
			}
		} while (ListErr != ERROR_NO_MORE_FILES);

		if(hDirCmd != INVALID_HANDLE_VALUE && hDirCmd != NULL)
			::FindClose(hDirCmd); hDirCmd = NULL;	/* You MUST close the file handle before leaving. */
	} catch(...) {/* finally. */
		if(hDirCmd != INVALID_HANDLE_VALUE && hDirCmd != NULL)
			::FindClose(hDirCmd); hDirCmd = NULL;	/* You MUST close the file handle before leaving. */
		throw;
	}
	/* Cleanup. */
	ListErr = 0;
	return TRUE;
}





/** PROGRESS DIALOG THREAD **/
unsigned __stdcall CProgressThread_ThreadFunc(void *ThisPtr)
{
	/* Get the pointer to the class. */
	CProgressThread *ThisClass = reinterpret_cast<CProgressThread *>(ThisPtr);

	/* Initialize COM for this thread. */
	CoAutoInitializer ComAutoInit;

	ThisClass->CProgressImpl.reset(new CProgressDialog);
	{
		MSG msg = {0};

		/* Create a message only window */
		WNDCLASSEX wndClassEx = {0};
		wndClassEx.cbSize = sizeof(WNDCLASSEX);
		wndClassEx.style = CS_GLOBALCLASS;
		wndClassEx.lpfnWndProc = CProgressThread_WndProc;
		wndClassEx.cbClsExtra = 0;
		wndClassEx.cbWndExtra = sizeof(HWND);
		wndClassEx.hInstance = ::g_hInst;
		wndClassEx.hIcon = NULL;
		wndClassEx.hbrBackground = NULL;
		wndClassEx.lpszMenuName = NULL;
		wndClassEx.lpszClassName = _T("MessageOnlyCustomClass");
		wndClassEx.hIconSm = NULL;
		::RegisterClassEx(&wndClassEx);

		HWND stathWnd = ::CreateWindowEx(WS_EX_NOACTIVATE | WS_EX_NOPARENTNOTIFY, _T("MessageOnlyCustomClass"),
			_T("ProgThreadWindow"), WS_DISABLED, 0, 0, 0, 0, NULL, NULL, ::g_hInst, ThisClass);

		::UpdateWindow(stathWnd);
		ThisClass->hWndProgress = stathWnd;

		/* Initialization is now complete, inform the main thread */
		::SetEvent(ThisClass->hStartupMutex);


		/* Now for the Dialog Message Loop. */
		while (::GetMessage(&msg, NULL, 0, 0) > 0)
		{
			::TranslateMessage(&msg);
			::DispatchMessage(&msg);
		}

		::UnregisterClass(_T("MessageOnlyCustomClass"), ::g_hInst);
	}
	ThisClass->DoCleanup();
	ThisClass->CProgressImpl.reset(NULL);
	/* Okay, we're finished, inform the main thread we're about to finish. */

	return 0;
}



/** PROGRESS DIALOG THREAD MESSAGE HANDLER **/
LRESULT CALLBACK CProgressThread_WndProc(HWND ThishWnd, UINT Msg, WPARAM wParam, LPARAM lParam)
{
	CProgressThread *ThisClass;
	if(Msg == WM_CREATE)
	{
		if(!::IsBadReadPtr(reinterpret_cast<void *>(lParam), sizeof(CProgressThread)))
		{
			ThisClass = *reinterpret_cast<CProgressThread **>(lParam);
			::SetProp(ThishWnd, _T("MainClass"), reinterpret_cast<HANDLE>(ThisClass));
		}
		/* initialise the progress bar */
		return DefWindowProc(ThishWnd, Msg, wParam, lParam);
	}
	ThisClass = reinterpret_cast<CProgressThread *>(::GetProp(ThishWnd, _T("MainClass")));
	switch(Msg)
	{
		case CProgressThread::WM_Count:
		{
			ThisClass->CProgressImpl->set_Count(static_cast<size_t>(wParam));
			break;
		}
		case CProgressThread::WM_Finish:
		{
			ThisClass->CProgressImpl->Finish();
			break;
		}

		case CProgressThread::WM_Increment:
		{
			if(::IsBadReadPtr(reinterpret_cast<void *>(lParam), sizeof(const std::basic_string<TCHAR>)) == FALSE)
			{
				/** The other thread allocated a pointer for us, it was new'ed. Once we have the string, it is
				*	our responsibility to delete the pointer.
				**/
				TCHAR *travelledPtr = reinterpret_cast<TCHAR *>(lParam);
				const std::basic_string<TCHAR> CountIn = travelledPtr;
				delete [] travelledPtr;
				ThisClass->CProgressImpl->IncProgressBar(CountIn);
			}
			break;
		}

		case CProgressThread::WM_Initialise:
		{/* Retrieve the HWND from wParam */
			HWND ownerhWnd = reinterpret_cast<HWND>(wParam);
			ThisClass->CProgressImpl->InitProgressDialog(ownerhWnd);
			break;
		}

		case WM_DESTROY:
		{
			::RemoveProp(ThishWnd, _T("MainClass"));
			::PostQuitMessage(0);
			break;
		}
		case WM_CLOSE:
		{
			::DestroyWindow(ThishWnd);
			break;
		}

		default:
		{
			break;
		}
	}
	return DefWindowProc(ThishWnd, Msg, wParam, lParam);
}



CProgressThread::CProgressThread ()
	: hWndProgress(NULL), hThread(NULL), ThrdID(0), CProgressImpl(NULL)
{
	this->ThrdPriority = ::GetThreadPriority(::GetCurrentThread());
	this->hStartupMutex = ::CreateEvent(NULL, TRUE, FALSE, NULL);
	this->Initialize();
}



CProgressThread::~CProgressThread()
{
	if(this->hThread != NULL)
	{/* The thread has had its chance to die. Now kill it. */
		DWORD lpExitCode = 0;
		::GetExitCodeThread(this->hThread, &lpExitCode);
		if(lpExitCode == STILL_ACTIVE)
		{/* We don't return 259 from our thread. */
			/* send a finish */
			this->Finish();
			::PostMessage(this->hWndProgress, WM_CLOSE, 0, 0);
			lpExitCode = ::WaitForSingleObject(hThread, 600000);
			if(lpExitCode != WAIT_OBJECT_0 && ::IsWindow(hWndProgress))
			{/* We gave it a chance to finish, and it didn't want to. Kill it--Leaky leaky */
				::TerminateThread(this->hThread, 0);
			}
		}
		::CloseHandle(this->hThread); hThread = NULL;
		/* The thread should either have gone, or it will die of its own accord. */
	}

	::CloseHandle(this->hStartupMutex); this->hStartupMutex = NULL;
	this->CProgressImpl.reset(NULL);
}

void CProgressThread::DoCleanup(void)
{
	/** IMPORTANT: This method is called in the context of the wrong thread.
	*	Only the CProgressDialog thread understands this.
	*	BUGBUG: Possible Race Condition exists for the thread.
	**/
	::CloseHandle(this->hThread); this->hThread = NULL;
	::ResetEvent(this->hStartupMutex);
	this->ThrdID = 0;
}



void CProgressThread::Initialize(void)
{
	this->CProgressImpl.reset(NULL);
	::ResetEvent(this->hStartupMutex);
	this->hThread = reinterpret_cast<HANDLE>(::_beginthreadex(NULL, 0, CProgressThread_ThreadFunc,
		reinterpret_cast<void*>(this), 0, &this->ThrdID));
	::SetThreadPriority(::GetCurrentThread(), THREAD_PRIORITY_LOWEST);
	/* The thread has completed initialization, and is fully operational. */
	::WaitForSingleObject(this->hStartupMutex, INFINITE);
	/* Tweak our thread priorities. */
}

void CProgressThread::InitProgressDialog(HWND TheirhWnd /* = NULL */)
{
	/* initially, hWndProgress will point to the owner window. After initialisation, the thread is ready. */
	::PostMessage(this->hWndProgress, this->WM_Initialise, reinterpret_cast<WPARAM>(TheirhWnd), 0);
}

void CProgressThread::set_Count(const size_t CountIn)
{
	if(this->CProgressImpl.get() != NULL)
	{
		::PostMessage(this->hWndProgress, this->WM_Count, CountIn, 0);
		this->CProgressImpl->set_Count(CountIn);
	}
}

size_t CProgressThread::get_Count(void) const
{
	/* This function can be called */
	if(this->CProgressImpl.get() != NULL)
		return (this->CProgressImpl->get_Count());
	else return 0;
}

BOOL CProgressThread::IncProgressBar(const std::basic_string<TCHAR> &CountIn)
{
	/**	Since CountIn needs to be read across BOTH threads, it must remain valid across both threads.
	*	Therefore, we must allocate RAM for the other thread. The other thread is responsible for deleting
	*	the data.
	**/
	TCHAR *CountThread = new TCHAR[CountIn.size() + 1];
	CountIn.copy(CountThread, CountIn.size(), 0);
	CountThread[CountIn.size()] = _T('\0');
	::PostMessage(this->hWndProgress, this->WM_Increment, static_cast<WPARAM>(1),
		reinterpret_cast<LPARAM>(CountThread));

	/* We do NOT delete CountThread. that's the responsibility of ProgressThread. */
	return !this->CProgressImpl->get_Cancelled();
}


void CProgressThread::Finish(void)
{
	/* Post a message to the main window */
	::SetThreadPriority(::GetCurrentThread(), THREAD_PRIORITY_NORMAL);
	::PostMessage(this->hWndProgress, this->WM_Finish, 0, 0);
}