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/pehacker/operations/add_imports_operation.h"
16 :
17 : #include "base/json/json_reader.h"
18 : #include "gmock/gmock.h"
19 : #include "gtest/gtest.h"
20 : #include "syzygy/core/unittest_util.h"
21 : #include "syzygy/pe/pe_transform_policy.h"
22 :
23 : namespace pehacker {
24 : namespace operations {
25 :
26 : namespace {
27 :
28 : using block_graph::BlockGraph;
29 : using testing::Return;
30 :
31 : const char kSimpleConfig[] =
32 : "{\n"
33 : " \"type\": \"add_imports\",\n"
34 : " \"modules\": [\n"
35 : " {\n"
36 : " \"module_name\": \"foo.dll\","
37 : " \"imports\": [\n"
38 : " { \"function_name\": \"bar\" },\n"
39 : " ]\n"
40 : " },\n"
41 : " ],\n"
42 : "}";
43 :
44 : class TestAddImportsOperation : public AddImportsOperation {
45 : public:
46 E : TestAddImportsOperation() { }
47 E : virtual ~TestAddImportsOperation() { }
48 :
49 : MOCK_METHOD4(ApplyTransform, bool(block_graph::BlockGraphTransformInterface*,
50 : const TransformPolicyInterface*,
51 : BlockGraph*,
52 E : BlockGraph::Block*));
53 :
54 : typedef AddImportsOperation::ImportedModuleMap ImportedModuleMap;
55 :
56 : using AddImportsOperation::add_imports_tx_;
57 : using AddImportsOperation::imported_modules_;
58 : using AddImportsOperation::imported_module_map_;
59 : };
60 :
61 : class AddImportsOperationTest : public testing::Test {
62 : public:
63 E : AddImportsOperationTest() : previous_log_level_(0) {
64 E : }
65 :
66 E : void SetUp() {
67 : // Silence logging.
68 E : previous_log_level_ = logging::GetMinLogLevel();
69 E : logging::SetMinLogLevel(logging::LOG_FATAL);
70 E : }
71 :
72 E : void TearDown() {
73 : // Restore logging to its previous level.
74 E : logging::SetMinLogLevel(previous_log_level_);
75 E : previous_log_level_ = 0;
76 E : }
77 :
78 E : void InitConfig() {
79 : scoped_ptr<base::Value> value(base::JSONReader::Read(
80 E : kSimpleConfig, base::JSON_ALLOW_TRAILING_COMMAS));
81 E : ASSERT_TRUE(value.get() != NULL);
82 E : base::DictionaryValue* dict = NULL;
83 E : ASSERT_TRUE(value->GetAsDictionary(&dict));
84 E : config_.reset(dict);
85 E : value.release();
86 E : }
87 :
88 : int previous_log_level_;
89 : scoped_ptr<base::DictionaryValue> config_;
90 : };
91 :
92 : } // namespace
93 :
94 E : TEST_F(AddImportsOperationTest, Name) {
95 E : TestAddImportsOperation op;
96 E : EXPECT_STREQ(TestAddImportsOperation::kName, op.name());
97 E : }
98 :
99 E : TEST_F(AddImportsOperationTest, Init) {
100 E : TestAddImportsOperation op;
101 E : ASSERT_NO_FATAL_FAILURE(InitConfig());
102 E : pe::PETransformPolicy policy;
103 E : EXPECT_TRUE(op.Init(&policy, config_.get()));
104 :
105 E : EXPECT_EQ(1u, op.imported_modules_.size());
106 E : EXPECT_EQ(1u, op.imported_module_map_.size());
107 :
108 : // Ensure the transform is appropriately configured.
109 : TestAddImportsOperation::ImportedModuleMap::iterator mod_it =
110 E : op.imported_module_map_.begin();
111 E : EXPECT_EQ("foo.dll", mod_it->first);
112 E : EXPECT_EQ(op.imported_modules_[0], mod_it->second);
113 E : EXPECT_EQ("foo.dll", mod_it->second->name());
114 E : EXPECT_EQ(1u, mod_it->second->size());
115 E : EXPECT_EQ("bar", mod_it->second->GetSymbolName(0));
116 : EXPECT_EQ(pe::transforms::ImportedModule::kAlwaysImport,
117 E : mod_it->second->GetSymbolMode(0));
118 E : }
119 :
120 E : TEST_F(AddImportsOperationTest, RunFails) {
121 E : TestAddImportsOperation op;
122 E : ASSERT_NO_FATAL_FAILURE(InitConfig());
123 E : pe::PETransformPolicy policy;
124 E : ASSERT_TRUE(op.Init(&policy, config_.get()));
125 :
126 E : BlockGraph bg;
127 E : BlockGraph::Block* header = bg.AddBlock(BlockGraph::DATA_BLOCK, 1, "header");
128 :
129 : EXPECT_CALL(op, ApplyTransform(&op.add_imports_tx_,
130 : &policy,
131 : &bg,
132 E : header)).WillOnce(Return(false));
133 :
134 E : EXPECT_FALSE(op.Apply(&policy, &bg, header));
135 E : }
136 :
137 E : TEST_F(AddImportsOperationTest, RunSucceeds) {
138 E : TestAddImportsOperation op;
139 E : ASSERT_NO_FATAL_FAILURE(InitConfig());
140 E : pe::PETransformPolicy policy;
141 E : ASSERT_TRUE(op.Init(&policy, config_.get()));
142 :
143 E : BlockGraph bg;
144 E : BlockGraph::Block* header = bg.AddBlock(BlockGraph::DATA_BLOCK, 1, "header");
145 :
146 : EXPECT_CALL(op, ApplyTransform(&op.add_imports_tx_,
147 : &policy,
148 : &bg,
149 E : header)).WillOnce(Return(true));
150 :
151 E : EXPECT_TRUE(op.Apply(&policy, &bg, header));
152 E : }
153 :
154 : } // namespace operations
155 : } // namespace pehacker
|