ReycleManger.dll: Rewrote bin change detection code to fix various issues

Plugins/PluginRecycleManager/RecycleManager.cpp

@@ -18,12 +18,10 @@
 
 #define WIN32_LEAN_AND_MEAN
 #include <Windows.h>
-#include <Sddl.h>
 #include <ShellAPI.h>
-#include <ShlObj.h>
 #include <process.h>
-#include <vector>
+#include <cstdio>
-#include "../../Library/RawString.h"
+#include <cstdlib>
 #include "../../Library/Export.h"	// Rainmeter's exported functions
 
 struct MeasureData
@@ -33,30 +31,12 @@ struct MeasureData
 	MeasureData() : count(false) {}
 };
 
-struct BinData
-{
-	union
-	{
-		ULONGLONG lastWrite;
-		UINT lastCount;
-	};
-
-	CRawString directory;
-	WCHAR drive;
-	bool isFAT;
-};
-
 DWORD WINAPI QueryRecycleBinThreadProc(void* pParam);
-HRESULT GetFolderCLSID(LPCWSTR pszPath, CLSID* pathCLSID);
+bool HasRecycleBinChanged();
-LPWSTR GetCurrentUserSid();
-CRawString GetRecycleBinDirectory(WCHAR drive, bool& isFAT);
 
-std::vector<BinData> g_BinData;
 double g_BinCount = 0;
 double g_BinSize = 0;
 
-bool g_IsXP = false;
-
 int g_UpdateCount = 0;
 int g_InstanceCount = 0;
 bool g_Thread = false;
@@ -87,19 +67,6 @@ PLUGIN_EXPORT void Initialize(void** data, void* rm)
 	MeasureData* measure = new MeasureData;
 	*data = measure;
 
-	if (g_InstanceCount == 0)
-	{
-		OSVERSIONINFOEX osvi = {sizeof(OSVERSIONINFOEX)};
-		if (GetVersionEx((OSVERSIONINFO*)&osvi))
-		{
-			if (osvi.dwMajorVersion == 5)
-			{
-				// Not checking for osvi.dwMinorVersion >= 1 because we won't run on pre-XP
-				g_IsXP = true;
-			}
-		}
-	}
-
 	++g_InstanceCount;
 }
 
