[Syzygy Instrumenter] Add FillerTransform.

syzygy/instrument/transforms/filler_transform_unittest.cc

Index: syzygy/instrument/transforms/filler_transform_unittest.cc
diff --git a/syzygy/instrument/transforms/filler_transform_unittest.cc b/syzygy/instrument/transforms/filler_transform_unittest.cc
new file mode 100644
index 0000000000000000000000000000000000000000..cebfa50e1a9451796a159b05a3cdce27bd64a383
--- /dev/null
+++ b/syzygy/instrument/transforms/filler_transform_unittest.cc
@@ -0,0 +1,302 @@
+// Copyright 2015 Google Inc. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "syzygy/instrument/transforms/filler_transform.h"
+
+#include <set>
+
+#include "base/memory/scoped_ptr.h"
+#include "gtest/gtest.h"
+#include "syzygy/assm/assembler_base.h"
+#include "syzygy/block_graph/basic_block.h"
+#include "syzygy/block_graph/basic_block_assembler.h"
+#include "syzygy/block_graph/basic_block_subgraph.h"
+#include "syzygy/block_graph/basic_block_test_util.h"
+#include "syzygy/block_graph/block_graph.h"
+#include "syzygy/instrument/transforms/unittest_util.h"
+
+namespace instrument {
+namespace transforms {
+namespace {
+
+using block_graph::BasicBlock;
+using block_graph::BlockGraph;
+
+const char kNopName[] = "NOP";
+
+// Finds all blocks whose names appear in @p target_list, and writes the mapping
+// fron name to block in @p result. Returns true iff all names in @p target_list
+// are found.
+bool FindAllBlocks(BlockGraph* block_graph,
+                   const std::vector<std::string>& target_list,
+                   std::map<std::string, BlockGraph::Block*>* result) {
+  DCHECK(result && result->empty());
+  std::set<std::string> name_set(target_list.begin(), target_list.end());
+  for (auto& it : block_graph->blocks_mutable()) {
+    std::string block_name = it.second.name();
+    if (name_set.find(block_name) != name_set.end()) {
+      (*result)[block_name] = &it.second;
+      name_set.erase(block_name);
+    }
+  }
+  return name_set.empty();
+}
+
+// Finds code offsets and sizes of every NOP in @p instructions and stores the
+// result in @p nop_offsets. Returns the number of @p instructions.
+size_t FindAllNops(BasicBlock::Instructions* instructions,
+                   std::map<size_t, int>* nop_offsets) {
+  // Find code offsets of every NOP that appears.
+  size_t offset = 0LL;
+  BasicBlock::Instructions::iterator inst_it = instructions->begin();
+  for (; inst_it != instructions->end(); ++inst_it, ++offset) {
+    if (!::strcmp(inst_it->GetName(), kNopName))
+      (*nop_offsets)[offset] = inst_it->size();
+  }
+  return offset;
+}
+
+class TestNopInjector : public NopInjector {
+ public:
+  explicit TestNopInjector(const NopSpec& nop_spec) : NopInjector(nop_spec) { }
+};
+
+class NopInjectorTest : public testing::Test {
+ public:
+  typedef testing::Test Super;
+
+  // Creates a new length @p n instruction list for testing.
+  void CreateInstructions(int n) {
+    using assm::eax;
+    using assm::ebp;
+    using assm::esp;
+    using assm::kSize32Bit;
+
+    instructions_.reset(new BasicBlock::Instructions);
+    if (n == 0)
+      return;
+
+    block_graph::BasicBlockAssembler
+        assm(instructions_->begin(), instructions_.get());
+    bool setup_stack = false;
+    --n;  // Reserve for ret.
+    if (n > 3) {
+      setup_stack = true;
+      n -= 3;  // Reserve for stack frame setup.
+    }
+    if (setup_stack) {
+      assm.push(ebp);
+      assm.mov(ebp, esp);
+    }
+    for (; n > 0; --n)
+      assm.mov(eax, block_graph::Immediate(n, kSize32Bit));
+    if (setup_stack)
+      assm.pop(ebp);
+    assm.ret(0);
+  }
+
+ protected:
+  scoped_ptr<BasicBlock::Instructions> instructions_;
+};
+
+class TestFillerTransform : public FillerTransform {
+ public:
+  explicit TestFillerTransform(const std::vector<std::string>& target_list)
+      : FillerTransform(target_list) { }
+};
+
+class FillerTransformTest : public testing::TestDllTransformTest {
+ public:
+  typedef testing::TestDllTransformTest Super;
+
+  // Checks that every basic code block in a block has NOP placed at expected
+  // code offsets.
+  void CheckBlockNops(BlockGraph::Block* block) {
+    using block_graph::BasicBlockDecomposer;
+    using block_graph::BasicBlockSubGraph;
+    using block_graph::BasicCodeBlock;
+
+    // Decompose target block to subgraph.
+    BasicBlockSubGraph subgraph;
+    BasicBlockDecomposer bb_decomposer(block, &subgraph);
+    ASSERT_TRUE(bb_decomposer.Decompose());
+
+    // Examine each code subblock.
+    BasicBlockSubGraph::BBCollection& basic_blocks =
+        subgraph.basic_blocks();
+    for (auto bb = basic_blocks.begin(); bb != basic_blocks.end(); ++bb) {
+      BasicCodeBlock* bc_block = BasicCodeBlock::Cast(*bb);
+      if (bc_block == nullptr)
+        continue;
+
+      std::map<size_t, int> nop_offsets;
+      size_t num_inst = FindAllNops(&bc_block->instructions(), &nop_offsets);
+      // The checks here depend on how NopInjector is initialized in
+      // FillerBasicBlockTransform.
+      EXPECT_TRUE(1 >= num_inst || nop_offsets[1] == 1);
+    }
+  }
+};
+
+}  // namespace
+
+// Sanity check for helper CreateInstructions().
+TEST_F(NopInjectorTest, CreateInstructions) {
+  for (int i = 0; i < 10; ++i) {
+    CreateInstructions(i);
+    size_t count = 0;
+    for (auto inst : *instructions_.get())
+      ++count;
+    EXPECT_EQ(i, count);
+  }
+}
+
+TEST_F(NopInjectorTest, InjectToStart) {
+  CreateInstructions(5);
+  TestNopInjector injector1({{0, NopInjector::NOP4}});
+  EXPECT_EQ(1U, injector1.Inject(instructions_.get()));
+  std::map<size_t, int> nop_offsets1;
+  EXPECT_EQ(6U, FindAllNops(instructions_.get(), &nop_offsets1));
+  EXPECT_EQ(1U, nop_offsets1.size());
+  EXPECT_EQ(4U, nop_offsets1[0]);
+
+  // Inject another NOP, when a NOP already exists.
+  TestNopInjector injector2({{0, NopInjector::NOP1}});
+  EXPECT_EQ(1U, injector2.Inject(instructions_.get()));
+  std::map<size_t, int> nop_offsets2;
+  EXPECT_EQ(7U, FindAllNops(instructions_.get(), &nop_offsets2));
+  EXPECT_EQ(2U, nop_offsets2.size());  // New + existing.
+  EXPECT_EQ(1U, nop_offsets2[0]);  // From |injector2|.
+  EXPECT_EQ(4U, nop_offsets2[1]);  // From |injector1|.
+}
+
+TEST_F(NopInjectorTest, InjectToBeforeEnd) {
+  CreateInstructions(7);
+  TestNopInjector injector({{6, NopInjector::NOP2}});
+  EXPECT_EQ(1U, injector.Inject(instructions_.get()));
+  std::map<size_t, int> nop_offsets;
+  EXPECT_EQ(8U, FindAllNops(instructions_.get(), &nop_offsets));
+  EXPECT_EQ(1U, nop_offsets.size());
+  EXPECT_EQ(2U, nop_offsets[6]);
+}
+
+TEST_F(NopInjectorTest, CannotInjectBeyondEnd) {
+  CreateInstructions(7);
+  TestNopInjector injector({{7, NopInjector::NOP1}, {17, NopInjector::NOP1}});
+  EXPECT_EQ(0U, injector.Inject(instructions_.get()));
+  std::map<size_t, int> nop_offsets;
+  EXPECT_EQ(7U, FindAllNops(instructions_.get(), &nop_offsets));
+  EXPECT_EQ(0U, nop_offsets.size());
+}
+
+TEST_F(NopInjectorTest, InjectToEmpty) {
+  CreateInstructions(0);
+  TestNopInjector injector({{0, NopInjector::NOP1}, {1, NopInjector::NOP2}});
+  EXPECT_EQ(0U, injector.Inject(instructions_.get()));
+  std::map<size_t, int> nop_offsets;
+  EXPECT_EQ(0U, FindAllNops(instructions_.get(), &nop_offsets));
+  EXPECT_EQ(0U, nop_offsets.size());
+}
+
+TEST_F(NopInjectorTest, InjectToSingle) {
+  CreateInstructions(1);
+  TestNopInjector injector({
+    {0, NopInjector::NOP5},
+    {1, NopInjector::NOP8},
+    {3, NopInjector::NOP2}});  // Gets ignored.
+  EXPECT_EQ(2U, injector.Inject(instructions_.get()));
+  std::map<size_t, int> nop_offsets;
+  EXPECT_EQ(3U, FindAllNops(instructions_.get(), &nop_offsets));
+  EXPECT_EQ(2U, nop_offsets.size());
+  EXPECT_EQ(5U, nop_offsets[0]);
+  EXPECT_EQ(8U, nop_offsets[1]);
+}
+
+TEST_F(NopInjectorTest, InjectNone) {
+  CreateInstructions(7);
+  TestNopInjector injector({});
+  EXPECT_EQ(0U, injector.Inject(instructions_.get()));
+  std::map<size_t, int> nop_offsets;
+  EXPECT_EQ(7U, FindAllNops(instructions_.get(), &nop_offsets));
+  EXPECT_EQ(0U, nop_offsets.size());
+}
+
+TEST_F(NopInjectorTest, InjectConsecutive) {
+  CreateInstructions(4);
+  TestNopInjector injector({
+      {0, NopInjector::NOP1},
+      {1, NopInjector::NOP2},
+      {2, NopInjector::NOP3},
+      {3, NopInjector::NOP5},
+      {4, NopInjector::NOP7},
+      {5, NopInjector::NOP11}});
+  EXPECT_EQ(6U, injector.Inject(instructions_.get()));
+  std::map<size_t, int> nop_offsets;
+  EXPECT_EQ(10U, FindAllNops(instructions_.get(), &nop_offsets));
+  EXPECT_EQ(6U, nop_offsets.size());
+  EXPECT_EQ(1U, nop_offsets[0]);
+  EXPECT_EQ(2U, nop_offsets[1]);
+  EXPECT_EQ(3U, nop_offsets[2]);
+  EXPECT_EQ(5U, nop_offsets[3]);
+  EXPECT_EQ(7U, nop_offsets[4]);
+  EXPECT_EQ(11U, nop_offsets[5]);
+}
+
+TEST_F(NopInjectorTest, InjectAlternate) {
+  CreateInstructions(4);
+  TestNopInjector injector({
+      {0, NopInjector::NOP10},
+      {2, NopInjector::NOP9},
+      {4, NopInjector::NOP8},
+      {6, NopInjector::NOP7},
+      {8, NopInjector::NOP6},  // Gets ignored.
+      {10, NopInjector::NOP5}});  // Gets ignored.
+  EXPECT_EQ(4U, injector.Inject(instructions_.get()));
+  std::map<size_t, int> nop_offsets;
+  EXPECT_EQ(8U, FindAllNops(instructions_.get(), &nop_offsets));
+  EXPECT_EQ(4U, nop_offsets.size());
+  EXPECT_EQ(10U, nop_offsets[0]);
+  EXPECT_EQ(9U, nop_offsets[2]);
+  EXPECT_EQ(8U, nop_offsets[4]);
+  EXPECT_EQ(7U, nop_offsets[6]);
+}
+
+TEST_F(FillerTransformTest, Apply) {
+  std::vector<std::string> target_list = {
+    "Used::M",
+    "TestUnusedFuncs"
+  };
+
+  ASSERT_NO_FATAL_FAILURE(DecomposeTestDll());
+
+  // Apply the transform.
+  TestFillerTransform tx(target_list);
+  ASSERT_TRUE(block_graph::ApplyBlockGraphTransform(
+      &tx, policy_, &block_graph_, header_block_));
+
+  // Check results. First find and store target blocks.
+  EXPECT_EQ(2U, tx.num_target_updated());
+  std::map<std::string, BlockGraph::Block*> target_map;
+  ASSERT_TRUE(FindAllBlocks(&block_graph_, target_list, &target_map));
+  ASSERT_TRUE(target_map.size() == target_list.size());
+
+  // Check the block for each target.
+  for (auto& it : target_map) {
+    BlockGraph::Block* target_block = it.second;
+    CheckBlockNops(target_block);
+  }
+}
+
+}  // namespace transforms
+}  // namespace instrument