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/files/scoped_temp_dir.h"
22 : #include "base/strings/string_util.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/pe_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::PEAddImportsTransform;
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(), 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 : if (base::StartsWith(it->second, kExportDLLPrefix,
110 E : base::CompareCase::INSENSITIVE_ASCII))
111 E : export_dll_imports.insert(it->second);
112 E : }
113 :
114 : EXPECT_THAT(export_dll_imports,
115 : testing::ElementsAre("export_dll.dll:#7",
116 : "export_dll.dll:function1",
117 : "export_dll.dll:function3",
118 E : "export_dll.dll:kExportedData"));
119 E : }
120 :
121 E : TEST_F(ThunkImportReferencesTransformTest, LookupDelayImportLocations) {
122 E : ImportAddressLocationNameMap all_import_locations;
123 : // This should enumerate all delay imports.
124 : ASSERT_TRUE(
125 : TestThunkImportReferencesTransform::LookupDelayImportLocations(
126 E : ModuleNameSet(), header_block_, &all_import_locations));
127 :
128 : // There should be precisely one ole32.dll import.
129 E : ASSERT_EQ(1, all_import_locations.size());
130 :
131 : ASSERT_STRCASEEQ("ole32.dll:CoCreateGuid",
132 E : all_import_locations.begin()->second.c_str());
133 E : }
134 :
135 E : TEST_F(ThunkImportReferencesTransformTest, ExcludeModule) {
136 E : TestThunkImportReferencesTransform transform;
137 :
138 E : EXPECT_TRUE(transform.modules_to_exclude_.empty());
139 E : transform.ExcludeModule("kernel32.dll");
140 : EXPECT_THAT(transform.modules_to_exclude_,
141 E : testing::ElementsAre("kernel32.dll"));
142 :
143 : // Make sure case is ignored.
144 E : transform.ExcludeModule("KERNEL32.DLL");
145 : EXPECT_THAT(transform.modules_to_exclude_,
146 E : testing::ElementsAre("kernel32.dll"));
147 :
148 : // Add another exclusion.
149 E : transform.ExcludeModule("export_dll.dll");
150 : EXPECT_THAT(transform.modules_to_exclude_,
151 E : testing::ElementsAre("export_dll.dll", "kernel32.dll"));
152 E : }
153 :
154 E : TEST_F(ThunkImportReferencesTransformTest, TestInstrumentation) {
155 : // Check that we don't have a thunks section yet.
156 E : EXPECT_EQ(NULL, block_graph_.FindSection(common::kThunkSectionName));
157 :
158 E : TestThunkImportReferencesTransform transform;
159 : // Exclude kernel32.dll.
160 E : transform.ExcludeModule("kernel32.dll");
161 :
162 : // Run the transform.
163 : ASSERT_TRUE(ApplyBlockGraphTransform(
164 E : &transform, policy_, &block_graph_, header_block_));
165 :
166 : // Check that we now have a thunks section.
167 : BlockGraph::Section* thunks_section =
168 E : block_graph_.FindSection(common::kThunkSectionName);
169 E : ASSERT_TRUE(thunks_section != NULL);
170 :
171 E : BlockGraph::SectionId thunks_section_id = thunks_section->id();
172 :
173 : // From here it gets a little tricky. We use the import/delay import lookups
174 : // to find where the import entries are, and what their names are. We then
175 : // use that information to make sure that only non-thunks reference
176 : // kernel32.dll imports, and that only thunks reference other imports.
177 E : ImportAddressLocationNameMap all_import_locations;
178 : ASSERT_TRUE(TestThunkImportReferencesTransform::LookupImportLocations(
179 E : ModuleNameSet(), header_block_, &all_import_locations));
180 : ASSERT_TRUE(TestThunkImportReferencesTransform::LookupDelayImportLocations(
181 E : ModuleNameSet(), header_block_, &all_import_locations));
182 :
183 : // Retrieve the import blocks.
184 E : BlockSet import_blocks;
185 : ASSERT_TRUE(TestThunkImportReferencesTransform::GetImportBlocks(
186 E : all_import_locations, &import_blocks));
187 : // There should be two import blocks, the IAT and the delay IAT.
188 E : ASSERT_EQ(2, import_blocks.size());
189 :
190 E : BlockSet::const_iterator block_it = import_blocks.begin();
191 E : for (; block_it != import_blocks.end(); ++block_it) {
192 E : const ReferrerSet& referrers = (*block_it)->referrers();
193 E : ReferrerSet::const_iterator ref_iter(referrers.begin()),
194 E : ref_end(referrers.end());
195 E : for (; ref_iter != ref_end; ++ref_iter) {
196 E : const BlockGraph::Block* referring_block = ref_iter->first;
197 E : BlockGraph::Offset referring_offset = ref_iter->second;
198 :
199 E : if (referring_block->type() != BlockGraph::CODE_BLOCK)
200 E : continue;
201 :
202 E : BlockGraph::Reference ref;
203 E : ASSERT_TRUE(referring_block->GetReference(referring_offset, &ref));
204 :
205 : ImportAddressLocation location(
206 E : std::make_pair(ref.referenced(), ref.offset()));
207 : ImportAddressLocationNameMap::const_iterator location_it(
208 E : all_import_locations.find(location));
209 :
210 E : ASSERT_TRUE(location_it != all_import_locations.end());
211 :
212 : // Compare on the import name prefix case-insensitively to determine
213 : // whether this is an eligible or excluded import.
214 : if (base::StartsWith(location_it->second, "kernel32.dll:",
215 E : base::CompareCase::INSENSITIVE_ASCII)) {
216 : // Excluded import, the referring block must not be in
217 : // the thunks section.
218 E : EXPECT_NE(thunks_section_id, referring_block->section());
219 E : } else {
220 : // Eligible import, if the referring block is not itself a thunk (such
221 : // as a linker-generated import or delay import thunk), it must be in
222 : // the thunks section.
223 E : if ((referring_block->attributes() & BlockGraph::THUNK) == 0) {
224 E : EXPECT_EQ(thunks_section_id, referring_block->section());
225 : }
226 : }
227 E : }
228 E : }
229 E : }
230 :
231 : } // namespace transforms
232 : } // namespace instrument
|