@@ -128,133 +95,30 @@ PLUGIN_EXPORT double Update(void* data)
 {
 	MeasureData* measure = (MeasureData*)data;
 
-	++g_UpdateCount;
+	if (TryEnterCriticalSection(&g_CriticalSection))
-	if (g_UpdateCount > g_InstanceCount &&
-		TryEnterCriticalSection(&g_CriticalSection))
 	{
-		g_UpdateCount = 0;
+		if (!g_Thread)
-
-		WCHAR buffer[128];
-		DWORD len = GetLogicalDriveStrings(128, buffer);
-
-		// Convert "A:\\\0B:\\\0" -> "AB\0"
-		int index = 0;
-		for (int i = 0; i < len; i += 4)
 		{
-			buffer[index] = buffer[i];
+			++g_UpdateCount;
-			++index;
+			if (g_UpdateCount > g_InstanceCount)
-		}
-		buffer[index] = L'\0';
-
-		const WCHAR DRIVE_HANDLED = 1;
-		bool changed = false;
-		auto iter = g_BinData.begin();
-		while (iter != g_BinData.end())
-		{
-			BinData& data = (*iter);
-
-			WCHAR* pos = wcschr(buffer, data.drive);
-			if (pos != NULL)
 			{
-				if (data.isFAT)
+				if (HasRecycleBinChanged())
 				{
-					// FAT/FAT32 doesn't update directory last write time.
+					// Delay next check.
-					// Use directory content count instead.
+					g_UpdateCount = g_InstanceCount * -2;
-					WCHAR filter[] = L"\0:\\$RECYCLE.BIN\\*";
-					WCHAR filterXP[] = L"\0:\\RECYCLED\\*";
-					filter[0] = *pos;
-					filterXP[0] = *pos;
 
-					WIN32_FIND_DATA fd;
+					DWORD id;
-					HANDLE hSearch = FindFirstFile(g_IsXP ? filterXP : filter, &fd);
+					HANDLE thread = CreateThread(NULL, 0, QueryRecycleBinThreadProc, NULL, 0, &id);
-					if (hSearch != INVALID_HANDLE_VALUE)
+					if (thread)
 					{
-						UINT count = 0;
+						CloseHandle(thread);
-						do
+						g_Thread = true;
-						{
-							++count;
-						}
-						while (FindNextFile(hSearch, &fd));
-
-						FindClose(hSearch);
-
-						if (count != data.lastCount)
-						{
-							data.lastCount = count;
-							changed = true;
-						}
 					}
 				}
-				else if (!data.directory.empty())
+				else
 				{
-					// Prevent the system from displaying message box.
+					g_UpdateCount = 0;
-					UINT oldMode = SetErrorMode(0);
-					SetErrorMode(oldMode | SEM_FAILCRITICALERRORS);
-
-					HANDLE bin = CreateFile(
-						data.directory.c_str(),
-						GENERIC_READ,
-						FILE_SHARE_READ | FILE_SHARE_WRITE,
-						NULL,
-						OPEN_EXISTING,
-						FILE_FLAG_BACKUP_SEMANTICS,
-						NULL);
-
-					if (bin)
-					{
-						ULONGLONG lastWrite;
-						GetFileTime(bin, NULL, NULL, (FILETIME*)&lastWrite);
-						if (data.lastWrite != lastWrite)
-						{
-							data.lastWrite = lastWrite;
-							changed = true;
-						}
-
-						CloseHandle(bin);
-					}
-
-					SetErrorMode(oldMode);
 				}
-
-				*pos = DRIVE_HANDLED;
-				++iter;
-			}
-			else
-			{
-				// Drive removed
-				changed = true;
-				iter = g_BinData.erase(iter);
-			}
-		}
-
-		for (int i = 0; i < index; ++i)
-		{
-			if (buffer[i] != DRIVE_HANDLED)
-			{
-				// New drive
-				g_BinData.push_back(BinData());
-				BinData& data = g_BinData.back();
-				data.drive = buffer[i];
-
-				WCHAR drive[] = L"\0:\\";
-				drive[0] = buffer[i];
-				if (GetDriveType(drive) == DRIVE_FIXED)
-				{
-					data.directory = GetRecycleBinDirectory(buffer[i], data.isFAT);
-				}
-			}
-		}
-
-		if (changed && !g_Thread)
-		{
-			g_UpdateCount = -8;
-
-			DWORD id;
-			HANDLE thread = CreateThread(NULL, 0, QueryRecycleBinThreadProc, NULL, 0, &id);
-			if (thread)
-			{
-				CloseHandle(thread);
-				g_Thread = true;
 			}
 		}
 
@@ -336,140 +200,80 @@ DWORD WINAPI QueryRecycleBinThreadProc(void* pParam)
 	return 0;
 }
 
-HRESULT GetFolderCLSID(LPCWSTR path, CLSID* clsid)
+bool HasRecycleBinChanged()
 {
-	LPITEMIDLIST pidl;
+	static DWORD s_LastVolumeCount = 0;
-	HRESULT hr = SHParseDisplayName(path, NULL, &pidl, 0, NULL);
+	static ULONGLONG s_LastWriteTime = 0;
-	if (SUCCEEDED(hr))
-	{
-		IShellFolder* sf;
-		LPCITEMIDLIST pidlLast;
-		hr = SHBindToParent(pidl, IID_IShellFolder, (void**)&sf, &pidlLast);
-		if (SUCCEEDED(hr))
-		{
-			SHDESCRIPTIONID did;
-			hr = SHGetDataFromIDList(sf, pidlLast, SHGDFIL_DESCRIPTIONID, &did, sizeof(did));
-			*clsid = did.clsid;
 
-			sf->Release();
+	bool changed = false;
+
+	// Check if items have been added to recycle bin since last check.
+	HKEY volumeKey;
+	const WCHAR* subKey = L"Software\\Microsoft\\Windows\\CurrentVersion\\Explorer\\BitBucket\\Volume";
+	LSTATUS ls = RegOpenKeyEx(HKEY_CURRENT_USER, subKey, 0, KEY_QUERY_VALUE | KEY_ENUMERATE_SUB_KEYS, &volumeKey);
+	if (ls == ERROR_SUCCESS)
+	{
+		DWORD volumeCount = 0;
+		RegQueryInfoKey(volumeKey, NULL, NULL, NULL, &volumeCount, NULL, NULL, NULL, NULL, NULL, NULL, NULL);
+		if (volumeCount != s_LastVolumeCount)
+		{
+			s_LastVolumeCount = volumeCount;
+			changed = true;
 		}
 
-		CoTaskMemFree(pidl);
+		WCHAR buffer[64];
-	}
+		DWORD bufferSize = _countof(buffer);
-	return hr;
+		DWORD index = 0;
-}
 
-// Return value must be freed with LocalFree
+		while ((ls = RegEnumKeyEx(volumeKey, index, buffer, &bufferSize, NULL, NULL, NULL, NULL)) == ERROR_SUCCESS)
-LPWSTR GetCurrentUserSid()
-{
-	LPWSTR sidStr = NULL;
-	HANDLE hToken;
-	if (OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &hToken))
-	{
-		DWORD dwBufSize = 0;
-		GetTokenInformation(hToken, TokenUser, NULL, 0, &dwBufSize);
-		BYTE* buf = new BYTE[dwBufSize];
-
-		if (GetTokenInformation(hToken, TokenUser, buf, dwBufSize, &dwBufSize))
 		{
-			TOKEN_USER* tu = (TOKEN_USER*)buf;
+			HKEY volumeSubKey;
-			ConvertSidToStringSid(tu->User.Sid, &sidStr);
+			ls = RegOpenKeyEx(volumeKey, buffer, 0, KEY_QUERY_VALUE, &volumeSubKey);
-		}
+			if (ls == ERROR_SUCCESS)
-
-		delete [] buf;
-
-		CloseHandle(hToken);
-	}
-
-	return sidStr;
-}
-
-CRawString GetRecycleBinDirectory(WCHAR drive, bool& isFAT)
-{
-	WCHAR search[] = L"\0:\\";
-	search[0] = drive;
-
-	WCHAR filesystem[16];
-	BOOL volumeResult = GetVolumeInformation(search, NULL, 0, NULL, NULL, NULL, filesystem, _countof(filesystem));
-
-	if (!volumeResult)
-	{
-		return NULL;
-	}
-
-	if (wcscmp(filesystem, L"NTFS") == 0)
-	{
-		isFAT = false;
-	}
-	else if (wcscmp(filesystem, L"FAT") == 0 || wcscmp(filesystem, L"FAT32") == 0)
-	{
-		isFAT = true;
-	}
-	else
-	{
-		RmLog(LOG_ERROR, L"RecycleManager.dll: Unsupported filesystem");
-		return NULL;
-	}
-
-	const WCHAR* binFolder;
-	if (g_IsXP)
-	{
-		binFolder = isFAT ? L"RECYCLED" : L"RECYCLER";
-	}
-	else
-	{
-		binFolder = L"$RECYCLE.BIN";
-	}
-
-	bool found = false;
-
-	WCHAR binPath[MAX_PATH];
-	_snwprintf_s(binPath, _TRUNCATE, L"%s%s\\", search, binFolder);
-
-	DWORD binAttributes = GetFileAttributes(binPath);
-	if (binAttributes != INVALID_FILE_ATTRIBUTES &&
-		binAttributes & (FILE_ATTRIBUTE_SYSTEM | FILE_ATTRIBUTE_HIDDEN | FILE_ATTRIBUTE_DIRECTORY))
-	{
-		if (isFAT)
-		{
-			if (_waccess(binPath, 0) != -1)
 			{
-				isFAT = true;
+				ULONGLONG lastWriteTime;
-				found = true;
+				ls = RegQueryInfoKey(volumeSubKey, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, (FILETIME*)&lastWriteTime);
-			}
+				if (ls == ERROR_SUCCESS)
-		}
-		else
-		{
-			// Get the correct, SID-specific bin for NTFS
-			LPWSTR currentSid = GetCurrentUserSid();
-			if (currentSid)
-			{
-				wcscat(binPath, currentSid);
-				binAttributes = GetFileAttributes(binPath);
-
-				if (binAttributes != INVALID_FILE_ATTRIBUTES &&
-					binAttributes & FILE_ATTRIBUTE_DIRECTORY)
 				{
-					CLSID id;
+					if (lastWriteTime > s_LastWriteTime)
-					HRESULT hr = GetFolderCLSID(binPath, &id);
-					if (SUCCEEDED(hr) && IsEqualGUID(CLSID_RecycleBin, id))
 					{
-						found = true;
+						s_LastWriteTime = lastWriteTime;
+						changed = true;
 					}
+
+					RegCloseKey(volumeSubKey);
+				}
+			}
+
+			bufferSize = _countof(buffer);
+			++index;
+		}
+
+		RegCloseKey(volumeKey);
+	}
+
+	if (!changed)
+	{
+		// Check if recycle bin has been emptied.
+		HKEY iconKey;
+		const WCHAR* subKey = L"Software\\Microsoft\\Windows\\CurrentVersion\\Explorer\\CLSID\\{645FF040-5081-101B-9F08-00AA002F954E}\\DefaultIcon";
+		ls = RegOpenKeyEx(HKEY_CURRENT_USER, subKey, 0, KEY_QUERY_VALUE, &iconKey);
+		if (ls == ERROR_SUCCESS)
+		{
+			ULONGLONG lastWriteTime;
+			ls = RegQueryInfoKey(iconKey, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, (FILETIME*)&lastWriteTime);
+			if (ls == ERROR_SUCCESS)
+			{
+				if (lastWriteTime > s_LastWriteTime)
+				{
+					s_LastWriteTime = lastWriteTime;
+					changed = true;
 				}
 
-				LocalFree(currentSid);
+				RegCloseKey(iconKey);
 			}
 		}
 	}
 
-	if (!found)
+	return changed;
-	{
-		RmLog(LOG_ERROR, L"RecycleManager.dll: Unable to find bin");
-		return NULL;
-	}
-	else
-	{
-		return binPath;
-	}
 }