1 : // Copyright 2012 Google Inc.
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 the Asan transform.
16 :
17 : #include "syzygy/instrument/transforms/asan_transform.h"
18 :
19 : #include <set>
20 : #include <vector>
21 :
22 : #include "base/scoped_native_library.h"
23 : #include "base/scoped_temp_dir.h"
24 : #include "base/win/pe_image.h"
25 : #include "gtest/gtest.h"
26 : #include "syzygy/block_graph/basic_block_assembler.h"
27 : #include "syzygy/core/unittest_util.h"
28 : #include "syzygy/pe/decomposer.h"
29 : #include "syzygy/pe/pe_file.h"
30 : #include "syzygy/pe/pe_relinker.h"
31 : #include "syzygy/pe/pe_utils.h"
32 : #include "syzygy/pe/unittest_util.h"
33 : #include "third_party/distorm/files/include/mnemonics.h"
34 :
35 : namespace instrument {
36 : namespace transforms {
37 :
38 : namespace {
39 :
40 : using block_graph::BasicBlock;
41 : using block_graph::BasicBlockSubGraph;
42 : using block_graph::BlockGraph;
43 :
44 : // A derived class to expose protected members for unit-testing.
45 : class TestAsanBasicBlockTransform : public AsanBasicBlockTransform {
46 : public:
47 : using AsanBasicBlockTransform::InstrumentBasicBlock;
48 :
49 E : explicit TestAsanBasicBlockTransform(BlockGraph::Reference* hook)
50 : : AsanBasicBlockTransform(hook) {
51 E : }
52 : };
53 :
54 : class AsanTransformTest : public testing::PELibUnitTest {
55 : public:
56 : AsanTransformTest() :
57 : dos_header_block_(NULL),
58 : basic_block_(0, "test block", BasicBlock::BASIC_CODE_BLOCK,
59 : BasicBlock::kNoOffset, kDataSize, kBlockData),
60 : bb_asm_(basic_block_.instructions().begin(),
61 E : &basic_block_.instructions()) {
62 E : }
63 :
64 E : void DecomposeTestDll() {
65 E : FilePath test_dll_path = ::testing::GetOutputRelativePath(kDllName);
66 :
67 E : ASSERT_TRUE(pe_file_.Init(test_dll_path));
68 :
69 E : pe::ImageLayout layout(&block_graph_);
70 E : pe::Decomposer decomposer(pe_file_);
71 E : ASSERT_TRUE(decomposer.Decompose(&layout));
72 :
73 : dos_header_block_ =
74 E : layout.blocks.GetBlockByAddress(core::RelativeAddress(0));
75 E : ASSERT_TRUE(dos_header_block_ != NULL);
76 E : }
77 :
78 E : void InitHookRefs() {
79 : hook_check_access_ = block_graph_.AddBlock(BlockGraph::CODE_BLOCK, 4,
80 : "hook_check_access"),
81 : // Set up the references to the hooks needed by SyzyAsan.
82 : hook_check_access_ref_ = BlockGraph::Reference(BlockGraph::ABSOLUTE_REF, 4,
83 E : hook_check_access_, 0, 0);
84 E : }
85 :
86 : // Some handy constants we'll use throughout the tests.
87 : // @{
88 : static const BasicBlock::Size kDataSize;
89 : static const uint8 kBlockData[];
90 : // @}
91 :
92 : protected:
93 : ScopedTempDir temp_dir_;
94 : pe::PEFile pe_file_;
95 : BlockGraph block_graph_;
96 : BlockGraph::Block* dos_header_block_;
97 : AsanTransform asan_transform_;
98 : BlockGraph::Block* hook_check_access_;
99 : BlockGraph::Reference hook_check_access_ref_;
100 : BasicBlock basic_block_;
101 : block_graph::BasicBlockAssembler bb_asm_;
102 :
103 : };
104 :
105 : const BasicBlock::Size AsanTransformTest::kDataSize = 32;
106 : const uint8 AsanTransformTest::kBlockData[AsanTransformTest::kDataSize] = {};
107 :
108 : } // namespace
109 :
110 E : TEST_F(AsanTransformTest, SetInstrumentDLLName) {
111 E : asan_transform_.set_instrument_dll_name("foo");
112 E : ASSERT_EQ(strcmp(asan_transform_.instrument_dll_name(), "foo"), 0);
113 E : }
114 :
115 E : TEST_F(AsanTransformTest, ApplyAsanTransform) {
116 E : ASSERT_NO_FATAL_FAILURE(DecomposeTestDll());
117 :
118 : ASSERT_TRUE(block_graph::ApplyBlockGraphTransform(
119 E : &asan_transform_, &block_graph_, dos_header_block_));
120 :
121 : // TODO(sebmarchand): Ensure that each memory access is instrumented by
122 : // decomposing each block of the new block-graph into basic blocks and walk
123 : // through their instructions. For now it's not possible due to an issue with
124 : // the labels in the new block-graph.
125 E : }
126 :
127 E : TEST_F(AsanTransformTest, InjectAsanHooks) {
128 : // Add a read access to the memory.
129 E : bb_asm_.mov(core::eax, block_graph::Operand(core::ebx));
130 : // Add a write access to the memory.
131 E : bb_asm_.mov(block_graph::Operand(core::ecx), core::edx);
132 :
133 : // Instrument this basic block.
134 E : InitHookRefs();
135 E : TestAsanBasicBlockTransform bb_transform(&hook_check_access_ref_);
136 E : ASSERT_TRUE(bb_transform.InstrumentBasicBlock(&basic_block_));
137 :
138 : // Ensure that the basic block is instrumented.
139 :
140 : // We had 2 instructions initially, and for each of them we add 3
141 : // instructions, so we expect to have 2 + 3 * 2 = 8 instructions.
142 E : ASSERT_EQ(basic_block_.instructions().size(), 8);
143 :
144 : // Walk through the instructions to ensure that the Asan hooks have been
145 : // injected.
146 : BasicBlock::Instructions::const_iterator iter_inst =
147 E : basic_block_.instructions().begin();
148 :
149 : // First we check if the first memory access is instrumented as a read
150 : // access.
151 E : ASSERT_TRUE((iter_inst++)->representation().opcode == I_PUSH);
152 E : ASSERT_TRUE((iter_inst++)->representation().opcode == I_LEA);
153 E : ASSERT_EQ(iter_inst->references().size(), 1);
154 : ASSERT_TRUE(
155 E : iter_inst->references().begin()->second.block() == hook_check_access_);
156 E : ASSERT_TRUE((iter_inst++)->representation().opcode == I_CALL);
157 E : ASSERT_TRUE((iter_inst++)->representation().opcode == I_MOV);
158 :
159 : // Then we check if the second memory access is well instrumented as a write
160 : // access.
161 E : ASSERT_TRUE((iter_inst++)->representation().opcode == I_PUSH);
162 E : ASSERT_TRUE((iter_inst++)->representation().opcode == I_LEA);
163 E : ASSERT_EQ(iter_inst->references().size(), 1);
164 : ASSERT_TRUE(
165 E : iter_inst->references().begin()->second.block() == hook_check_access_);
166 E : ASSERT_TRUE((iter_inst++)->representation().opcode == I_CALL);
167 E : ASSERT_TRUE((iter_inst++)->representation().opcode == I_MOV);
168 :
169 E : ASSERT_TRUE(iter_inst == basic_block_.instructions().end());
170 E : }
171 :
172 E : TEST_F(AsanTransformTest, InstrumentDifferentKindOfInstructions) {
173 E : uint32 instrumentable_instructions = 0;
174 :
175 : // Generate a bunch of instrumentable and non instrumentable instructions.
176 E : bb_asm_.mov(core::eax, block_graph::Operand(core::ebx));
177 E : instrumentable_instructions++;
178 E : bb_asm_.mov(block_graph::Operand(core::ecx), core::edx);
179 E : instrumentable_instructions++;
180 :
181 : // Non-instrumentable.
182 E : bb_asm_.call(block_graph::Operand(core::ecx));
183 E : bb_asm_.push(block_graph::Operand(core::eax));
184 E : instrumentable_instructions++;
185 :
186 : // Non-instrumentable.
187 E : bb_asm_.lea(core::eax, block_graph::Operand(core::ecx));
188 :
189 : uint32 expected_instructions_count = basic_block_.instructions().size()
190 E : + 3 * instrumentable_instructions;
191 : // Instrument this basic block.
192 E : InitHookRefs();
193 E : TestAsanBasicBlockTransform bb_transform(&hook_check_access_ref_);
194 E : ASSERT_TRUE(bb_transform.InstrumentBasicBlock(&basic_block_));
195 E : ASSERT_EQ(basic_block_.instructions().size(), expected_instructions_count);
196 E : }
197 :
198 : namespace {
199 : using base::win::PEImage;
200 : typedef std::set<std::string> StringSet;
201 :
202 : bool EnumImports(const PEImage &image, LPCSTR module,
203 : DWORD ordinal, LPCSTR name, DWORD hint,
204 E : PIMAGE_THUNK_DATA iat, PVOID cookie) {
205 E : StringSet* modules = reinterpret_cast<StringSet*>(cookie);
206 :
207 E : if (strcmp("asan_rtl.dll", module) == 0)
208 E : modules->insert(name);
209 :
210 E : return true;
211 E : }
212 :
213 : };
214 :
215 E : TEST_F(AsanTransformTest, ImportsAreRedirected) {
216 E : pe::PERelinker relinker;
217 :
218 E : ASSERT_TRUE(temp_dir_.CreateUniqueTempDir());
219 :
220 E : relinker.set_input_path(::testing::GetOutputRelativePath(kDllName));
221 E : relinker.set_output_path(temp_dir_.path().Append(kDllName));
222 :
223 E : relinker.AppendTransform(&asan_transform_);
224 E : ASSERT_TRUE(relinker.Init());
225 E : ASSERT_TRUE(relinker.Relink());
226 :
227 : // Load the transformed module without resolving its dependencies.
228 : base::NativeLibrary lib =
229 : ::LoadLibraryEx(relinker.output_path().value().c_str(),
230 : NULL,
231 E : DONT_RESOLVE_DLL_REFERENCES);
232 E : ASSERT_TRUE(lib != NULL);
233 : // Make sure it's unloaded on failure.
234 E : base::ScopedNativeLibrary lib_keeper(lib);
235 :
236 E : PEImage image(lib);
237 E : ASSERT_TRUE(image.VerifyMagic());
238 E : StringSet imports;
239 E : ASSERT_TRUE(image.EnumAllImports(&EnumImports, &imports));
240 :
241 : // This isn't strictly speaking a full test, as we only check that the new
242 : // imports have been added. It's however more trouble than it's worth to
243 : // test this fully for now.
244 E : StringSet expected;
245 E : expected.insert("asan_HeapCreate");
246 E : expected.insert("asan_HeapDestroy");
247 E : expected.insert("asan_HeapAlloc");
248 E : expected.insert("asan_HeapReAlloc");
249 E : expected.insert("asan_HeapFree");
250 E : expected.insert("asan_HeapSize");
251 E : expected.insert("asan_HeapValidate");
252 E : expected.insert("asan_HeapCompact");
253 E : expected.insert("asan_HeapLock");
254 E : expected.insert("asan_HeapUnlock");
255 E : expected.insert("asan_HeapWalk");
256 E : expected.insert("asan_HeapSetInformation");
257 E : expected.insert("asan_HeapQueryInformation");
258 E : expected.insert("asan_check_access");
259 :
260 E : EXPECT_EQ(expected, imports);
261 E : }
262 :
263 : } // namespace transforms
264 : } // namespace instrument
|