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 : #include "syzygy/pe/transforms/coff_rename_symbols_transform.h"
16 :
17 : #include "gmock/gmock.h"
18 : #include "gtest/gtest.h"
19 : #include "syzygy/block_graph/block_hash.h"
20 : #include "syzygy/block_graph/unittest_util.h"
21 : #include "syzygy/core/unittest_util.h"
22 : #include "syzygy/pe/coff_decomposer.h"
23 : #include "syzygy/pe/coff_utils.h"
24 : #include "syzygy/pe/pe_utils.h"
25 : #include "syzygy/pe/unittest_util.h"
26 :
27 : namespace pe {
28 : namespace transforms {
29 :
30 : namespace {
31 :
32 : using block_graph::BlockGraph;
33 : using core::RelativeAddress;
34 :
35 : class TestCoffRenameSymbolsTransform : public CoffRenameSymbolsTransform {
36 : public:
37 : typedef CoffRenameSymbolsTransform::SymbolMap SymbolMap;
38 :
39 : using CoffRenameSymbolsTransform::mappings_;
40 : };
41 :
42 : class CoffRenameSymbolsTransformTest : public testing::CoffUnitTest {
43 : public:
44 E : virtual void SetUp() override {
45 E : testing::CoffUnitTest::SetUp();
46 E : ASSERT_NO_FATAL_FAILURE(DecomposeOriginal());
47 E : }
48 : };
49 :
50 : const char kFunction1Name[] = "__imp_?function1@@YAHXZ";
51 : const char kFunction2Name[] = "?function2@@YAHXZ";
52 : const char kFunction3Name[] = "?function3@@YAHXZ";
53 : const char kFunction4Name[] = "?function4@@YAHXZ"; // Does not exist.
54 : const char kFoo[] = "foo"; // Does not exist, less than 8 characters.
55 : const char kMemset[] = "_memset"; // Exists, and is multiply defined.
56 :
57 : } // namespace
58 :
59 : std::ostream& operator<<(std::ostream& os,
60 : const block_graph::BlockHash& hash) {
61 : os << "BlockHash(" << hash.md5_digest.a << ")";
62 : return os;
63 : }
64 :
65 E : TEST_F(CoffRenameSymbolsTransformTest, ApplyMissingSymbolFails) {
66 E : TestCoffRenameSymbolsTransform tx;
67 E : EXPECT_TRUE(tx.mappings_.empty());
68 :
69 E : TestCoffRenameSymbolsTransform::SymbolMap expected_mappings;
70 E : expected_mappings.push_back(std::make_pair(kFunction4Name, kFunction1Name));
71 :
72 E : tx.AddSymbolMapping(kFunction4Name, kFunction1Name);
73 E : EXPECT_THAT(tx.mappings_, testing::ContainerEq(expected_mappings));
74 :
75 : BlockGraph::Block* symbols_block;
76 : BlockGraph::Block* strings_block;
77 : ASSERT_TRUE(FindCoffSpecialBlocks(
78 E : &block_graph_, NULL, &symbols_block, &strings_block));
79 E : block_graph::BlockHash symbols_hash_before(symbols_block);
80 E : block_graph::BlockHash strings_hash_before(strings_block);
81 :
82 E : EXPECT_TRUE(tx.symbols_must_exist());
83 E : EXPECT_FALSE(tx.TransformBlockGraph(&policy_, &block_graph_, headers_block_));
84 :
85 : // The block contents should not have changed.
86 E : block_graph::BlockHash symbols_hash_after(symbols_block);
87 E : block_graph::BlockHash strings_hash_after(strings_block);
88 E : EXPECT_EQ(symbols_hash_before, symbols_hash_after);
89 E : EXPECT_EQ(strings_hash_before, strings_hash_after);
90 :
91 E : ASSERT_NO_FATAL_FAILURE(TestRoundTrip());
92 E : }
93 :
94 E : TEST_F(CoffRenameSymbolsTransformTest, ApplyMissingSymbolSucceeds) {
95 E : TestCoffRenameSymbolsTransform tx;
96 E : EXPECT_TRUE(tx.mappings_.empty());
97 :
98 E : TestCoffRenameSymbolsTransform::SymbolMap expected_mappings;
99 E : expected_mappings.push_back(std::make_pair(kFunction4Name, kFunction1Name));
100 :
101 E : tx.AddSymbolMapping(kFunction4Name, kFunction1Name);
102 E : EXPECT_THAT(tx.mappings_, testing::ContainerEq(expected_mappings));
103 :
104 : BlockGraph::Block* symbols_block;
105 : BlockGraph::Block* strings_block;
106 : ASSERT_TRUE(FindCoffSpecialBlocks(
107 E : &block_graph_, NULL, &symbols_block, &strings_block));
108 E : block_graph::BlockHash symbols_hash_before(symbols_block);
109 E : block_graph::BlockHash strings_hash_before(strings_block);
110 :
111 E : tx.set_symbols_must_exist(false);
112 E : EXPECT_FALSE(tx.symbols_must_exist());
113 E : EXPECT_TRUE(tx.TransformBlockGraph(&policy_, &block_graph_, headers_block_));
114 :
115 : // The block contents should not have changed.
116 E : block_graph::BlockHash symbols_hash_after(symbols_block);
117 E : block_graph::BlockHash strings_hash_after(strings_block);
118 E : EXPECT_EQ(symbols_hash_before, symbols_hash_after);
119 E : EXPECT_EQ(strings_hash_before, strings_hash_after);
120 :
121 E : ASSERT_NO_FATAL_FAILURE(TestRoundTrip());
122 E : }
123 :
124 E : TEST_F(CoffRenameSymbolsTransformTest, ApplyExistingSymbols) {
125 E : TestCoffRenameSymbolsTransform tx;
126 E : EXPECT_TRUE(tx.mappings_.empty());
127 :
128 E : TestCoffRenameSymbolsTransform::SymbolMap expected_mappings;
129 E : expected_mappings.push_back(std::make_pair(kFunction1Name, kFunction2Name));
130 :
131 E : tx.AddSymbolMapping(kFunction1Name, kFunction2Name);
132 E : EXPECT_THAT(tx.mappings_, testing::ContainerEq(expected_mappings));
133 :
134 : BlockGraph::Block* symbols_block;
135 : BlockGraph::Block* strings_block;
136 : ASSERT_TRUE(FindCoffSpecialBlocks(
137 E : &block_graph_, NULL, &symbols_block, &strings_block));
138 E : size_t symbols_before = symbols_block->size();
139 E : size_t strings_before = strings_block->size();
140 :
141 E : EXPECT_TRUE(tx.symbols_must_exist());
142 E : EXPECT_TRUE(tx.TransformBlockGraph(&policy_, &block_graph_, headers_block_));
143 :
144 E : EXPECT_EQ(symbols_before, symbols_block->size());
145 E : EXPECT_EQ(strings_before, strings_block->size());
146 :
147 E : ASSERT_NO_FATAL_FAILURE(TestRoundTrip());
148 E : }
149 :
150 E : TEST_F(CoffRenameSymbolsTransformTest, ApplyMultiplyDefinedSource) {
151 E : TestCoffRenameSymbolsTransform tx;
152 E : EXPECT_TRUE(tx.mappings_.empty());
153 :
154 E : TestCoffRenameSymbolsTransform::SymbolMap expected_mappings;
155 E : expected_mappings.push_back(std::make_pair(kMemset, kFunction2Name));
156 :
157 E : tx.AddSymbolMapping(kMemset, kFunction2Name);
158 E : EXPECT_THAT(tx.mappings_, testing::ContainerEq(expected_mappings));
159 :
160 : BlockGraph::Block* symbols_block;
161 : BlockGraph::Block* strings_block;
162 : ASSERT_TRUE(FindCoffSpecialBlocks(
163 E : &block_graph_, NULL, &symbols_block, &strings_block));
164 E : size_t symbols_before = symbols_block->size();
165 E : size_t strings_before = strings_block->size();
166 :
167 E : EXPECT_TRUE(tx.symbols_must_exist());
168 E : EXPECT_TRUE(tx.TransformBlockGraph(&policy_, &block_graph_, headers_block_));
169 :
170 E : EXPECT_EQ(symbols_before, symbols_block->size());
171 E : EXPECT_EQ(strings_before, strings_block->size());
172 :
173 E : ASSERT_NO_FATAL_FAILURE(TestRoundTrip());
174 E : }
175 :
176 E : TEST_F(CoffRenameSymbolsTransformTest, ApplyMultiplyDefinedDestination) {
177 E : TestCoffRenameSymbolsTransform tx;
178 E : EXPECT_TRUE(tx.mappings_.empty());
179 :
180 E : TestCoffRenameSymbolsTransform::SymbolMap expected_mappings;
181 E : expected_mappings.push_back(std::make_pair(kFunction2Name, kMemset));
182 :
183 E : tx.AddSymbolMapping(kFunction2Name, kMemset);
184 E : EXPECT_THAT(tx.mappings_, testing::ContainerEq(expected_mappings));
185 :
186 : BlockGraph::Block* symbols_block;
187 : BlockGraph::Block* strings_block;
188 : ASSERT_TRUE(FindCoffSpecialBlocks(
189 E : &block_graph_, NULL, &symbols_block, &strings_block));
190 E : size_t symbols_before = symbols_block->size();
191 E : size_t strings_before = strings_block->size();
192 :
193 E : EXPECT_TRUE(tx.symbols_must_exist());
194 E : EXPECT_TRUE(tx.TransformBlockGraph(&policy_, &block_graph_, headers_block_));
195 :
196 E : EXPECT_EQ(symbols_before, symbols_block->size());
197 E : EXPECT_EQ(strings_before, strings_block->size());
198 :
199 E : ASSERT_NO_FATAL_FAILURE(TestRoundTrip());
200 E : }
201 :
202 E : TEST_F(CoffRenameSymbolsTransformTest, ApplyNewSymbolLong) {
203 E : TestCoffRenameSymbolsTransform tx;
204 E : EXPECT_TRUE(tx.mappings_.empty());
205 :
206 E : TestCoffRenameSymbolsTransform::SymbolMap expected_mappings;
207 E : expected_mappings.push_back(std::make_pair(kFunction3Name, kFunction4Name));
208 :
209 E : tx.AddSymbolMapping(kFunction3Name, kFunction4Name);
210 E : EXPECT_THAT(tx.mappings_, testing::ContainerEq(expected_mappings));
211 :
212 : BlockGraph::Block* symbols_block;
213 : BlockGraph::Block* strings_block;
214 : ASSERT_TRUE(FindCoffSpecialBlocks(
215 E : &block_graph_, NULL, &symbols_block, &strings_block));
216 E : size_t symbols_before = symbols_block->size();
217 E : size_t strings_before = strings_block->size();
218 :
219 E : EXPECT_TRUE(tx.symbols_must_exist());
220 E : EXPECT_TRUE(tx.TransformBlockGraph(&policy_, &block_graph_, headers_block_));
221 :
222 : // Expect only one symbol to have been created, and a corresponding string.
223 E : EXPECT_EQ(symbols_before + sizeof(IMAGE_SYMBOL), symbols_block->size());
224 : EXPECT_EQ(strings_before + ::strlen(kFunction4Name) + 1,
225 E : strings_block->size());
226 :
227 E : ASSERT_NO_FATAL_FAILURE(TestRoundTrip());
228 E : }
229 :
230 E : TEST_F(CoffRenameSymbolsTransformTest, ApplyNewSymbolShort) {
231 E : TestCoffRenameSymbolsTransform tx;
232 E : EXPECT_TRUE(tx.mappings_.empty());
233 :
234 E : TestCoffRenameSymbolsTransform::SymbolMap expected_mappings;
235 E : expected_mappings.push_back(std::make_pair(kFunction1Name, kFoo));
236 :
237 E : tx.AddSymbolMapping(kFunction1Name, kFoo);
238 E : EXPECT_THAT(tx.mappings_, testing::ContainerEq(expected_mappings));
239 :
240 : BlockGraph::Block* symbols_block;
241 : BlockGraph::Block* strings_block;
242 : ASSERT_TRUE(FindCoffSpecialBlocks(
243 E : &block_graph_, NULL, &symbols_block, &strings_block));
244 E : size_t symbols_before = symbols_block->size();
245 E : size_t strings_before = strings_block->size();
246 :
247 E : EXPECT_TRUE(tx.symbols_must_exist());
248 E : EXPECT_TRUE(tx.TransformBlockGraph(&policy_, &block_graph_, headers_block_));
249 :
250 : // Expect only one symbol to have been created, but no new string.
251 E : EXPECT_EQ(symbols_before + sizeof(IMAGE_SYMBOL), symbols_block->size());
252 E : EXPECT_EQ(strings_before, strings_block->size());
253 :
254 E : ASSERT_NO_FATAL_FAILURE(TestRoundTrip());
255 E : }
256 :
257 : } // namespace transforms
258 : } // namespace pe
|