Ignore:
Timestamp:
12/11/2008 12:19:53 PM (6 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:
1 edited

Legend:

Unmodified
Added
Removed
  • branches/eraser6/ShellExt/CtxMenu.cpp

    r771 r789  
    467467        std::wstring commandAction; 
    468468        std::wstring commandLine; 
     469        bool commandElevate = false; 
    469470        HRESULT result = E_INVALIDARG; 
    470471        switch (VerbMenuIndices[LOWORD(pCmdInfo->lpVerb)]) 
     
    513514                commandAction = L"addtask"; 
    514515 
     516                //Erasing unused space requires elevation 
     517                commandElevate = true; 
     518 
    515519                //Add every item onto the command line 
    516520                for (std::list<std::wstring>::const_iterator i = SelectedFiles.begin(); 
     
    532536        } 
    533537 
    534         //Execute the command. 
    535         if (!commandLine.empty()) 
    536         { 
    537             //Get the path to this DLL so we can look for Eraser.exe 
    538             wchar_t fileName[MAX_PATH]; 
    539             DWORD fileNameLength = GetModuleFileName(theApp.m_hInstance, fileName, 
    540                 sizeof(fileName) / sizeof(fileName[0])); 
    541             if (!fileNameLength || fileNameLength >= sizeof(fileName) / sizeof(fileName[0])) 
    542                 return E_UNEXPECTED; 
    543              
    544             //Trim to the last \, then append Eraser.exe 
    545             std::wstring eraserPath(fileName, fileNameLength); 
    546             std::wstring::size_type lastBackslash = eraserPath.rfind('\\'); 
    547             if (lastBackslash == std::wstring::npos) 
    548                 return E_INVALIDARG; 
    549  
    550             eraserPath.erase(eraserPath.begin() + lastBackslash + 1, eraserPath.end()); 
    551             if (eraserPath.empty()) 
    552                 return E_UNEXPECTED; 
    553  
    554             eraserPath += L"Eraser.exe"; 
    555  
     538        try 
     539        { 
     540            RunEraser(commandAction, commandLine, commandElevate, pCmdInfo->hwnd, 
     541                pCmdInfo->nShow); 
     542        } 
     543        catch (const std::wstring& e) 
     544        { 
     545            if (!(pCmdInfo->fMask & CMIC_MASK_FLAG_NO_UI)) 
     546            { 
     547                MessageBox(pCmdInfo->hwnd, e.c_str(), LoadString(IDS_ERASERSHELLEXT).c_str(), 
     548                    MB_OK | MB_ICONERROR); 
     549            } 
     550        } 
     551 
     552        return result; 
     553    } 
     554 
     555    CCtxMenu::Actions CCtxMenu::GetApplicableActions() 
     556    { 
     557        unsigned result = 0; 
     558         
     559        //First decide the actions which are applicable to the current invocation 
     560        //reason. 
     561        switch (InvokeReason) 
     562        { 
     563        case INVOKEREASON_RECYCLEBIN: 
     564            result |= ACTION_ERASE | ACTION_ERASE_ON_RESTART; 
     565            break; 
     566        case INVOKEREASON_FILEFOLDER: 
     567            result |= ACTION_ERASE | ACTION_ERASE_ON_RESTART | ACTION_ERASE_UNUSED_SPACE; 
     568        case INVOKEREASON_DRAGDROP: 
     569            result |= ACTION_SECURE_MOVE; 
     570        } 
     571 
     572        //Remove actions that don't apply to the current invocation reason. 
     573        for (std::list<std::wstring>::const_iterator i = SelectedFiles.begin(); 
     574            i != SelectedFiles.end(); ++i) 
     575        { 
     576            //Remove trailing slashes if they are directories. 
     577            std::wstring item(*i); 
     578 
     579            //Check if the path is a path to a volume, if it is not, remove the 
     580            //erase unused space verb. 
     581            wchar_t volumeName[MAX_PATH]; 
     582            if (!GetVolumeNameForVolumeMountPoint(item.c_str(), volumeName, 
     583                sizeof(volumeName) / sizeof(volumeName[0]))) 
     584            { 
     585                result &= ~ACTION_ERASE_UNUSED_SPACE; 
     586            } 
     587        } 
     588 
     589        return static_cast<Actions>(result); 
     590    } 
     591 
     592    std::wstring CCtxMenu::LoadString(UINT stringID) 
     593    { 
     594        //Get a pointer to the buffer containing the string (we're copying it anyway) 
     595        wchar_t* buffer = NULL; 
     596        DWORD lastCount = ::LoadString(theApp.m_hInstance, stringID, 
     597            reinterpret_cast<wchar_t*>(&buffer), 0); 
     598 
     599        if (lastCount > 0) 
     600            return std::wstring(buffer, lastCount); 
     601        return std::wstring(); 
     602    } 
     603 
     604    std::wstring CCtxMenu::FormatString(const std::wstring& formatString, ...) 
     605    { 
     606        std::vector<wchar_t> formatStr(formatString.length() + 1); 
     607        wcscpy_s(&formatStr.front(), formatStr.size(), formatString.c_str()); 
     608 
     609        std::vector<wchar_t> buffer(formatStr.size()); 
     610        for ( ; ; ) 
     611        { 
     612            buffer.resize(buffer.size() * 2); 
     613            va_list arguments; 
     614            va_start(arguments, formatString); 
     615            int result = vswprintf_s(&buffer.front(), buffer.size(), &formatStr.front(), 
     616                arguments); 
     617            va_end(arguments); 
     618 
     619            if (result > 0 && static_cast<unsigned>(result) < buffer.size()) 
     620            { 
     621                break; 
     622            } 
     623        } 
     624 
     625        //Return the result as a wstring 
     626        std::wstring result; 
     627        if (buffer.size() > 0) 
     628            result = &buffer.front(); 
     629        return result; 
     630    } 
     631 
     632    std::wstring CCtxMenu::GetHKeyPath(HKEY handle) 
     633    { 
     634        if (NtQueryKey == NULL) 
     635            NtQueryKey = reinterpret_cast<pNtQueryKey>(GetProcAddress( 
     636                LoadLibrary(L"Ntdll.dll"), "NtQueryKey")); 
     637 
     638        //Keep querying for the key information until enough buffer space has been allocated. 
     639        std::vector<char> buffer(sizeof(KEY_NODE_INFORMATION)); 
     640        NTSTATUS queryResult = STATUS_BUFFER_TOO_SMALL; 
     641        ULONG keyInfoSize = 0; 
     642 
     643        while (queryResult == STATUS_BUFFER_TOO_SMALL || queryResult == STATUS_BUFFER_OVERFLOW) 
     644        { 
     645            buffer.resize(buffer.size() + keyInfoSize); 
     646            ZeroMemory(&buffer.front(), buffer.size()); 
     647            queryResult = NtQueryKey(handle, KeyNodeInformation, &buffer.front(), 
     648                static_cast<ULONG>(buffer.size()), &keyInfoSize); 
     649        } 
     650 
     651        if (queryResult != STATUS_SUCCESS) 
     652            return std::wstring(); 
     653 
     654        KEY_NODE_INFORMATION* keyInfo = reinterpret_cast<KEY_NODE_INFORMATION*>( 
     655            &buffer.front()); 
     656        return keyInfo->Name; 
     657    } 
     658 
     659    void CCtxMenu::RunEraser(const std::wstring& action, const std::wstring& parameters, 
     660        bool elevated, HWND parent, int show) 
     661    { 
     662        //Get the path to this DLL so we can look for Eraser.exe 
     663        wchar_t fileName[MAX_PATH]; 
     664        DWORD fileNameLength = GetModuleFileName(theApp.m_hInstance, fileName, 
     665            sizeof(fileName) / sizeof(fileName[0])); 
     666        if (!fileNameLength || fileNameLength >= sizeof(fileName) / sizeof(fileName[0])) 
     667            throw LoadString(IDS_ERROR_CANNOTFINDERASER); 
     668         
     669        //Trim to the last \, then append Eraser.exe 
     670        std::wstring eraserPath(fileName, fileNameLength); 
     671        std::wstring::size_type lastBackslash = eraserPath.rfind('\\'); 
     672        if (lastBackslash == std::wstring::npos) 
     673            throw LoadString(IDS_ERROR_CANNOTFINDERASER); 
     674 
     675        eraserPath.erase(eraserPath.begin() + lastBackslash + 1, eraserPath.end()); 
     676        if (eraserPath.empty()) 
     677            throw LoadString(IDS_ERROR_CANNOTFINDERASER); 
     678 
     679        eraserPath += L"Eraser.exe"; 
     680 
     681        std::wostringstream finalCmdLine; 
     682        finalCmdLine << L"\"" << eraserPath << L"\" \"" << action << L"\" -q " 
     683                     << parameters; 
     684        std::wstring cmdLine(finalCmdLine.str()); 
     685        std::vector<wchar_t> buffer(cmdLine.length() + 1); 
     686        wcscpy_s(&buffer.front(), cmdLine.length() + 1, cmdLine.c_str()); 
     687 
     688#if 0 
     689        //If the process must be elevated we use ShellExecute with the runas verb 
     690        //to elevate the new process. 
     691        if (elevated) 
     692        { 
     693            int result = reinterpret_cast<int>(ShellExecute(parent, L"runas", 
     694                eraserPath.c_str(), &buffer.front(), NULL, show)); 
     695            if (result <= 32) 
     696                throw LoadString(IDS_ERROR_UNKNOWN); 
     697        } 
     698 
     699        //If the process isn't to be elevated, we use CreateProcess so we can get 
     700        //read the output from the child process 
     701        else 
     702#endif 
     703        { 
    556704            //Create the process. 
    557705            STARTUPINFO startupInfo; 
     
    559707            startupInfo.cb = sizeof(startupInfo); 
    560708            startupInfo.dwFlags = STARTF_USESHOWWINDOW; 
    561             startupInfo.wShowWindow = static_cast<WORD>(pCmdInfo->nShow); 
     709            startupInfo.wShowWindow = static_cast<WORD>(show); 
    562710            startupInfo.hStdInput = startupInfo.hStdOutput = startupInfo.hStdError = 
    563711                INVALID_HANDLE_VALUE; 
     
    581729            PROCESS_INFORMATION processInfo; 
    582730            ZeroMemory(&processInfo, sizeof(processInfo)); 
    583              
    584             std::wostringstream finalCmdLine; 
    585             finalCmdLine << L"\"" << eraserPath << L"\" \"" << commandAction 
    586                          << L"\" -q " << commandLine; 
    587             std::wstring cmdLine(finalCmdLine.str()); 
    588             std::vector<wchar_t> buffer(cmdLine.length() + 1); 
    589             wcscpy_s(&buffer.front(), cmdLine.length() + 1, cmdLine.c_str()); 
    590  
    591731            if (!CreateProcess(NULL, &buffer.front(), NULL, NULL, true, CREATE_NO_WINDOW, 
    592732                NULL, NULL, &startupInfo, &processInfo)) 
    593733            { 
    594                 if (!(pCmdInfo->fMask & CMIC_MASK_FLAG_NO_UI)) 
    595                 { 
    596                     MessageBox(pCmdInfo->hwnd, LoadString(IDS_ERROR_CANNOTFINDERASER).c_str(), 
    597                         LoadString(IDS_ERASERSHELLEXT).c_str(), MB_OK | MB_ICONERROR); 
    598                 } 
    599  
    600                 return E_UNEXPECTED; 
     734                throw LoadString(IDS_ERROR_CANNOTFINDERASER); 
    601735            } 
    602736 
    603737            //Wait for the process to finish. 
    604738            Handle<HANDLE> hProcess(processInfo.hProcess), 
    605                            hThread(processInfo.hThread); 
     739                           hThread(processInfo.hThread); 
    606740            CloseHandle(writePipe); 
    607741            WaitForSingleObject(hProcess, static_cast<DWORD>(-1)); 
     
    619753                    wchar_t wBuffer[8192]; 
    620754                    if (!mbstowcs_s(&lastConvert, wBuffer, sizeof(wBuffer) / sizeof(wBuffer[0]), 
    621                         buffer, lastRead)) 
     755                        buffer, lastRead)) 
    622756                    { 
    623757                        output += std::wstring(wBuffer, lastConvert); 
     
    626760 
    627761                //Show the error message. 
    628                 MessageBox(pCmdInfo->hwnd, output.c_str(), LoadString(IDS_ERASERSHELLEXT).c_str(), 
    629                     MB_OK | MB_ICONERROR); 
    630             } 
    631         } 
    632  
    633         return result; 
    634     } 
    635  
    636     CCtxMenu::Actions CCtxMenu::GetApplicableActions() 
    637     { 
    638         unsigned result = 0; 
    639          
    640         //First decide the actions which are applicable to the current invocation 
    641         //reason. 
    642         switch (InvokeReason) 
    643         { 
    644         case INVOKEREASON_RECYCLEBIN: 
    645             result |= ACTION_ERASE | ACTION_ERASE_ON_RESTART; 
    646             break; 
    647         case INVOKEREASON_FILEFOLDER: 
    648             result |= ACTION_ERASE | ACTION_ERASE_ON_RESTART | ACTION_ERASE_UNUSED_SPACE; 
    649         case INVOKEREASON_DRAGDROP: 
    650             result |= ACTION_SECURE_MOVE; 
    651         } 
    652  
    653         //Remove actions that don't apply to the current invocation reason. 
    654         for (std::list<std::wstring>::const_iterator i = SelectedFiles.begin(); 
    655             i != SelectedFiles.end(); ++i) 
    656         { 
    657             //Remove trailing slashes if they are directories. 
    658             std::wstring item(*i); 
    659  
    660             //Check if the path is a path to a volume, if it is not, remove the 
    661             //erase unused space verb. 
    662             wchar_t volumeName[MAX_PATH]; 
    663             if (!GetVolumeNameForVolumeMountPoint(item.c_str(), volumeName, 
    664                 sizeof(volumeName) / sizeof(volumeName[0]))) 
    665             { 
    666                 result &= ~ACTION_ERASE_UNUSED_SPACE; 
    667             } 
    668         } 
    669  
    670         return static_cast<Actions>(result); 
    671     } 
    672  
    673     std::wstring CCtxMenu::LoadString(UINT stringID) 
    674     { 
    675         //Get a pointer to the buffer containing the string (we're copying it anyway) 
    676         wchar_t* buffer = NULL; 
    677         DWORD lastCount = ::LoadString(theApp.m_hInstance, stringID, 
    678             reinterpret_cast<wchar_t*>(&buffer), 0); 
    679  
    680         if (lastCount > 0) 
    681             return std::wstring(buffer, lastCount); 
    682         return std::wstring(); 
    683     } 
    684  
    685     std::wstring CCtxMenu::FormatString(std::wstring formatString, ...) 
    686     { 
    687         std::vector<wchar_t> formatStr(formatString.length() + 1); 
    688         wcscpy_s(&formatStr.front(), formatStr.size(), formatString.c_str()); 
    689  
    690         std::vector<wchar_t> buffer(formatStr.size()); 
    691         for ( ; ; ) 
    692         { 
    693             buffer.resize(buffer.size() * 2); 
    694             va_list arguments; 
    695             va_start(arguments, formatString); 
    696             int result = vswprintf_s(&buffer.front(), buffer.size(), &formatStr.front(), 
    697                 arguments); 
    698             va_end(arguments); 
    699  
    700             if (result > 0 && static_cast<unsigned>(result) < buffer.size()) 
    701             { 
    702                 break; 
    703             } 
    704         } 
    705  
    706         //Return the result as a wstring 
    707         std::wstring result; 
    708         if (buffer.size() > 0) 
    709             result = &buffer.front(); 
    710         return result; 
    711     } 
    712  
    713     std::wstring CCtxMenu::GetHKeyPath(HKEY handle) 
    714     { 
    715         if (NtQueryKey == NULL) 
    716             NtQueryKey = reinterpret_cast<pNtQueryKey>(GetProcAddress( 
    717                 LoadLibrary(L"Ntdll.dll"), "NtQueryKey")); 
    718  
    719         //Keep querying for the key information until enough buffer space has been allocated. 
    720         std::vector<char> buffer(sizeof(KEY_NODE_INFORMATION)); 
    721         NTSTATUS queryResult = STATUS_BUFFER_TOO_SMALL; 
    722         ULONG keyInfoSize = 0; 
    723  
    724         while (queryResult == STATUS_BUFFER_TOO_SMALL || queryResult == STATUS_BUFFER_OVERFLOW) 
    725         { 
    726             buffer.resize(buffer.size() + keyInfoSize); 
    727             ZeroMemory(&buffer.front(), buffer.size()); 
    728             queryResult = NtQueryKey(handle, KeyNodeInformation, &buffer.front(), 
    729                 static_cast<ULONG>(buffer.size()), &keyInfoSize); 
    730         } 
    731  
    732         if (queryResult != STATUS_SUCCESS) 
    733             return std::wstring(); 
    734  
    735         KEY_NODE_INFORMATION* keyInfo = reinterpret_cast<KEY_NODE_INFORMATION*>( 
    736             &buffer.front()); 
    737         return keyInfo->Name; 
     762                throw output; 
     763            } 
     764        } 
    738765    } 
    739766 
Note: See TracChangeset for help on using the changeset viewer.