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/pe/transforms/add_hot_patching_metadata_transform.h"
16 :
17 : #include <unordered_set>
18 :
19 : #include "gtest/gtest.h"
20 : #include "syzygy/block_graph/hot_patching_metadata.h"
21 : #include "syzygy/common/defs.h"
22 : #include "syzygy/instrument/transforms/unittest_util.h"
23 :
24 : namespace pe {
25 : namespace transforms {
26 :
27 : namespace {
28 :
29 : using block_graph::BlockGraph;
30 :
31 : // Expose the function needed for unittesting.
32 : class TestAddHotPatchingMetadataTransform
33 : : public AddHotPatchingMetadataTransform {
34 : public:
35 : using AddHotPatchingMetadataTransform::CalculateCodeSize;
36 : };
37 :
38 : // Inserts all blocks that are safe to decompose into a container.
39 : // @param block_container The blocks will be inserted into this container.
40 : void BuildListOfDecomposableBlocks(
41 : const block_graph::TransformPolicyInterface* policy,
42 : BlockGraph* block_graph,
43 E : AddHotPatchingMetadataTransform::BlockVector* block_container) {
44 : // Make sure that test.dll has been decomposed first.
45 E : ASSERT_NE(0U, block_graph->blocks().size());
46 :
47 : // Collect decomposable code blocks to test_blocks_ vector.
48 E : for (auto &entry : block_graph->blocks_mutable()) {
49 E : BlockGraph::Block* block = &entry.second;
50 :
51 E : if (policy->BlockIsSafeToBasicBlockDecompose(block))
52 E : block_container->push_back(block);
53 E : }
54 :
55 : // Check that there are decomposable code blocks.
56 E : EXPECT_NE(0U, block_container->size());
57 E : }
58 :
59 : // A test fixture which knows how to decompose the "standard" test dll.
60 : class AddHotPatchingMetadataTransformTest : public testing::PELibUnitTest {
61 : public:
62 E : AddHotPatchingMetadataTransformTest() : layout_(&block_graph_) {}
63 :
64 : // The block graph for test_dll.dll.
65 : BlockGraph block_graph_;
66 :
67 : // The layout of test_dll.dll.
68 : pe::ImageLayout layout_;
69 :
70 : // The policy objects restricting how the transform is applied.
71 : pe::PETransformPolicy pe_policy_;
72 :
73 : // The PEFile instance referring to test_dll.
74 : pe::PEFile pe_file_;
75 : };
76 :
77 : } // namespace
78 :
79 E : TEST(AddHotPatchingMetadataTransformSimpleTest, SetBlocksPrepared) {
80 : // Create an empty container.
81 E : AddHotPatchingMetadataTransform::BlockVector cont;
82 :
83 : // Test set_blocks_prepared.
84 E : AddHotPatchingMetadataTransform hpt;
85 E : EXPECT_EQ(nullptr, hpt.blocks_prepared());
86 E : hpt.set_blocks_prepared(&cont);
87 E : EXPECT_EQ(&cont, hpt.blocks_prepared());
88 E : }
89 :
90 E : TEST(AddHotPatchingMetadataTransformSimpleTest, CalculateCodeSize) {
91 E : BlockGraph block_graph;
92 : BlockGraph::Block* block = block_graph.AddBlock(BlockGraph::CODE_BLOCK,
93 : 100U,
94 E : "dummy");
95 :
96 : // Add some data to the block.
97 : static const uint8 buffer[50];
98 E : block->SetData(buffer, sizeof(buffer));
99 :
100 : // The whole data should be considered as code if there are no labels.
101 E : ASSERT_EQ(50U, TestAddHotPatchingMetadataTransform::CalculateCodeSize(block));
102 :
103 : // A code label should not change the code size.
104 E : block->SetLabel(20U, "CODE", BlockGraph::CODE_LABEL);
105 E : ASSERT_EQ(50U, TestAddHotPatchingMetadataTransform::CalculateCodeSize(block));
106 :
107 : // A data label should limit the code size.
108 E : block->SetLabel(30U, "DATA", BlockGraph::DATA_LABEL);
109 E : ASSERT_EQ(30U, TestAddHotPatchingMetadataTransform::CalculateCodeSize(block));
110 :
111 : // A data label with other attributes set should limit the code size.
112 : block->SetLabel(29U,
113 : "DATA",
114 E : BlockGraph::DATA_LABEL | BlockGraph::JUMP_TABLE_LABEL);
115 E : ASSERT_EQ(29U, TestAddHotPatchingMetadataTransform::CalculateCodeSize(block));
116 :
117 : // A debug-end label at the end should be ignored.
118 E : block->SetLabel(49U, "DEBUG-END", BlockGraph::DEBUG_END_LABEL);
119 E : ASSERT_EQ(29U, TestAddHotPatchingMetadataTransform::CalculateCodeSize(block));
120 E : }
121 :
122 E : TEST_F(AddHotPatchingMetadataTransformTest, TransformBlockGraph) {
123 E : DecomposeTestDll(&pe_file_, &layout_);
124 :
125 : BlockGraph::Block* header_block = layout_.blocks.GetBlockByAddress(
126 E : core::RelativeAddress(0));
127 E : ASSERT_NE(nullptr, header_block);
128 :
129 : // Save section and block map size for comparison later.
130 E : size_t sections_before_size = block_graph_.sections().size();
131 E : size_t blocks_before_size = block_graph_.blocks().size();
132 :
133 : // Create the transform.
134 E : AddHotPatchingMetadataTransform hpt;
135 :
136 : // Initialize a blocks_prepared with some sample blocks
137 E : AddHotPatchingMetadataTransform::BlockVector cont;
138 E : BuildListOfDecomposableBlocks(&pe_policy_, &block_graph_, &cont);
139 E : hpt.set_blocks_prepared(&cont);
140 :
141 : // Add hot patching section.
142 E : hpt.TransformBlockGraph(&pe_policy_, &block_graph_, header_block);
143 :
144 : // Check that one new section and one new block is created.
145 E : ASSERT_EQ(sections_before_size + 1, block_graph_.sections().size());
146 E : ASSERT_EQ(blocks_before_size + 1, block_graph_.blocks().size());
147 :
148 : // Retrieve the new section.
149 : BlockGraph::Section* hp_metadata_section =
150 E : block_graph_.FindSection(common::kHotPatchingMetadataSectionName);
151 E : ASSERT_NE(nullptr, hp_metadata_section);
152 :
153 : // Retrieve the new block.
154 E : BlockGraph::Block* hp_metadata_block = nullptr;
155 E : for (auto &item: block_graph_.blocks_mutable()) {
156 E : if (item.second.name() == common::kHotPatchingMetadataSectionName) {
157 : // This is the new block.
158 E : hp_metadata_block = &item.second;
159 E : break;
160 : }
161 E : }
162 E : ASSERT_NE(nullptr, hp_metadata_block);
163 :
164 : // Check hot patching metadata header.
165 : const block_graph::HotPatchingMetadataHeader* hp_metadata_header =
166 : reinterpret_cast<const block_graph::HotPatchingMetadataHeader*>(
167 E : hp_metadata_block->data());
168 E : ASSERT_NE(nullptr, hp_metadata_header);
169 : EXPECT_EQ(block_graph::kHotPatchingMetadataVersion,
170 E : hp_metadata_header->version);
171 : EXPECT_EQ(hpt.blocks_prepared()->size(),
172 E : hp_metadata_header->number_of_blocks);
173 :
174 : // Locate the block metadata array.
175 : // The (hp_metadata_header + 1) expression is a pointer pointing to the
176 : // location after the header.
177 : const block_graph::HotPatchingBlockMetadata* hp_block_metadata_arr =
178 : reinterpret_cast<const block_graph::HotPatchingBlockMetadata*>(
179 E : hp_metadata_header + 1);
180 :
181 : // The new block should have a reference to each of the test blocks. This test
182 : // uses the assumption that the references to the blocks will be in the same
183 : // order as the blocks in blocks_prepared_.
184 : ASSERT_EQ(hpt.blocks_prepared()->size(),
185 E : hp_metadata_block->references().size());
186 E : int i = 0;
187 E : for (const auto& ref_entry : hp_metadata_block->references()) {
188 E : BlockGraph::Offset ref_offset = ref_entry.first;
189 E : const BlockGraph::Reference& ref = ref_entry.second;
190 :
191 : // Check reference offset.
192 E : EXPECT_EQ(
193 : reinterpret_cast<const uint8*>(
194 : &hp_block_metadata_arr[i].relative_address) -
195 : reinterpret_cast<const uint8*>(hp_metadata_header),
196 : ref_offset);
197 :
198 : // Check reference.
199 E : EXPECT_EQ(0, ref.base());
200 E : EXPECT_EQ(0, ref.offset());
201 E : EXPECT_EQ(hpt.blocks_prepared()->operator[](i), ref.referenced());
202 E : EXPECT_EQ(4U, ref.size());
203 E : EXPECT_EQ(BlockGraph::RELATIVE_REF, ref.type());
204 :
205 : // Check if the code and data size information is correct.
206 : const BlockGraph::Block* original_block =
207 E : hpt.blocks_prepared()->operator[](i);
208 E : EXPECT_EQ(original_block->data_size(),
209 : hp_block_metadata_arr[i].block_size);
210 E : EXPECT_EQ(TestAddHotPatchingMetadataTransform::CalculateCodeSize(
211 : original_block),
212 : hp_block_metadata_arr[i].code_size);
213 :
214 E : ++i;
215 E : }
216 E : }
217 :
218 E : TEST_F(AddHotPatchingMetadataTransformTest, TransformBlockGraphEmpty) {
219 E : DecomposeTestDll(&pe_file_, &layout_);
220 :
221 : BlockGraph::Block* header_block = layout_.blocks.GetBlockByAddress(
222 E : core::RelativeAddress(0));
223 E : ASSERT_NE(nullptr, header_block);
224 :
225 : // Save section and block map size for comparison later.
226 E : size_t sections_before_size = block_graph_.sections().size();
227 E : size_t blocks_before_size = block_graph_.blocks().size();
228 :
229 : // Create the transform.
230 E : AddHotPatchingMetadataTransform hpt;
231 :
232 : // Initialize a blocks_prepared with empty container
233 E : AddHotPatchingMetadataTransform::BlockVector cont;
234 E : hpt.set_blocks_prepared(&cont);
235 :
236 : // Add hot patching section with some sample blocks.
237 E : hpt.TransformBlockGraph(&pe_policy_, &block_graph_, header_block);
238 :
239 : // Check that no sections or blocks have been added.
240 E : EXPECT_EQ(sections_before_size, block_graph_.sections().size());
241 E : EXPECT_EQ(blocks_before_size, block_graph_.blocks().size());
242 E : }
243 :
244 : } // namespace transforms
245 : } // namespace pe
|