1    :  // Copyright 2015 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/refinery/process_state/process_state_util.h"
  16    :  
  17    :  #include "base/strings/utf_string_conversions.h"
  18    :  #include "syzygy/core/address.h"
  19    :  #include "syzygy/refinery/process_state/layer_traits.h"
  20    :  
  21  m :  namespace refinery {
  22    :  
  23  m :  namespace {
  24    :  
  25  m :  template <typename RecordType>
  26  m :  RecordType* CreateRecord(const AddressRange& range,
  27  m :                           ProcessState* process_state) {
  28  m :    DCHECK(range.IsValid());
  29  m :    DCHECK(process_state);
  30    :  
  31  m :    scoped_refptr<ProcessState::Layer<RecordType>> layer;
  32  m :    process_state->FindOrCreateLayer(&layer);
  33    :  
  34  m :    ProcessState::Layer<RecordType>::RecordPtr record;
  35  m :    layer->CreateRecord(range, &record);
  36    :  
  37  m :    return record->mutable_data();
  38  m :  }
  39    :  
  40  m :  }  // namespace
  41    :  
  42  m :  ModuleLayerAccessor::ModuleLayerAccessor(ProcessState* process_state)
  43  m :      : process_state_(process_state) {
  44  m :    DCHECK(process_state);
  45  m :  }
  46    :  
  47  m :  void ModuleLayerAccessor::AddModuleRecord(const AddressRange& range,
  48  m :                                            const uint32_t checksum,
  49  m :                                            const uint32_t timestamp,
  50  m :                                            const std::wstring& path) {
  51  m :    DCHECK(range.IsValid());
  52    :  
  53    :    // Note: we set the preferred loading address to 0.
  54  m :    pe::PEFile::Signature signature(path, core::AbsoluteAddress(0U), range.size(),
  55  m :                                    checksum, timestamp);
  56    :  
  57  m :    ModuleLayerPtr layer;
  58  m :    process_state_->FindOrCreateLayer(&layer);
  59  m :    ModuleId id = layer->mutable_data()->FindOrIndex(signature);
  60    :  
  61  m :    Module* module_proto = CreateRecord<Module>(range, process_state_);
  62  m :    module_proto->set_module_id(id);
  63  m :  }
  64    :  
  65  m :  bool ModuleLayerAccessor::GetModuleSignature(const Address va,
  66  m :                                               pe::PEFile::Signature* signature) {
  67  m :    DCHECK(signature);
  68    :  
  69    :    // Find the module record corresponding to the virtual address.
  70  m :    ModuleRecordPtr module_record;
  71  m :    if (!process_state_->FindSingleRecord(va, &module_record))
  72  m :      return false;
  73    :  
  74    :    // Retrieve the signature.
  75  m :    const Module& module = module_record->data();
  76  m :    if (!GetModuleSignature(module.module_id(), signature))
  77  m :      return false;
  78    :  
  79    :    // Set the signature's address.
  80  m :    const AddressRange& module_range = module_record->range();
  81  m :    if (!base::IsValueInRangeForNumericType<uint32_t>(module_range.start())) {
  82  m :      LOG(ERROR) << "PE::Signature doesn't support 64bit addresses. Address: "
  83  m :                 << module_range.start();
  84  m :      return false;
  85  m :    }
  86  m :    signature->base_address =
  87  m :        core::AbsoluteAddress(base::checked_cast<uint32_t>(module_range.start()));
  88    :  
  89  m :    return true;
  90  m :  }
  91    :  
  92  m :  bool ModuleLayerAccessor::GetModuleSignature(const ModuleId id,
  93  m :                                               pe::PEFile::Signature* signature) {
  94  m :    DCHECK_NE(kNoModuleId, id);
  95  m :    DCHECK(signature);
  96    :  
  97  m :    ModuleLayerPtr layer;
  98  m :    process_state_->FindOrCreateLayer(&layer);
  99  m :    return layer->data().Find(id, signature);
 100  m :  }
 101    :  
 102  m :  ModuleId ModuleLayerAccessor::GetModuleId(const Address va) {
 103  m :    ModuleRecordPtr module_record;
 104  m :    if (!process_state_->FindSingleRecord(va, &module_record))
 105  m :      return kNoModuleId;
 106  m :    return module_record->data().module_id();
 107  m :  }
 108    :  
 109  m :  ModuleId ModuleLayerAccessor::GetModuleId(
 110  m :      const pe::PEFile::Signature& signature) {
 111  m :    ModuleLayerPtr layer;
 112  m :    process_state_->FindOrCreateLayer(&layer);
 113  m :    return layer->data().Find(signature);
 114  m :  }
 115    :  
 116  m :  bool AddTypedBlockRecord(const AddressRange& range,
 117  m :                           base::StringPiece16 data_name,
 118  m :                           ModuleId module_id,
 119  m :                           TypeId type_id,
 120  m :                           ProcessState* process_state) {
 121  m :    DCHECK(range.IsValid());
 122  m :    DCHECK(process_state);
 123    :  
 124  m :    TypedBlock* typedblock_proto = CreateRecord<TypedBlock>(range, process_state);
 125    :  
 126  m :    std::string data_name_narrow;
 127  m :    if (!base::UTF16ToUTF8(data_name.data(), data_name.size(), &data_name_narrow))
 128  m :      return false;
 129  m :    typedblock_proto->set_data_name(data_name_narrow);
 130    :  
 131  m :    typedblock_proto->set_module_id(module_id);
 132  m :    typedblock_proto->set_type_id(type_id);
 133    :  
 134  m :    return true;
 135  m :  }
 136    :  
 137  m :  }  // namespace refinery