Kill all browser instances at test start up in addition to at tear down which should improve test reliability.

Make running the Chrome Frame integration tests slightly less onerous by not killing non-ChromeFrame instances of Chrome.

BUG=All chrome.exe instances would die when running these tests, which made them annoying to run.
TEST=Only Chrome Frame instances of Chrome will be killed when running the integration tests.

Review URL: http://codereview.chromium.org/467036

git-svn-id: svn://svn.chromium.org/chrome/trunk/src@33991 0039d316-1c4b-4281-b951-d872f2087c98

4 files changed, 192 insertions, 8 deletions

diff --git a/chrome_frame/test/chrome_frame_unittests.cc b/chrome_frame/test/chrome_frame_unittests.cc
index fece220..4d62e74 100644
--- a/chrome_frame/test/chrome_frame_unittests.cc
+++ b/chrome_frame/test/chrome_frame_unittests.cc
@@ -24,6 +24,7 @@
 #include "chrome_frame/crash_reporting/vectored_handler-impl.h"
 #include "chrome_frame/test/chrome_frame_test_utils.h"
 #include "chrome_frame/test_utils.h"
+#include "chrome/common/chrome_switches.h"
 #include "chrome/installer/util/install_util.h"
 #include "chrome/installer/util/helper.h"
 
@@ -78,7 +79,24 @@ _ATL_FUNC_INFO WebBrowserEventSink::kNewWindow3Info = {
 _ATL_FUNC_INFO WebBrowserEventSink::kVoidMethodInfo = {
     CC_STDCALL, VT_EMPTY, 0, {NULL}};
 
+void ChromeFrameTestWithWebServer::CloseAllBrowsers() {
+  // Web browsers tend to relaunch themselves in other processes, meaning the
+  // KillProcess stuff above might not have actually cleaned up all our browser
+  // instances, so make really sure browsers are dead.
+  base::KillProcesses(chrome_frame_test::kIEImageName, 0, NULL);
+  base::KillProcesses(chrome_frame_test::kIEBrokerImageName, 0, NULL);
+  base::KillProcesses(chrome_frame_test::kFirefoxImageName, 0, NULL);
+  base::KillProcesses(chrome_frame_test::kSafariImageName, 0, NULL);
+
+  // Endeavour to only kill off Chrome Frame derived Chrome processes.
+  KillAllNamedProcessesWithArgument(chrome_frame_test::kChromeImageName,
+                                    UTF8ToWide(switches::kChromeFrame));
+}
+
 void ChromeFrameTestWithWebServer::SetUp() {
+  // Make sure our playground is clean before we start.
+  CloseAllBrowsers();
+
   server_.SetUp();
   results_dir_ = server_.GetDataDir();
   file_util::AppendToPath(&results_dir_, L"dump");
@@ -87,14 +105,7 @@ void ChromeFrameTestWithWebServer::SetUp() {
 void ChromeFrameTestWithWebServer::TearDown() {
   CloseBrowser();
 
-  // Web browsers tend to relaunch themselves in other processes, meaning the
-  // KillProcess stuff above might not have actually cleaned up all our browser
-  // instances, so make really sure browsers are dead.
-  base::KillProcesses(chrome_frame_test::kIEImageName, 0, NULL);
-  base::KillProcesses(chrome_frame_test::kIEBrokerImageName, 0, NULL);
-  base::KillProcesses(chrome_frame_test::kFirefoxImageName, 0, NULL);
-  base::KillProcesses(chrome_frame_test::kSafariImageName, 0, NULL);
-  base::KillProcesses(chrome_frame_test::kChromeImageName, 0, NULL);
+  CloseAllBrowsers();
 
   server_.TearDown();
 }
diff --git a/chrome_frame/test/chrome_frame_unittests.h b/chrome_frame/test/chrome_frame_unittests.h
index c39058b..b009faf 100644
--- a/chrome_frame/test/chrome_frame_unittests.h
+++ b/chrome_frame/test/chrome_frame_unittests.h
@@ -63,6 +63,9 @@ class ChromeFrameTestWithWebServer: public testing::Test {
   void VersionTest(BrowserKind browser, const wchar_t* page,
                    const wchar_t* result_file_to_check);
 
+  // Closes all browsers in preparation for a test and during cleanup.
+  void CloseAllBrowsers();
+
   void CloseBrowser();
 
   // Ensures (well, at least tries to ensure) that the browser window has focus.
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),
+                             &params,
+                             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;
+}
diff --git a/chrome_frame/test_utils.h b/chrome_frame/test_utils.h
index 260b4b5..9dc3a7e 100644
--- a/chrome_frame/test_utils.h
+++ b/chrome_frame/test_utils.h
@@ -76,5 +76,11 @@ class DispCallback
   Method method_;
 };
 
+// Kills all running processes named |process_name| that have the string
+// |argument| on their command line. Useful for killing all Chrome Frame
+// instances of Chrome that all have --chrome-frame in their command line.
+bool KillAllNamedProcessesWithArgument(const std::wstring& process_name,
+                                       const std::wstring& argument);
+
 
 #endif  // CHROME_FRAME_TEST_UTILS_H_