1 : // Copyright 2012 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 : // Unittests for ThunkImportReferencesTransform.
16 :
17 : #include "syzygy/instrument/transforms/thunk_import_references_transform.h"
18 :
19 : #include <vector>
20 :
21 : #include "base/string_util.h"
22 : #include "base/files/scoped_temp_dir.h"
23 : #include "gmock/gmock.h"
24 : #include "gtest/gtest.h"
25 : #include "syzygy/block_graph/typed_block.h"
26 : #include "syzygy/common/defs.h"
27 : #include "syzygy/core/unittest_util.h"
28 : #include "syzygy/instrument/transforms/unittest_util.h"
29 : #include "syzygy/pe/decomposer.h"
30 : #include "syzygy/pe/transforms/add_imports_transform.h"
31 :
32 : namespace instrument {
33 : namespace transforms {
34 :
35 : namespace {
36 :
37 : using block_graph::BlockGraph;
38 : using block_graph::TypedBlock;
39 : using core::AbsoluteAddress;
40 : using pe::transforms::AddImportsTransform;
41 :
42 : typedef BlockGraph::Block::ReferrerSet ReferrerSet;
43 :
44 : // Expose protected members for testing.
45 : class TestThunkImportReferencesTransform
46 : : public ThunkImportReferencesTransform {
47 : public:
48 : // Make data members public for testing.
49 : using ThunkImportReferencesTransform::add_imports_transform;
50 : using ThunkImportReferencesTransform::modules_to_exclude_;
51 :
52 : // Make member types public for testing.
53 : using ThunkImportReferencesTransform::BlockSet;
54 : using ThunkImportReferencesTransform::ImportAddressLocation;
55 : using ThunkImportReferencesTransform::ImportAddressLocationNameMap;
56 : using ThunkImportReferencesTransform::ModuleNameSet;
57 :
58 : // Make member functions public for testing.
59 : using ThunkImportReferencesTransform::GetImportBlocks;
60 : using ThunkImportReferencesTransform::LookupImportLocations;
61 : using ThunkImportReferencesTransform::LookupDelayImportLocations;
62 : };
63 :
64 : class ThunkImportReferencesTransformTest
65 : : public testing::TestDllTransformTest {
66 : public:
67 : typedef TestThunkImportReferencesTransform::BlockSet BlockSet;
68 : typedef TestThunkImportReferencesTransform::ModuleNameSet ModuleNameSet;
69 : typedef TestThunkImportReferencesTransform::ImportAddressLocation
70 : ImportAddressLocation;
71 : typedef TestThunkImportReferencesTransform::ImportAddressLocationNameMap
72 : ImportAddressLocationNameMap;
73 :
74 E : ThunkImportReferencesTransformTest() : input_image_layout_(&block_graph_) {
75 E : }
76 :
77 E : virtual void SetUp() {
78 E : ASSERT_NO_FATAL_FAILURE(DecomposeTestDll());
79 E : }
80 :
81 : protected:
82 : pe::ImageLayout input_image_layout_;
83 : };
84 :
85 : } // namespace
86 :
87 E : TEST_F(ThunkImportReferencesTransformTest, Initialization) {
88 E : TestThunkImportReferencesTransform transform;
89 :
90 : EXPECT_STREQ(TestThunkImportReferencesTransform::kDefaultInstrumentDll,
91 E : transform.instrument_dll_name());
92 E : EXPECT_TRUE(transform.modules_to_exclude_.empty());
93 E : }
94 :
95 E : TEST_F(ThunkImportReferencesTransformTest, LookupImportLocations) {
96 E : ImportAddressLocationNameMap all_import_locations;
97 : // This should enumerate all imports.
98 : ASSERT_TRUE(
99 : TestThunkImportReferencesTransform::LookupImportLocations(
100 E : ModuleNameSet(), dos_header_block_, &all_import_locations));
101 :
102 : // Check that we found all the functions imported from export_dll.dll.
103 E : std::set<std::string> export_dll_imports;
104 : ImportAddressLocationNameMap::const_iterator it =
105 E : all_import_locations.begin();
106 : static const char kExportDLLPrefix[] = "export_dll.dll:";
107 :
108 E : for (; it != all_import_locations.end(); ++it) {
109 E : if (StartsWithASCII(it->second, kExportDLLPrefix, false))
110 E : export_dll_imports.insert(it->second);
111 E : }
112 :
113 : EXPECT_THAT(export_dll_imports,
114 : testing::ElementsAre("export_dll.dll:#7",
115 : "export_dll.dll:function1",
116 : "export_dll.dll:function3",
117 E : "export_dll.dll:kExportedData"));
118 E : }
119 :
120 E : TEST_F(ThunkImportReferencesTransformTest, LookupDelayImportLocations) {
121 E : ImportAddressLocationNameMap all_import_locations;
122 : // This should enumerate all delay imports.
123 : ASSERT_TRUE(
124 : TestThunkImportReferencesTransform::LookupDelayImportLocations(
125 E : ModuleNameSet(), dos_header_block_, &all_import_locations));
126 :
127 : // There should be precisely one ole32.dll import.
128 E : ASSERT_EQ(1, all_import_locations.size());
129 :
130 : ASSERT_STRCASEEQ("ole32.dll:CoCreateGuid",
131 E : all_import_locations.begin()->second.c_str());
132 E : }
133 :
134 E : TEST_F(ThunkImportReferencesTransformTest, ExcludeModule) {
135 E : TestThunkImportReferencesTransform transform;
136 :
137 E : EXPECT_TRUE(transform.modules_to_exclude_.empty());
138 E : transform.ExcludeModule("kernel32.dll");
139 : EXPECT_THAT(transform.modules_to_exclude_,
140 E : testing::ElementsAre("kernel32.dll"));
141 :
142 : // Make sure case is ignored.
143 E : transform.ExcludeModule("KERNEL32.DLL");
144 : EXPECT_THAT(transform.modules_to_exclude_,
145 E : testing::ElementsAre("kernel32.dll"));
146 :
147 : // Add another exclusion.
148 E : transform.ExcludeModule("export_dll.dll");
149 : EXPECT_THAT(transform.modules_to_exclude_,
150 E : testing::ElementsAre("export_dll.dll", "kernel32.dll"));
151 E : }
152 :
153 E : TEST_F(ThunkImportReferencesTransformTest, TestInstrumentation) {
154 : // Check that we don't have a thunks section yet.
155 E : EXPECT_EQ(NULL, block_graph_.FindSection(common::kThunkSectionName));
156 :
157 E : TestThunkImportReferencesTransform transform;
158 : // Exclude kernel32.dll.
159 E : transform.ExcludeModule("kernel32.dll");
160 :
161 : // Run the transform.
162 : ASSERT_TRUE(ApplyBlockGraphTransform(
163 E : &transform, &block_graph_, dos_header_block_));
164 :
165 : // Check that we now have a thunks section.
166 : BlockGraph::Section* thunks_section =
167 E : block_graph_.FindSection(common::kThunkSectionName);
168 E : ASSERT_TRUE(thunks_section != NULL);
169 :
170 E : BlockGraph::SectionId thunks_section_id = thunks_section->id();
171 :
172 : // From here it gets a little tricky. We use the import/delay import lookups
173 : // to find where the import entries are, and what their names are. We then
174 : // use that information to make sure that only non-thunks reference
175 : // kernel32.dll imports, and that only thunks reference other imports.
176 E : ImportAddressLocationNameMap all_import_locations;
177 : ASSERT_TRUE(TestThunkImportReferencesTransform::LookupImportLocations(
178 E : ModuleNameSet(), dos_header_block_, &all_import_locations));
179 : ASSERT_TRUE(TestThunkImportReferencesTransform::LookupDelayImportLocations(
180 E : ModuleNameSet(), dos_header_block_, &all_import_locations));
181 :
182 : // Retrieve the import blocks.
183 E : BlockSet import_blocks;
184 : ASSERT_TRUE(TestThunkImportReferencesTransform::GetImportBlocks(
185 E : all_import_locations, &import_blocks));
186 : // There should be two import blocks, the IAT and the delay IAT.
187 E : ASSERT_EQ(2, import_blocks.size());
188 :
189 E : BlockSet::const_iterator block_it = import_blocks.begin();
190 E : for (; block_it != import_blocks.end(); ++block_it) {
191 E : const ReferrerSet& referrers = (*block_it)->referrers();
192 E : ReferrerSet::const_iterator ref_iter(referrers.begin()),
193 E : ref_end(referrers.end());
194 E : for (; ref_iter != ref_end; ++ref_iter) {
195 E : const BlockGraph::Block* referring_block = ref_iter->first;
196 E : BlockGraph::Offset referring_offset = ref_iter->second;
197 :
198 E : if (referring_block->type() != BlockGraph::CODE_BLOCK)
199 E : continue;
200 :
201 E : BlockGraph::Reference ref;
202 E : ASSERT_TRUE(referring_block->GetReference(referring_offset, &ref));
203 :
204 : ImportAddressLocation location(
205 E : std::make_pair(ref.referenced(), ref.offset()));
206 : ImportAddressLocationNameMap::const_iterator location_it(
207 E : all_import_locations.find(location));
208 :
209 E : ASSERT_TRUE(location_it != all_import_locations.end());
210 :
211 : // Compare on the import name prefix case-insensitively to determine
212 : // whether this is an eligible or excluded import.
213 E : if (StartsWithASCII(location_it->second, "kernel32.dll:", false)) {
214 : // Excluded import, the referring block must not be in
215 : // the thunks section.
216 E : EXPECT_NE(thunks_section_id, referring_block->section());
217 E : } else {
218 : // Eligible import, if the referring block is not itself a thunk (such
219 : // as a linker-generated import or delay import thunk), it must be in
220 : // the thunks section.
221 E : if ((referring_block->attributes() & BlockGraph::THUNK) == 0) {
222 E : EXPECT_EQ(thunks_section_id, referring_block->section());
223 : }
224 : }
225 E : }
226 E : }
227 E : }
228 :
229 : } // namespace transforms
230 : } // namespace instrument
|