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 <psapi.h>
  18    :  
  19    :  #include "base/logging.h"
  20    :  #include "base/win/pe_image.h"
  21    :  #include "sawbuck/common/com_utils.h"
  22    :  #include "syzygy/common/path_util.h"
  23    :  #include "syzygy/trace/client/rpc_session.h"
  24    :  
  25    :  namespace agent {
  26    :  namespace common {
  27    :  
  28    :  namespace {
  29    :  
  30    :  // Accessing a module acquired from process iteration calls is inherently racy,
  31    :  // as we don't hold any kind of reference to the module, and so the module
  32    :  // could be unloaded while we're accessing it. In practice this shouldn't
  33    :  // happen to us, as we'll be running under the loader's lock in all cases.
  34    :  bool CaptureModuleInformation(const IMAGE_NT_HEADERS* nt_headers,
  35  i :                                TraceModuleData* module_event) {
  36  i :    DCHECK(nt_headers != NULL);
  37  i :    DCHECK(module_event != NULL);
  38    :  
  39  i :    __try {
  40  i :      module_event->module_base_size = nt_headers->OptionalHeader.SizeOfImage;
  41  i :      module_event->module_checksum = nt_headers->OptionalHeader.CheckSum;
  42  i :      module_event->module_time_date_stamp = nt_headers->FileHeader.TimeDateStamp;
  43  i :    } __except(EXCEPTION_EXECUTE_HANDLER) {
  44  i :      return false;
  45  i :    }
  46    :  
  47  i :    return true;
  48  i :  }
  49    :  
  50    :  }  // namespace
  51    :  
  52  E :  void GetProcessModules(ModuleVector* modules) {
  53  E :    DCHECK(modules != NULL);
  54    :  
  55  E :    modules->resize(128);
  56  E :    while (true) {
  57  E :      DWORD bytes = sizeof(modules->at(0)) * modules->size();
  58  E :      DWORD needed_bytes = 0;
  59    :      BOOL success = ::EnumProcessModules(::GetCurrentProcess(),
  60    :                                          &modules->at(0),
  61    :                                          bytes,
  62  E :                                          &needed_bytes);
  63  E :      if (success && bytes >= needed_bytes) {
  64    :        // Success - break out of the loop.
  65    :        // Resize our module vector to the returned size.
  66  E :        modules->resize(needed_bytes / sizeof(modules->at(0)));
  67  E :        return;
  68    :      }
  69    :  
  70    :      // Resize our module vector with the needed size and little slop.
  71  i :      modules->resize(needed_bytes / sizeof(modules->at(0)) + 4);
  72  i :    }
  73  E :  }
  74    :  
  75    :  bool LogModule(HMODULE module,
  76    :                 trace::client::RpcSession* session,
  77  E :                 trace::client::TraceFileSegment* segment) {
  78  E :    DCHECK(module != NULL);
  79  E :    DCHECK(session != NULL);
  80  E :    DCHECK(segment != NULL);
  81    :  
  82    :    // Make sure the event we're about to write will fit.
  83    :    if (!segment->CanAllocate(sizeof(TraceModuleData)) ||
  84  E :        !session->ExchangeBuffer(segment)) {
  85    :      // Failed to allocate a new segment.
  86  i :      LOG(ERROR) << "Failed to exchange buffer.";
  87  i :      return false;
  88    :    }
  89    :  
  90  E :    DCHECK(segment->CanAllocate(sizeof(TraceModuleData)));
  91    :  
  92    :    // Allocate a record in the log.
  93    :    TraceModuleData* module_event = reinterpret_cast<TraceModuleData*>(
  94    :        segment->AllocateTraceRecordImpl(
  95  E :            TRACE_PROCESS_ATTACH_EVENT, sizeof(TraceModuleData)));
  96  E :    DCHECK(module_event != NULL);
  97    :  
  98    :    // Populate the log record.
  99  E :    base::win::PEImage image(module);
 100  E :    module_event->module_base_addr = module;
 101  E :    if (!CaptureModuleInformation(image.GetNTHeaders(), module_event)) {
 102  i :      LOG(ERROR) << "Failed to capture module information.";
 103  i :      return false;
 104    :    }
 105    :  
 106  E :    wchar_t module_name[MAX_PATH] = { 0 };
 107    :    if (::GetMappedFileName(::GetCurrentProcess(), module,
 108  E :                            module_name, arraysize(module_name)) == 0) {
 109  i :      DWORD error = ::GetLastError();
 110  i :      LOG(ERROR) << "Failed to get module name: " << com::LogWe(error) << ".";
 111  i :      return false;
 112    :    }
 113  E :    base::FilePath device_path(module_name);
 114  E :    base::FilePath drive_path;
 115  E :    if (!::common::ConvertDevicePathToDrivePath(device_path, &drive_path)) {
 116  i :      LOG(ERROR) << "ConvertDevicePathToDrivePath failed.";
 117  i :      return false;
 118    :    }
 119    :    ::wcsncpy(module_event->module_name, drive_path.value().c_str(),
 120  E :              arraysize(module_event->module_name));
 121    :  
 122  E :    module_event->module_exe[0] = L'\0';
 123    :  
 124  E :    return true;
 125  E :  }
 126    :  
 127    :  }  // namespace common
 128    :  }  // namespace agent