1    :  // Copyright 2012 Google Inc.
   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/instrument/mutators/add_bb_ranges_stream.h"
  16    :  
  17    :  #include "gmock/gmock.h"
  18    :  #include "gtest/gtest.h"
  19    :  #include "syzygy/common/basic_block_frequency_data.h"
  20    :  #include "syzygy/pdb/pdb_byte_stream.h"
  21    :  #include "syzygy/pdb/unittest_util.h"
  22    :  
  23    :  namespace instrument {
  24    :  namespace mutators {
  25    :  
  26    :  namespace {
  27    :  
  28    :  typedef AddBasicBlockRangesStreamPdbMutator::RelativeAddressRange
  29    :      RelativeAddressRange;
  30    :  typedef AddBasicBlockRangesStreamPdbMutator::RelativeAddressRangeVector
  31    :      RelativeAddressRangeVector;
  32    :  
  33    :  using core::RelativeAddress;
  34    :  using common::kBasicBlockRangesStreamName;
  35    :  
  36    :  }  // namespace
  37    :  
  38  E :  TEST(AddBasicBlockRangesStreamPdbMutatorTest, FailsIfStreamAlreadyExists) {
  39  E :    RelativeAddressRangeVector bb_addresses;
  40  E :    AddBasicBlockRangesStreamPdbMutator mutator(bb_addresses);
  41    :  
  42  E :    bb_addresses.push_back(RelativeAddressRange(RelativeAddress(0x11111111), 4));
  43  E :    bb_addresses.push_back(RelativeAddressRange(RelativeAddress(0x22222222), 4));
  44    :  
  45  E :    pdb::PdbFile pdb_file;
  46  E :    ASSERT_NO_FATAL_FAILURE(testing::InitMockPdbFile(&pdb_file));
  47    :  
  48    :    // Add a dummy stream with the same name as the one we want to add.
  49  E :    pdb::PdbInfoHeader70 pdb_header = {};
  50  E :    pdb::NameStreamMap name_stream_map;
  51    :    EXPECT_TRUE(pdb::ReadHeaderInfoStream(pdb_file, &pdb_header,
  52  E :                                          &name_stream_map));
  53  E :    scoped_refptr<pdb::PdbStream> stream(new pdb::PdbByteStream);
  54  E :    size_t stream_id = pdb_file.AppendStream(stream.get());
  55  E :    name_stream_map[kBasicBlockRangesStreamName] = stream_id;
  56    :    EXPECT_TRUE(pdb::WriteHeaderInfoStream(pdb_header, name_stream_map,
  57  E :                                           &pdb_file));
  58    :  
  59  E :    EXPECT_FALSE(mutator.MutatePdb(&pdb_file));
  60  E :  }
  61    :  
  62  E :  TEST(AddBasicBlockRangesStreamPdbMutatorTest, DoesNotAddEmptyStream) {
  63  E :    RelativeAddressRangeVector bb_addresses;
  64  E :    AddBasicBlockRangesStreamPdbMutator mutator(bb_addresses);
  65    :  
  66  E :    pdb::PdbFile pdb_file;
  67  E :    ASSERT_NO_FATAL_FAILURE(testing::InitMockPdbFile(&pdb_file));
  68    :  
  69  E :    EXPECT_TRUE(mutator.MutatePdb(&pdb_file));
  70    :  
  71  E :    pdb::PdbInfoHeader70 pdb_header = {};
  72  E :    pdb::NameStreamMap name_stream_map;
  73    :    EXPECT_TRUE(pdb::ReadHeaderInfoStream(pdb_file, &pdb_header,
  74  E :                                          &name_stream_map));
  75    :  
  76    :    // We expect no named stream to have been added.
  77  E :    EXPECT_EQ(0u, name_stream_map.count(kBasicBlockRangesStreamName));
  78  E :  }
  79    :  
  80  E :  TEST(AddBasicBlockRangesStreamPdbMutatorTest, AddsStream) {
  81  E :    RelativeAddressRangeVector bb_addresses;
  82  E :    AddBasicBlockRangesStreamPdbMutator mutator(bb_addresses);
  83    :  
  84  E :    bb_addresses.push_back(RelativeAddressRange(RelativeAddress(0x11111111), 4));
  85  E :    bb_addresses.push_back(RelativeAddressRange(RelativeAddress(0x22222222), 4));
  86    :  
  87  E :    pdb::PdbFile pdb_file;
  88  E :    ASSERT_NO_FATAL_FAILURE(testing::InitMockPdbFile(&pdb_file));
  89    :  
  90  E :    EXPECT_TRUE(mutator.MutatePdb(&pdb_file));
  91    :  
  92  E :    pdb::PdbInfoHeader70 pdb_header = {};
  93  E :    pdb::NameStreamMap name_stream_map;
  94    :    EXPECT_TRUE(pdb::ReadHeaderInfoStream(pdb_file, &pdb_header,
  95  E :                                          &name_stream_map));
  96    :  
  97    :    // We expect the named stream to have been added.
  98  E :    EXPECT_EQ(1u, name_stream_map.count(kBasicBlockRangesStreamName));
  99    :  
 100    :    // Get the stream.
 101  E :    size_t stream_id = name_stream_map[kBasicBlockRangesStreamName];
 102  E :    scoped_refptr<pdb::PdbStream> stream = pdb_file.GetStream(stream_id);
 103  E :    EXPECT_TRUE(stream.get() != NULL);
 104    :  
 105    :    // Validate the stream contents.
 106  E :    RelativeAddressRangeVector bb_addresses2;
 107  E :    EXPECT_TRUE(stream->Seek(0));
 108  E :    EXPECT_TRUE(stream->Read(&bb_addresses2));
 109  E :    EXPECT_THAT(bb_addresses, testing::ContainerEq(bb_addresses2));
 110  E :  }
 111    :  
 112    :  }  // namespace mutators
 113    :  }  // namespace instrument