1    :  // Copyright 2012 Google Inc. All Rights Reserved.
   2    :  //
   3    :  // Licensed under the Apache License, Version 2.0 (the "License");
   4    :  // you may not use this file except in compliance with the License.
   5    :  // You may obtain a copy of the License at
   6    :  //
   7    :  //     http://www.apache.org/licenses/LICENSE-2.0
   8    :  //
   9    :  // Unless required by applicable law or agreed to in writing, software
  10    :  // distributed under the License is distributed on an "AS IS" BASIS,
  11    :  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12    :  // See the License for the specific language governing permissions and
  13    :  // limitations under the License.
  14    :  
  15    :  #include "syzygy/agent/common/process_utils.h"
  16    :  
  17    :  #include <algorithm>
  18    :  
  19    :  #include "base/environment.h"
  20    :  #include "base/file_util.h"
  21    :  #include "base/utf_string_conversions.h"
  22    :  #include "base/files/scoped_temp_dir.h"
  23    :  #include "gtest/gtest.h"
  24    :  #include "syzygy/trace/client/rpc_session.h"
  25    :  #include "syzygy/trace/common/unittest_util.h"
  26    :  #include "syzygy/trace/parse/unittest_util.h"
  27    :  
  28    :  namespace agent {
  29    :  namespace common {
  30    :  
  31    :  namespace {
  32    :  
  33    :  using testing::_;
  34    :  using testing::StrictMockParseEventHandler;
  35    :  using trace::parser::Parser;
  36    :  
  37  E :  MATCHER_P(ModuleAtAddress, module, "") {
  38  E :    return arg->module_base_addr == module;
  39  E :  }
  40    :  
  41    :  // TODO(chrisha): Create agent_test_utils, and centralize the basic
  42    :  //     functionality of this fixture. All agents can then use that.
  43    :  //     Don't forget the 'ModuleAtAddress' above as well.
  44    :  class ProcessUtilsTest : public testing::Test {
  45    :   public:
  46  E :    virtual void SetUp() OVERRIDE {
  47  E :      testing::Test::SetUp();
  48    :  
  49    :      // Call trace files will be stuffed here.
  50  E :      ASSERT_TRUE(temp_dir_.CreateUniqueTempDir());
  51    :  
  52  E :      service_.SetEnvironment();
  53  E :    }
  54    :  
  55  E :    virtual void TearDown() OVERRIDE {
  56  E :      service_.Stop();
  57  E :    }
  58    :  
  59  E :    void StartService() {
  60  E :      service_.Start(temp_dir_.path());
  61  E :    }
  62    :  
  63  E :    void StopService() {
  64  E :      service_.Stop();
  65  E :    }
  66    :  
  67  E :    void ReplayLogs(size_t files_expected) {
  68    :      // Stop the service if it's running.
  69  E :      ASSERT_NO_FATAL_FAILURE(StopService());
  70    :  
  71  E :      Parser parser;
  72  E :      ASSERT_TRUE(parser.Init(&handler_));
  73    :  
  74    :      // Queue up the trace file(s) we engendered.
  75    :      file_util::FileEnumerator enumerator(temp_dir_.path(),
  76    :                                           false,
  77  E :                                           file_util::FileEnumerator::FILES);
  78  E :      size_t num_files = 0;
  79  E :      while (true) {
  80  E :        base::FilePath trace_file = enumerator.Next();
  81  E :        if (trace_file.empty())
  82  E :          break;
  83  E :        ASSERT_TRUE(parser.OpenTraceFile(trace_file));
  84  E :        ++num_files;
  85  E :      }
  86    :  
  87  E :      EXPECT_EQ(files_expected, num_files);
  88    :  
  89  E :      if (num_files > 0)
  90  E :        ASSERT_TRUE(parser.Consume());
  91  E :    }
  92    :  
  93    :   protected:
  94    :    // The directory where trace file output will be written.
  95    :    base::ScopedTempDir temp_dir_;
  96    :  
  97    :    // The handler to which the trace file parser will delegate events.
  98    :    StrictMockParseEventHandler handler_;
  99    :  
 100    :    // Our call trace service process instance.
 101    :    testing::CallTraceService service_;
 102    :  };
 103    :  
 104    :  }  // namespace
 105    :  
 106  E :  TEST_F(ProcessUtilsTest, GetProcessModules) {
 107  E :    ModuleVector modules;
 108    :  
 109  E :    GetProcessModules(&modules);
 110    :  
 111    :    // Make sure our own module is in the list.
 112  E :    HMODULE exe_module = ::GetModuleHandle(NULL);
 113    :    EXPECT_TRUE(
 114  E :        std::find(modules.begin(), modules.end(), exe_module) != modules.end());
 115    :  
 116    :    // We have some imports, so there should be
 117    :    // more than just our own module here.
 118  E :    EXPECT_LT(1U, modules.size());
 119  E :  }
 120    :  
 121  E :  TEST_F(ProcessUtilsTest, LogModule) {
 122  E :    ASSERT_NO_FATAL_FAILURE(StartService());
 123    :  
 124    :    // Initialize the session.
 125  E :    trace::client::RpcSession session;
 126  E :    trace::client::TraceFileSegment segment;
 127  E :    std::string id = trace::client::GetInstanceIdForThisModule();
 128  E :    session.set_instance_id(UTF8ToWide(id));
 129  E :    session.CreateSession(&segment);
 130    :  
 131  E :    HMODULE self = ::GetModuleHandle(NULL);
 132  E :    DWORD process_id = ::GetCurrentProcessId();
 133  E :    DWORD thread_id = ::GetCurrentThreadId();
 134    :  
 135  E :    ASSERT_TRUE(LogModule(self, &session, &segment));
 136    :  
 137  E :    ASSERT_NO_FATAL_FAILURE(StopService());
 138    :  
 139  E :    EXPECT_CALL(handler_, OnProcessStarted(_, process_id, _));
 140    :    EXPECT_CALL(handler_, OnProcessAttach(_,
 141    :                                          process_id,
 142    :                                          thread_id,
 143  E :                                          ModuleAtAddress(self)));;
 144  E :    EXPECT_CALL(handler_, OnProcessEnded(_, process_id));
 145    :  
 146  E :    ASSERT_NO_FATAL_FAILURE(ReplayLogs(1));
 147  E :  }
 148    :  
 149    :  }  // namespace common
 150    :  }  // namespace agent