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 <vector>
  18    :  
  19    :  #include "base/strings/utf_string_conversions.h"
  20    :  #include "gtest/gtest.h"
  21    :  #include "syzygy/refinery/core/address.h"
  22    :  #include "syzygy/refinery/process_state/layer_data.h"
  23    :  #include "syzygy/refinery/process_state/process_state.h"
  24    :  #include "syzygy/refinery/process_state/refinery.pb.h"
  25    :  #include "syzygy/refinery/types/type.h"
  26    :  
  27  m :  namespace refinery {
  28    :  
  29  m :  namespace {
  30    :  
  31  m :  const Address kAddress = 0x0000CAFE;  // Fits 32-bit.
  32  m :  const Size kSize = 42U;
  33  m :  const uint32_t kChecksum = 11U;
  34  m :  const uint32_t kTimestamp = 22U;
  35  m :  const wchar_t kPath[] = L"c:\\path\\ModuleName";
  36  m :  const char kDataName[] = "data_name";
  37  m :  const ModuleId kModuleId = 100;
  38  m :  const TypeId kTypeId = 42;
  39    :  
  40  m :  }  // namespace
  41    :  
  42  m :  TEST(ModuleLayerAccessorTest, AddModuleRecord) {
  43  m :    ProcessState state;
  44  m :    ModuleLayerAccessor accessor(&state);
  45  m :    accessor.AddModuleRecord(AddressRange(kAddress, kSize), kChecksum, kTimestamp,
  46  m :                             kPath);
  47    :  
  48    :    // Validate a record was added.
  49  m :    ModuleLayerPtr module_layer;
  50  m :    ASSERT_TRUE(state.FindLayer(&module_layer));
  51  m :    std::vector<ModuleRecordPtr> matching_records;
  52  m :    module_layer->GetRecordsAt(kAddress, &matching_records);
  53  m :    ASSERT_EQ(1, matching_records.size());
  54    :  
  55    :    // Validate the record.
  56  m :    ModuleRecordPtr record = matching_records[0];
  57  m :    ASSERT_EQ(AddressRange(kAddress, kSize), record->range());
  58  m :    const Module& module = matching_records[0]->data();
  59  m :    ASSERT_NE(kNoModuleId, module.module_id());
  60    :  
  61    :    // Validate the layer data contains the module information.
  62  m :    pe::PEFile::Signature signature;
  63  m :    ASSERT_TRUE(module_layer->data().Find(module.module_id(), &signature));
  64  m :    ASSERT_EQ(kPath, signature.path);
  65  m :    ASSERT_EQ(0U, signature.base_address.value());
  66  m :    ASSERT_EQ(kSize, signature.module_size);
  67  m :    ASSERT_EQ(kChecksum, signature.module_checksum);
  68  m :    ASSERT_EQ(kTimestamp, signature.module_time_date_stamp);
  69    :  
  70  m :    ASSERT_EQ(module.module_id(), module_layer->data().Find(signature));
  71  m :  }
  72    :  
  73  m :  TEST(ModuleLayerAccessorTest, GetModuleSignatureVATest) {
  74  m :    ProcessState state;
  75  m :    ModuleLayerAccessor accessor(&state);
  76  m :    pe::PEFile::Signature signature;
  77    :  
  78    :    // Fails when VA doesn't correspond to a module.
  79  m :    ASSERT_FALSE(accessor.GetModuleSignature(kAddress, &signature));
  80    :  
  81    :    // Add a module.
  82  m :    accessor.AddModuleRecord(AddressRange(kAddress, kSize), kChecksum, kTimestamp,
  83  m :                             kPath);
  84    :  
  85    :    // Fails outside the module's range.
  86  m :    ASSERT_FALSE(accessor.GetModuleSignature(kAddress - 1, &signature));
  87  m :    ASSERT_FALSE(accessor.GetModuleSignature(kAddress + kSize, &signature));
  88    :  
  89    :    // Succeeds within the module's range.
  90  m :    ASSERT_TRUE(accessor.GetModuleSignature(kAddress, &signature));
  91  m :    ASSERT_TRUE(accessor.GetModuleSignature(kAddress + kSize - 1, &signature));
  92    :  
  93    :    // Validate signature on the last hit.
  94  m :    ASSERT_EQ(kAddress, signature.base_address.value());
  95  m :    ASSERT_EQ(kSize, signature.module_size);
  96  m :    ASSERT_EQ(kChecksum, signature.module_checksum);
  97  m :    ASSERT_EQ(kTimestamp, signature.module_time_date_stamp);
  98  m :    ASSERT_EQ(kPath, signature.path);
  99  m :  }
 100    :  
 101  m :  TEST(ModuleLayerAccessorTest, GetModuleSignatureIdTest) {
 102  m :    ProcessState state;
 103  m :    ModuleLayerAccessor accessor(&state);
 104    :  
 105    :    // Add a module and get its id.
 106  m :    accessor.AddModuleRecord(AddressRange(kAddress, kSize), kChecksum, kTimestamp,
 107  m :                             kPath);
 108  m :    ModuleId module_id = accessor.GetModuleId(kAddress);
 109    :  
 110    :    // Validate.
 111  m :    pe::PEFile::Signature signature;
 112  m :    ASSERT_TRUE(accessor.GetModuleSignature(module_id, &signature));
 113    :  
 114  m :    ASSERT_EQ(0U, signature.base_address.value());
 115  m :    ASSERT_EQ(kSize, signature.module_size);
 116  m :    ASSERT_EQ(kChecksum, signature.module_checksum);
 117  m :    ASSERT_EQ(kTimestamp, signature.module_time_date_stamp);
 118  m :    ASSERT_EQ(kPath, signature.path);
 119  m :  }
 120    :  
 121  m :  TEST(ModuleLayerAccessorTest, GetModuleIdTest) {
 122  m :    ProcessState state;
 123  m :    ModuleLayerAccessor accessor(&state);
 124    :  
 125    :    // Not hitting a module case.
 126  m :    ASSERT_EQ(kNoModuleId, accessor.GetModuleId(kAddress));
 127    :  
 128    :    // Hitting a module case.
 129  m :    accessor.AddModuleRecord(AddressRange(kAddress, kSize), kChecksum, kTimestamp,
 130  m :                             kPath);
 131  m :    ModuleId module_id = accessor.GetModuleId(kAddress);
 132  m :    ASSERT_NE(kNoModuleId, module_id);
 133    :  
 134    :    // Consistency check: the signature associated to module_id must be equal to
 135    :    // that associated with va, up to the base address being 0.
 136  m :    pe::PEFile::Signature sig_from_va;
 137  m :    ASSERT_TRUE(accessor.GetModuleSignature(kAddress, &sig_from_va));
 138  m :    sig_from_va.base_address = core::AbsoluteAddress(0U);
 139    :  
 140  m :    pe::PEFile::Signature sig_from_id;
 141  m :    ModuleLayerPtr layer;
 142  m :    state.FindOrCreateLayer(&layer);
 143  m :    ASSERT_TRUE(layer->data().Find(module_id, &sig_from_id));
 144    :  
 145  m :    ASSERT_EQ(sig_from_va, sig_from_id);
 146  m :  }
 147    :  
 148  m :  TEST(AddTypedBlockRecord, BasicTest) {
 149  m :    ProcessState state;
 150  m :    AddTypedBlockRecord(AddressRange(kAddress, kSize),
 151  m :                        base::ASCIIToUTF16(kDataName), kModuleId, kTypeId,
 152  m :                        &state);
 153    :  
 154    :    // Validate a record was added.
 155  m :    TypedBlockLayerPtr layer;
 156  m :    ASSERT_TRUE(state.FindLayer(&layer));
 157  m :    std::vector<TypedBlockRecordPtr> matching_records;
 158  m :    layer->GetRecordsAt(kAddress, &matching_records);
 159  m :    ASSERT_EQ(1, matching_records.size());
 160    :  
 161    :    // Validate range.
 162  m :    TypedBlockRecordPtr record = matching_records[0];
 163  m :    ASSERT_EQ(AddressRange(kAddress, kSize), record->range());
 164    :  
 165    :    // Validate TypedBlock proto.
 166  m :    TypedBlock* proto = record->mutable_data();
 167  m :    ASSERT_EQ(kDataName, proto->data_name());
 168  m :    ASSERT_EQ(kTypeId, proto->type_id());
 169  m :    ASSERT_EQ(kModuleId, proto->module_id());
 170  m :  }
 171    :  
 172  m :  }  // namespace refinery