Changeset 789 for branches/eraser6/ShellExt/CtxMenu.cpp

Timestamp:

12/11/2008 12:19:53 PM (4 years ago)

Author:

lowjoel

Message:

-Replace all functions taking a string parameter to take a const reference to reduce copying
-Factor out the call to launch Eraser, allowing for the spawned instance to run elevated (though it is currently unimplemented)

File:

• branches/eraser6/ShellExt/CtxMenu.cpp (modified) (7 diffs)

branches/eraser6/ShellExt/CtxMenu.cpp

@@ -467 +467 @@
         std::wstring commandAction;
         std::wstring commandLine;
+        bool commandElevate = false;
         HRESULT result = E_INVALIDARG;
         switch (VerbMenuIndices[LOWORD(pCmdInfo->lpVerb)])
@@ -513 +514 @@
                 commandAction = L"addtask";
 
+                //Erasing unused space requires elevation
+                commandElevate = true;
+
                 //Add every item onto the command line
                 for (std::list<std::wstring>::const_iterator i = SelectedFiles.begin();
@@ -532 +536 @@
         }
 
-        //Execute the command.
-        if (!commandLine.empty())
-        {
-            //Get the path to this DLL so we can look for Eraser.exe
-            wchar_t fileName[MAX_PATH];
-            DWORD fileNameLength = GetModuleFileName(theApp.m_hInstance, fileName,
-                sizeof(fileName) / sizeof(fileName[0]));
-            if (!fileNameLength || fileNameLength >= sizeof(fileName) / sizeof(fileName[0]))
-                return E_UNEXPECTED;
-            
-            //Trim to the last \, then append Eraser.exe
-            std::wstring eraserPath(fileName, fileNameLength);
-            std::wstring::size_type lastBackslash = eraserPath.rfind('\\');
-            if (lastBackslash == std::wstring::npos)
-                return E_INVALIDARG;
-
-            eraserPath.erase(eraserPath.begin() + lastBackslash + 1, eraserPath.end());
-            if (eraserPath.empty())
-                return E_UNEXPECTED;
-
-            eraserPath += L"Eraser.exe";
-
+        try
+        {
+            RunEraser(commandAction, commandLine, commandElevate, pCmdInfo->hwnd,
+                pCmdInfo->nShow);
+        }
+        catch (const std::wstring& e)
+        {
+            if (!(pCmdInfo->fMask & CMIC_MASK_FLAG_NO_UI))
+            {
+                MessageBox(pCmdInfo->hwnd, e.c_str(), LoadString(IDS_ERASERSHELLEXT).c_str(),
+                    MB_OK | MB_ICONERROR);
+            }
+        }
+
+        return result;
+    }
+
+    CCtxMenu::Actions CCtxMenu::GetApplicableActions()
+    {
+        unsigned result = 0;
+        
+        //First decide the actions which are applicable to the current invocation
+        //reason.
+        switch (InvokeReason)
+        {
+        case INVOKEREASON_RECYCLEBIN:
+            result |= ACTION_ERASE | ACTION_ERASE_ON_RESTART;
+            break;
+        case INVOKEREASON_FILEFOLDER:
+            result |= ACTION_ERASE | ACTION_ERASE_ON_RESTART | ACTION_ERASE_UNUSED_SPACE;
+        case INVOKEREASON_DRAGDROP:
+            result |= ACTION_SECURE_MOVE;
+        }
+
+        //Remove actions that don't apply to the current invocation reason.
+        for (std::list<std::wstring>::const_iterator i = SelectedFiles.begin();
+            i != SelectedFiles.end(); ++i)
+        {
+            //Remove trailing slashes if they are directories.
+            std::wstring item(*i);
+
+            //Check if the path is a path to a volume, if it is not, remove the
+            //erase unused space verb.
+            wchar_t volumeName[MAX_PATH];
+            if (!GetVolumeNameForVolumeMountPoint(item.c_str(), volumeName,
+                sizeof(volumeName) / sizeof(volumeName[0])))
+            {
+                result &= ~ACTION_ERASE_UNUSED_SPACE;
+            }
+        }
+
+        return static_cast<Actions>(result);
+    }
+
+    std::wstring CCtxMenu::LoadString(UINT stringID)
+    {
+        //Get a pointer to the buffer containing the string (we're copying it anyway)
+        wchar_t* buffer = NULL;
+        DWORD lastCount = ::LoadString(theApp.m_hInstance, stringID,
+            reinterpret_cast<wchar_t*>(&buffer), 0);
+
+        if (lastCount > 0)
+            return std::wstring(buffer, lastCount);
+        return std::wstring();
+    }
+
+    std::wstring CCtxMenu::FormatString(const std::wstring& formatString, ...)
+    {
+        std::vector<wchar_t> formatStr(formatString.length() + 1);
+        wcscpy_s(&formatStr.front(), formatStr.size(), formatString.c_str());
+
+        std::vector<wchar_t> buffer(formatStr.size());
+        for ( ; ; )
+        {
+            buffer.resize(buffer.size() * 2);
+            va_list arguments;
+            va_start(arguments, formatString);
+            int result = vswprintf_s(&buffer.front(), buffer.size(), &formatStr.front(),
+                arguments);
+            va_end(arguments);
+
+            if (result > 0 && static_cast<unsigned>(result) < buffer.size())
+            {
+                break;
+            }
+        }
+
+        //Return the result as a wstring
+        std::wstring result;
+        if (buffer.size() > 0)
+            result = &buffer.front();
+        return result;
+    }
+
+    std::wstring CCtxMenu::GetHKeyPath(HKEY handle)
+    {
+        if (NtQueryKey == NULL)
+            NtQueryKey = reinterpret_cast<pNtQueryKey>(GetProcAddress(
+                LoadLibrary(L"Ntdll.dll"), "NtQueryKey"));
+
+        //Keep querying for the key information until enough buffer space has been allocated.
+        std::vector<char> buffer(sizeof(KEY_NODE_INFORMATION));
+        NTSTATUS queryResult = STATUS_BUFFER_TOO_SMALL;
+        ULONG keyInfoSize = 0;
+
+        while (queryResult == STATUS_BUFFER_TOO_SMALL || queryResult == STATUS_BUFFER_OVERFLOW)
+        {
+            buffer.resize(buffer.size() + keyInfoSize);
+            ZeroMemory(&buffer.front(), buffer.size());
+            queryResult = NtQueryKey(handle, KeyNodeInformation, &buffer.front(),
+                static_cast<ULONG>(buffer.size()), &keyInfoSize);
+        }
+
+        if (queryResult != STATUS_SUCCESS)
+            return std::wstring();
+
+        KEY_NODE_INFORMATION* keyInfo = reinterpret_cast<KEY_NODE_INFORMATION*>(
+            &buffer.front());
+        return keyInfo->Name;
+    }
+
+    void CCtxMenu::RunEraser(const std::wstring& action, const std::wstring& parameters,
+        bool elevated, HWND parent, int show)
+    {
+        //Get the path to this DLL so we can look for Eraser.exe
+        wchar_t fileName[MAX_PATH];
+        DWORD fileNameLength = GetModuleFileName(theApp.m_hInstance, fileName,
+            sizeof(fileName) / sizeof(fileName[0]));
+        if (!fileNameLength || fileNameLength >= sizeof(fileName) / sizeof(fileName[0]))
+            throw LoadString(IDS_ERROR_CANNOTFINDERASER);
+        
+        //Trim to the last \, then append Eraser.exe
+        std::wstring eraserPath(fileName, fileNameLength);
+        std::wstring::size_type lastBackslash = eraserPath.rfind('\\');
+        if (lastBackslash == std::wstring::npos)
+            throw LoadString(IDS_ERROR_CANNOTFINDERASER);
+
+        eraserPath.erase(eraserPath.begin() + lastBackslash + 1, eraserPath.end());
+        if (eraserPath.empty())
+            throw LoadString(IDS_ERROR_CANNOTFINDERASER);
+
+        eraserPath += L"Eraser.exe";
+
+        std::wostringstream finalCmdLine;
+        finalCmdLine << L"\"" << eraserPath << L"\" \"" << action << L"\" -q "
+                     << parameters;
+        std::wstring cmdLine(finalCmdLine.str());
+        std::vector<wchar_t> buffer(cmdLine.length() + 1);
+        wcscpy_s(&buffer.front(), cmdLine.length() + 1, cmdLine.c_str());
+
+#if 0
+        //If the process must be elevated we use ShellExecute with the runas verb
+        //to elevate the new process.
+        if (elevated)
+        {
+            int result = reinterpret_cast<int>(ShellExecute(parent, L"runas",
+                eraserPath.c_str(), &buffer.front(), NULL, show));
+            if (result <= 32)
+                throw LoadString(IDS_ERROR_UNKNOWN);
+        }
+
+        //If the process isn't to be elevated, we use CreateProcess so we can get
+        //read the output from the child process
+        else
+#endif
+        {
             //Create the process.
             STARTUPINFO startupInfo;
@@ -559 +707 @@
             startupInfo.cb = sizeof(startupInfo);
             startupInfo.dwFlags = STARTF_USESHOWWINDOW;
-            startupInfo.wShowWindow = static_cast<WORD>(pCmdInfo->nShow);
+            startupInfo.wShowWindow = static_cast<WORD>(show);
             startupInfo.hStdInput = startupInfo.hStdOutput = startupInfo.hStdError =
                 INVALID_HANDLE_VALUE;
@@ -581 +729 @@
             PROCESS_INFORMATION processInfo;
             ZeroMemory(&processInfo, sizeof(processInfo));
-            
-            std::wostringstream finalCmdLine;
-            finalCmdLine << L"\"" << eraserPath << L"\" \"" << commandAction
-                         << L"\" -q " << commandLine;
-            std::wstring cmdLine(finalCmdLine.str());
-            std::vector<wchar_t> buffer(cmdLine.length() + 1);
-            wcscpy_s(&buffer.front(), cmdLine.length() + 1, cmdLine.c_str());
-
             if (!CreateProcess(NULL, &buffer.front(), NULL, NULL, true, CREATE_NO_WINDOW,
                 NULL, NULL, &startupInfo, &processInfo))
             {
-                if (!(pCmdInfo->fMask & CMIC_MASK_FLAG_NO_UI))
-                {
-                    MessageBox(pCmdInfo->hwnd, LoadString(IDS_ERROR_CANNOTFINDERASER).c_str(),
-                        LoadString(IDS_ERASERSHELLEXT).c_str(), MB_OK | MB_ICONERROR);
-                }
-
-                return E_UNEXPECTED;
+                throw LoadString(IDS_ERROR_CANNOTFINDERASER);
             }
 
             //Wait for the process to finish.
             Handle<HANDLE> hProcess(processInfo.hProcess),
-                           hThread(processInfo.hThread);
+                           hThread(processInfo.hThread);
             CloseHandle(writePipe);
             WaitForSingleObject(hProcess, static_cast<DWORD>(-1));
@@ -619 +753 @@
                     wchar_t wBuffer[8192];
                     if (!mbstowcs_s(&lastConvert, wBuffer, sizeof(wBuffer) / sizeof(wBuffer[0]),
-                        buffer, lastRead))
+                        buffer, lastRead))
                     {
                         output += std::wstring(wBuffer, lastConvert);
@@ -626 +760 @@
 
                 //Show the error message.
-                MessageBox(pCmdInfo->hwnd, output.c_str(), LoadString(IDS_ERASERSHELLEXT).c_str(),
-                    MB_OK | MB_ICONERROR);
-            }
-        }
-
-        return result;
-    }
-
-    CCtxMenu::Actions CCtxMenu::GetApplicableActions()
-    {
-        unsigned result = 0;
-        
-        //First decide the actions which are applicable to the current invocation
-        //reason.
-        switch (InvokeReason)
-        {
-        case INVOKEREASON_RECYCLEBIN:
-            result |= ACTION_ERASE | ACTION_ERASE_ON_RESTART;
-            break;
-        case INVOKEREASON_FILEFOLDER:
-            result |= ACTION_ERASE | ACTION_ERASE_ON_RESTART | ACTION_ERASE_UNUSED_SPACE;
-        case INVOKEREASON_DRAGDROP:
-            result |= ACTION_SECURE_MOVE;
-        }
-
-        //Remove actions that don't apply to the current invocation reason.
-        for (std::list<std::wstring>::const_iterator i = SelectedFiles.begin();
-            i != SelectedFiles.end(); ++i)
-        {
-            //Remove trailing slashes if they are directories.
-            std::wstring item(*i);
-
-            //Check if the path is a path to a volume, if it is not, remove the
-            //erase unused space verb.
-            wchar_t volumeName[MAX_PATH];
-            if (!GetVolumeNameForVolumeMountPoint(item.c_str(), volumeName,
-                sizeof(volumeName) / sizeof(volumeName[0])))
-            {
-                result &= ~ACTION_ERASE_UNUSED_SPACE;
-            }
-        }
-
-        return static_cast<Actions>(result);
-    }
-
-    std::wstring CCtxMenu::LoadString(UINT stringID)
-    {
-        //Get a pointer to the buffer containing the string (we're copying it anyway)
-        wchar_t* buffer = NULL;
-        DWORD lastCount = ::LoadString(theApp.m_hInstance, stringID,
-            reinterpret_cast<wchar_t*>(&buffer), 0);
-
-        if (lastCount > 0)
-            return std::wstring(buffer, lastCount);
-        return std::wstring();
-    }
-
-    std::wstring CCtxMenu::FormatString(std::wstring formatString, ...)
-    {
-        std::vector<wchar_t> formatStr(formatString.length() + 1);
-        wcscpy_s(&formatStr.front(), formatStr.size(), formatString.c_str());
-
-        std::vector<wchar_t> buffer(formatStr.size());
-        for ( ; ; )
-        {
-            buffer.resize(buffer.size() * 2);
-            va_list arguments;
-            va_start(arguments, formatString);
-            int result = vswprintf_s(&buffer.front(), buffer.size(), &formatStr.front(),
-                arguments);
-            va_end(arguments);
-
-            if (result > 0 && static_cast<unsigned>(result) < buffer.size())
-            {
-                break;
-            }
-        }
-
-        //Return the result as a wstring
-        std::wstring result;
-        if (buffer.size() > 0)
-            result = &buffer.front();
-        return result;
-    }
-
-    std::wstring CCtxMenu::GetHKeyPath(HKEY handle)
-    {
-        if (NtQueryKey == NULL)
-            NtQueryKey = reinterpret_cast<pNtQueryKey>(GetProcAddress(
-                LoadLibrary(L"Ntdll.dll"), "NtQueryKey"));
-
-        //Keep querying for the key information until enough buffer space has been allocated.
-        std::vector<char> buffer(sizeof(KEY_NODE_INFORMATION));
-        NTSTATUS queryResult = STATUS_BUFFER_TOO_SMALL;
-        ULONG keyInfoSize = 0;
-
-        while (queryResult == STATUS_BUFFER_TOO_SMALL || queryResult == STATUS_BUFFER_OVERFLOW)
-        {
-            buffer.resize(buffer.size() + keyInfoSize);
-            ZeroMemory(&buffer.front(), buffer.size());
-            queryResult = NtQueryKey(handle, KeyNodeInformation, &buffer.front(),
-                static_cast<ULONG>(buffer.size()), &keyInfoSize);
-        }
-
-        if (queryResult != STATUS_SUCCESS)
-            return std::wstring();
-
-        KEY_NODE_INFORMATION* keyInfo = reinterpret_cast<KEY_NODE_INFORMATION*>(
-            &buffer.front());
-        return keyInfo->Name;
+                throw output;
+            }
+        }
     }