diff options
Diffstat (limited to 'chrome_frame/test_utils.cc')
| -rw-r--r-- | chrome_frame/test_utils.cc | 164 |
1 files changed, 164 insertions, 0 deletions
diff --git a/chrome_frame/test_utils.cc b/chrome_frame/test_utils.cc index c5e1f8e..bbe6c86 100644 --- a/chrome_frame/test_utils.cc +++ b/chrome_frame/test_utils.cc @@ -6,11 +6,17 @@ #include <atlbase.h> #include <atlwin.h> +#include <winternl.h> #include "base/file_path.h" #include "base/file_util.h" +#include "base/logging.h" #include "base/path_service.h" +#include "base/process_util.h" +#include "base/scoped_handle.h" +#include "base/string_util.h" #include "chrome/common/chrome_paths.h" +#include "chrome/common/chrome_switches.h" #include "testing/gtest/include/gtest/gtest.h" // Statics @@ -94,3 +100,161 @@ void ScopedChromeFrameRegistrar::RegisterReferenceChromeFrameBuild() { std::wstring ScopedChromeFrameRegistrar::GetChromeFrameDllPath() const { return new_chrome_frame_dll_path_; } + +// TODO(robertshield): The following could be factored out into its own file. +namespace { + +typedef LONG WINAPI +NtQueryInformationProcess( + IN HANDLE ProcessHandle, + IN PROCESSINFOCLASS ProcessInformationClass, + OUT PVOID ProcessInformation, + IN ULONG ProcessInformationLength, + OUT PULONG ReturnLength OPTIONAL +); + +// Get the function pointer to NtQueryInformationProcess in NTDLL.DLL +static bool GetQIP(NtQueryInformationProcess** qip_func_ptr) { + static NtQueryInformationProcess* qip_func = + reinterpret_cast<NtQueryInformationProcess*>( + GetProcAddress(GetModuleHandle(L"ntdll.dll"), + "NtQueryInformationProcess")); + DCHECK(qip_func) << "Could not get pointer to NtQueryInformationProcess."; + *qip_func_ptr = qip_func; + return qip_func != NULL; +} + +// Get the command line of a process +bool GetCommandLineForProcess(uint32 process_id, std::wstring* cmd_line) { + DCHECK(process_id != 0); + DCHECK(cmd_line); + + // Open the process + ScopedHandle process_handle(::OpenProcess( + PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, + false, + process_id)); + if (!process_handle) { + DLOG(ERROR) << "Failed to open process " << process_id << ", last error = " + << GetLastError(); + } + + // Obtain Process Environment Block + NtQueryInformationProcess* qip_func = NULL; + if (process_handle) { + GetQIP(&qip_func); + } + + // Read the address of the process params from the peb. + DWORD process_params_address = 0; + if (qip_func) { + PROCESS_BASIC_INFORMATION info = { 0 }; + // NtQueryInformationProcess returns an NTSTATUS for whom negative values + // are negative. Just check for that instead of pulling in DDK macros. + if ((qip_func(process_handle.Get(), + ProcessBasicInformation, + &info, + sizeof(info), + NULL)) < 0) { + DLOG(ERROR) << "Failed to invoke NtQueryProcessInformation, last error = " + << GetLastError(); + } else { + BYTE* peb = reinterpret_cast<BYTE*>(info.PebBaseAddress); + + // The process command line parameters are (or were once) located at + // the base address of the PEB + 0x10 for 32 bit processes. 64 bit + // processes have a different PEB struct as per + // http://msdn.microsoft.com/en-us/library/aa813706(VS.85).aspx. + // TODO(robertshield): See about doing something about this. + SIZE_T bytes_read = 0; + if (!::ReadProcessMemory(process_handle.Get(), + peb + 0x10, + &process_params_address, + sizeof(process_params_address), + &bytes_read)) { + DLOG(ERROR) << "Failed to read process params address, last error = " + << GetLastError(); + } + } + } + + // Copy all the process parameters into a buffer. + bool success = false; + std::wstring buffer; + if (process_params_address) { + SIZE_T bytes_read; + RTL_USER_PROCESS_PARAMETERS params = { 0 }; + if (!::ReadProcessMemory(process_handle.Get(), + reinterpret_cast<void*>(process_params_address), + ¶ms, + sizeof(params), + &bytes_read)) { + DLOG(ERROR) << "Failed to read RTL_USER_PROCESS_PARAMETERS, " + << "last error = " << GetLastError(); + } else { + // Read the command line parameter + const int max_cmd_line_len = std::min( + static_cast<int>(params.CommandLine.MaximumLength), + 4096); + buffer.resize(max_cmd_line_len + 1); + if (!::ReadProcessMemory(process_handle.Get(), + params.CommandLine.Buffer, + &buffer[0], + max_cmd_line_len, + &bytes_read)) { + DLOG(ERROR) << "Failed to copy process command line, " + << "last error = " << GetLastError(); + } else { + *cmd_line = buffer; + success = true; + } + } + } + + return success; +} + +// Used to filter processes by process ID. +class ArgumentFilter : public base::ProcessFilter { + public: + explicit ArgumentFilter(const std::wstring& argument) + : argument_to_find_(argument) {} + + // Returns true to indicate set-inclusion and false otherwise. This method + // should not have side-effects and should be idempotent. + virtual bool Includes(base::ProcessId pid, base::ProcessId parent_pid) const { + bool found = false; + std::wstring command_line; + if (GetCommandLineForProcess(pid, &command_line)) { + std::wstring::const_iterator it = + std::search(command_line.begin(), + command_line.end(), + argument_to_find_.begin(), + argument_to_find_.end(), + CaseInsensitiveCompareASCII<wchar_t>()); + found = (it != command_line.end()); + } + return found; + } + + protected: + std::wstring argument_to_find_; +}; + +} // namespace + +bool KillAllNamedProcessesWithArgument(const std::wstring& process_name, + const std::wstring& argument) { + bool result = true; + const ProcessEntry* entry; + ArgumentFilter filter(argument); + base::NamedProcessIterator iter(process_name, &filter); + while (entry = iter.NextProcessEntry()) { + if (!base::KillProcessById((*entry).th32ProcessID, 0, true)) { + DLOG(ERROR) << "Failed to kill process " << (*entry).th32ProcessID; + result = false; + } + } + + return result; +} |