1 : // Copyright 2013 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 : // Jump table case count instrumentation transform unit-tests.
16 :
17 : #include "syzygy/instrument/transforms/jump_table_count_transform.h"
18 :
19 : #include "gtest/gtest.h"
20 : #include "syzygy/block_graph/basic_block.h"
21 : #include "syzygy/block_graph/basic_block_decomposer.h"
22 : #include "syzygy/block_graph/basic_block_subgraph.h"
23 : #include "syzygy/block_graph/block_graph.h"
24 : #include "syzygy/block_graph/block_util.h"
25 : #include "syzygy/block_graph/typed_block.h"
26 : #include "syzygy/common/indexed_frequency_data.h"
27 : #include "syzygy/instrument/transforms/unittest_util.h"
28 :
29 : #include "mnemonics.h" // NOLINT
30 :
31 : namespace instrument {
32 : namespace transforms {
33 :
34 : namespace {
35 :
36 : using block_graph::BasicBlock;
37 : using block_graph::BasicCodeBlock;
38 : using block_graph::BasicBlockDecomposer;
39 : using block_graph::BasicBlockSubGraph;
40 : using block_graph::BlockGraph;
41 : using block_graph::Instruction;
42 : using common::IndexedFrequencyData;
43 :
44 : class TestJumpTableCaseCountTransform : public JumpTableCaseCountTransform {
45 : public:
46 : using JumpTableCaseCountTransform::add_frequency_data;
47 : using JumpTableCaseCountTransform::jump_table_case_counter_hook_ref;
48 : using JumpTableCaseCountTransform::thunk_section;
49 :
50 E : BlockGraph::Block* frequency_data_block() {
51 E : return add_frequency_data()->frequency_data_block();
52 E : }
53 :
54 E : BlockGraph::Block* frequency_data_buffer_block() {
55 E : return add_frequency_data()->frequency_data_buffer_block();
56 E : }
57 : };
58 :
59 : typedef testing::TestDllTransformTest JumpTableCaseCountTransformTest;
60 :
61 : // Ensures that the @p block is a jump table case count thunk.
62 E : void CheckBlockIsAThunk(BlockGraph::Block* block) {
63 : // Decompose the block to basic-blocks.
64 E : BasicBlockSubGraph subgraph;
65 E : BasicBlockDecomposer bb_decomposer(block, &subgraph);
66 E : ASSERT_TRUE(bb_decomposer.Decompose());
67 :
68 E : ASSERT_EQ(1, subgraph.basic_blocks().size());
69 : const BasicCodeBlock* bb = BasicCodeBlock::Cast(
70 E : *subgraph.basic_blocks().begin());
71 E : ASSERT_TRUE(bb != NULL);
72 E : ASSERT_FALSE(bb->is_padding());
73 :
74 E : ASSERT_EQ(2U, bb->instructions().size());
75 : BasicBlock::Instructions::const_iterator inst_iter =
76 E : bb->instructions().begin();
77 :
78 : // Instruction 1 should push the case id.
79 E : EXPECT_EQ(I_PUSH, inst_iter->representation().opcode);
80 :
81 : // Instruction 2 should call the jump table counter hook.
82 E : ++inst_iter;
83 E : EXPECT_EQ(I_CALL, inst_iter->representation().opcode);
84 E : }
85 :
86 : } // namespace
87 :
88 E : TEST_F(JumpTableCaseCountTransformTest, Apply) {
89 E : ASSERT_NO_FATAL_FAILURE(DecomposeTestDll());
90 :
91 : // Apply the transform.
92 E : TestJumpTableCaseCountTransform tx;
93 : ASSERT_TRUE(block_graph::ApplyBlockGraphTransform(
94 E : &tx, &policy_, &block_graph_, dos_header_block_));
95 E : ASSERT_TRUE(tx.frequency_data_block() != NULL);
96 E : ASSERT_TRUE(tx.thunk_section() != NULL);
97 E : ASSERT_TRUE(tx.jump_table_case_counter_hook_ref() != NULL);
98 E : ASSERT_TRUE(tx.jump_table_case_counter_hook_ref()->IsValid());
99 :
100 : // Validate the jump table frequency data structure.
101 E : block_graph::ConstTypedBlock<IndexedFrequencyData> frequency_data;
102 E : ASSERT_TRUE(frequency_data.Init(0, tx.frequency_data_block()));
103 E : EXPECT_EQ(sizeof(uint32), frequency_data->frequency_size);
104 E : EXPECT_EQ(common::kJumpTableCountAgentId, frequency_data->agent_id);
105 E : EXPECT_EQ(common::kJumpTableFrequencyDataVersion, frequency_data->version);
106 E : EXPECT_EQ(IndexedFrequencyData::JUMP_TABLE, frequency_data->data_type);
107 E : EXPECT_EQ(sizeof(IndexedFrequencyData), tx.frequency_data_block()->size());
108 : EXPECT_EQ(sizeof(IndexedFrequencyData),
109 E : tx.frequency_data_block()->data_size());
110 : EXPECT_TRUE(frequency_data.HasReferenceAt(
111 E : frequency_data.OffsetOf(frequency_data->frequency_data)));
112 E : BlockGraph::Reference frequency_data_reference;
113 : ASSERT_TRUE(frequency_data.block()->GetReference(
114 : frequency_data.OffsetOf(frequency_data->frequency_data),
115 E : &frequency_data_reference));
116 : EXPECT_EQ(frequency_data_reference.referenced(),
117 E : tx.frequency_data_buffer_block());
118 : EXPECT_EQ(frequency_data->num_entries * frequency_data->frequency_size,
119 E : tx.frequency_data_buffer_block()->size());
120 :
121 : // Examine each eligible block to verify that all the jump tables have
122 : // been instrumented.
123 E : size_t jump_table_entries = 0;
124 : for (BlockGraph::BlockMap::const_iterator block_iter(
125 E : block_graph_.blocks().begin());
126 E : block_iter != block_graph_.blocks().end();
127 E : ++block_iter) {
128 E : const BlockGraph::Block& block = block_iter->second;
129 :
130 : // Skip non-code blocks.
131 E : if (block.type() != BlockGraph::CODE_BLOCK)
132 E : continue;
133 :
134 : // We don't want to check the thunk blocks.
135 E : if (block.section() == tx.thunk_section()->id())
136 E : continue;
137 :
138 : // Iterate over the labels to find the jump tables.
139 : for (BlockGraph::Block::LabelMap::const_iterator iter_label(
140 E : block.labels().begin());
141 E : iter_label != block.labels().end();
142 E : ++iter_label) {
143 E : if (!iter_label->second.has_attributes(BlockGraph::JUMP_TABLE_LABEL))
144 E : continue;
145 :
146 E : size_t table_size = 0;
147 : ASSERT_TRUE(
148 E : block_graph::GetJumpTableSize(&block, iter_label, &table_size));
149 :
150 : BlockGraph::Block::ReferenceMap::const_iterator iter_ref =
151 E : block.references().find(iter_label->first);
152 E : ASSERT_TRUE(iter_ref != block.references().end());
153 :
154 : // Iterate over the references and ensure that they are thunked.
155 E : for (size_t i = 0; i < table_size; ++i) {
156 E : BlockGraph::Block* ref_block = iter_ref->second.referenced();
157 E : CheckBlockIsAThunk(ref_block);
158 E : ++iter_ref;
159 E : }
160 :
161 E : jump_table_entries += table_size;
162 E : }
163 E : }
164 E : DCHECK_EQ(frequency_data->num_entries, jump_table_entries);
165 E : }
166 :
167 : } // namespace transforms
168 : } // namespace instrument
|