1 : // Copyright 2014 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/coff_utils.h"
16 :
17 : #include "base/bind.h"
18 : #include "gmock/gmock.h"
19 : #include "gtest/gtest.h"
20 : #include "syzygy/pe/unittest_util.h"
21 :
22 : namespace pe {
23 :
24 : namespace {
25 :
26 : using block_graph::BlockGraph;
27 : using testing::_;
28 : using testing::Return;
29 :
30 : static const char kFunction2[] = "?function2@@YAHXZ";
31 : static const char kDebugS[] = ".debug$S";
32 :
33 : class LenientCoffUtilsTest : public testing::CoffUnitTest {
34 : public:
35 : MOCK_METHOD3(VisitCoffSymbol, bool(BlockGraph::Block*,
36 : BlockGraph::Block*,
37 E : BlockGraph::Offset));
38 : };
39 : typedef testing::StrictMock<LenientCoffUtilsTest> CoffUtilsTest;
40 :
41 : typedef std::set<std::string> StringSet;
42 : bool VisitCoffSymbolAndGrabName(StringSet* names,
43 : BlockGraph::Block* symbols_block,
44 : BlockGraph::Block* strings_block,
45 E : BlockGraph::Offset symbol_offset) {
46 E : DCHECK_NE(reinterpret_cast<StringSet*>(NULL), names);
47 E : DCHECK_NE(reinterpret_cast<BlockGraph::Block*>(NULL), symbols_block);
48 E : DCHECK_NE(reinterpret_cast<BlockGraph::Block*>(NULL), strings_block);
49 :
50 E : base::StringPiece name;
51 : EXPECT_TRUE(GetCoffSymbolName(symbols_block, strings_block, symbol_offset,
52 E : &name));
53 E : EXPECT_FALSE(name.empty());
54 E : names->insert(name.as_string());
55 E : return true;
56 E : }
57 :
58 : } // namespace
59 :
60 E : TEST_F(CoffUtilsTest, FindCoffSpecialBlocks) {
61 E : BlockGraph::Block* actual_headers_block = NULL;
62 E : BlockGraph::Block* actual_symbols_block = NULL;
63 E : BlockGraph::Block* actual_strings_block = NULL;
64 :
65 : BlockGraph::Block* headers_block =
66 : block_graph_.AddBlock(
67 : BlockGraph::DATA_BLOCK,
68 : sizeof(IMAGE_FILE_HEADER) + 12 * sizeof(IMAGE_SECTION_HEADER),
69 E : "COFF Headers");
70 E : ASSERT_TRUE(headers_block != NULL);
71 E : headers_block->set_attribute(BlockGraph::COFF_HEADERS);
72 :
73 : // FindCoffSpecialBlocks() should fail even if we don't request the other
74 : // special blocks.
75 : EXPECT_FALSE(FindCoffSpecialBlocks(&block_graph_,
76 : &actual_headers_block,
77 : &actual_symbols_block,
78 E : &actual_strings_block));
79 : EXPECT_FALSE(FindCoffSpecialBlocks(&block_graph_,
80 E : &actual_headers_block, NULL, NULL));
81 :
82 : BlockGraph::Block* symbols_block =
83 : block_graph_.AddBlock(BlockGraph::DATA_BLOCK,
84 : 30 * sizeof(IMAGE_SYMBOL),
85 E : "COFF Symbol Table");
86 E : ASSERT_TRUE(symbols_block != NULL);
87 E : symbols_block->set_attribute(BlockGraph::COFF_SYMBOL_TABLE);
88 :
89 : EXPECT_FALSE(FindCoffSpecialBlocks(&block_graph_,
90 : &actual_headers_block,
91 : &actual_symbols_block,
92 E : &actual_strings_block));
93 : EXPECT_FALSE(FindCoffSpecialBlocks(&block_graph_,
94 : &actual_headers_block,
95 : &actual_symbols_block,
96 E : NULL));
97 :
98 : BlockGraph::Block* strings_block =
99 E : block_graph_.AddBlock(BlockGraph::DATA_BLOCK, 242, "COFF String Table");
100 E : ASSERT_TRUE(strings_block != NULL);
101 E : strings_block->set_attribute(BlockGraph::COFF_STRING_TABLE);
102 :
103 : EXPECT_TRUE(FindCoffSpecialBlocks(&block_graph_,
104 : &actual_headers_block,
105 : &actual_symbols_block,
106 E : &actual_strings_block));
107 E : EXPECT_EQ(headers_block, actual_headers_block);
108 E : EXPECT_EQ(symbols_block, actual_symbols_block);
109 E : EXPECT_EQ(strings_block, actual_strings_block);
110 E : }
111 :
112 E : TEST_F(CoffUtilsTest, VisitCoffSymbols) {
113 E : ASSERT_NO_FATAL_FAILURE(DecomposeOriginal());
114 :
115 E : BlockGraph::Block* symbols_block = NULL;
116 E : BlockGraph::Block* strings_block = NULL;
117 : ASSERT_TRUE(FindCoffSpecialBlocks(&block_graph_,
118 : NULL,
119 : &symbols_block,
120 E : &strings_block));
121 :
122 : VisitCoffSymbolCallback callback = base::Bind(
123 E : &CoffUtilsTest::VisitCoffSymbol, base::Unretained(this));
124 :
125 : // Expect the visitor to fail if the callback does.
126 : EXPECT_CALL(*this, VisitCoffSymbol(symbols_block,
127 : strings_block,
128 E : _)).WillOnce(Return(false));
129 E : EXPECT_FALSE(VisitCoffSymbols(callback, &block_graph_));
130 :
131 : // Now expect the visitor to succeed.
132 : EXPECT_CALL(*this, VisitCoffSymbol(symbols_block,
133 : strings_block,
134 : _)).
135 E : WillRepeatedly(Return(true));
136 E : EXPECT_TRUE(VisitCoffSymbols(callback, &block_graph_));
137 E : }
138 :
139 E : TEST_F(CoffUtilsTest, GetCoffSymbolName) {
140 E : ASSERT_NO_FATAL_FAILURE(DecomposeOriginal());
141 :
142 E : StringSet names;
143 : VisitCoffSymbolCallback callback = base::Bind(
144 E : &VisitCoffSymbolAndGrabName, base::Unretained(&names));
145 :
146 E : EXPECT_TRUE(VisitCoffSymbols(callback, &block_graph_));
147 E : EXPECT_FALSE(names.empty());
148 E : }
149 :
150 E : TEST_F(CoffUtilsTest, FindCoffSymbolInvalid) {
151 E : ASSERT_NO_FATAL_FAILURE(DecomposeOriginal());
152 :
153 E : CoffSymbolOffsets offsets;
154 E : EXPECT_TRUE(FindCoffSymbol("_foo_bar_baz", &block_graph_, &offsets));
155 E : EXPECT_TRUE(offsets.empty());
156 E : }
157 :
158 E : TEST_F(CoffUtilsTest, FindCoffSymbolDuplicate) {
159 E : ASSERT_NO_FATAL_FAILURE(DecomposeOriginal());
160 :
161 E : CoffSymbolOffsets offsets;
162 E : EXPECT_TRUE(FindCoffSymbol(kDebugS, &block_graph_, &offsets));
163 E : EXPECT_LT(1u, offsets.size());
164 E : }
165 :
166 E : TEST_F(CoffUtilsTest, FindCoffSymbolUnique) {
167 E : ASSERT_NO_FATAL_FAILURE(DecomposeOriginal());
168 :
169 E : CoffSymbolOffsets offsets;
170 E : EXPECT_TRUE(FindCoffSymbol(kFunction2, &block_graph_, &offsets));
171 E : EXPECT_EQ(1u, offsets.size());
172 E : }
173 :
174 E : TEST_F(CoffUtilsTest, BuildCoffSymbolNameOffsetMap) {
175 E : ASSERT_NO_FATAL_FAILURE(DecomposeOriginal());
176 :
177 E : CoffSymbolNameOffsetMap map;
178 E : EXPECT_TRUE(BuildCoffSymbolNameOffsetMap(&block_graph_, &map));
179 E : EXPECT_FALSE(map.empty());
180 :
181 E : CoffSymbolNameOffsetMap::const_iterator it = map.find("_foo_bar_baz");
182 E : EXPECT_TRUE(it == map.end());
183 :
184 E : it = map.find(kFunction2);
185 E : ASSERT_TRUE(it != map.end());
186 E : EXPECT_FALSE(it->second.empty());
187 :
188 E : it = map.find(kDebugS);
189 E : ASSERT_TRUE(it != map.end());
190 E : EXPECT_LT(1u, it->second.size());
191 E : }
192 :
193 : } // namespace pe
|