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/instrument/instrumenters/archive_instrumenter.h"
16 :
17 : #include "base/files/file_util.h"
18 : #include "gtest/gtest.h"
19 : #include "syzygy/ar/unittest_util.h"
20 : #include "syzygy/core/unittest_util.h"
21 : #include "syzygy/instrument/instrumenters/asan_instrumenter.h"
22 : #include "syzygy/pe/unittest_util.h"
23 :
24 : namespace instrument {
25 : namespace instrumenters {
26 :
27 : namespace {
28 :
29 : // Some global state that is updated by IdentityInstrumenter.
30 : // NOTE: Because of these this unittest is not thread safe! We could do this
31 : // with a complicate usage of gmock, but that's overkill for this
32 : // scenario.
33 : size_t constructor_count = 0;
34 : size_t parse_count = 0;
35 : size_t instrument_count = 0;
36 E : std::set<base::FilePath> input_images;
37 E : std::set<base::FilePath> output_images;
38 :
39 : // An identity instrumenter. Simply copies the input-image to the output-image.
40 : class IdentityInstrumenter : public InstrumenterInterface {
41 : public:
42 E : IdentityInstrumenter() {
43 E : ++constructor_count;
44 E : }
45 :
46 : virtual bool ParseCommandLine(
47 E : const base::CommandLine* command_line) override {
48 E : ++parse_count;
49 E : input_image_ = command_line->GetSwitchValuePath("input-image");
50 E : output_image_ = command_line->GetSwitchValuePath("output-image");
51 E : input_images.insert(input_image_);
52 E : output_images.insert(output_image_);
53 E : return true;
54 E : }
55 :
56 E : virtual bool Instrument() override {
57 E : ++instrument_count;
58 E : base::CopyFile(input_image_, output_image_);
59 E : return true;
60 E : }
61 :
62 : base::FilePath input_image_;
63 : base::FilePath output_image_;
64 : };
65 :
66 E : InstrumenterInterface* IdentityInstrumenterFactory() {
67 E : return new IdentityInstrumenter();
68 E : }
69 :
70 E : InstrumenterInterface* AsanInstrumenterFactory() {
71 E : return new AsanInstrumenter();
72 E : }
73 :
74 : class ArchiveInstrumenterTest : public testing::PELibUnitTest {
75 : public:
76 E : virtual void SetUp() override {
77 E : testing::PELibUnitTest::SetUp();
78 :
79 E : CreateTemporaryDir(&temp_dir_);
80 : test_dll_dll_ = testing::GetExeTestDataRelativePath(
81 E : testing::kTestDllName);
82 E : zlib_lib_ = testing::GetSrcRelativePath(testing::kArchiveFile);
83 E : output_image_ = temp_dir_.Append(L"output.dat");
84 :
85 : command_line_.reset(
86 E : new base::CommandLine(base::FilePath(L"instrumenter.exe")));
87 E : command_line_->AppendSwitchPath("output-image", output_image_);
88 :
89 : // Reset function counts.
90 E : constructor_count = 0;
91 E : parse_count = 0;
92 E : instrument_count = 0;
93 E : input_images.clear();
94 E : output_images.clear();
95 E : }
96 :
97 E : virtual void TearDown() override { testing::PELibUnitTest::TearDown(); }
98 :
99 : base::FilePath temp_dir_;
100 : base::FilePath test_dll_dll_;
101 : base::FilePath zlib_lib_;
102 : base::FilePath output_image_;
103 :
104 : scoped_ptr<base::CommandLine> command_line_;
105 : };
106 :
107 : } // namespace
108 :
109 E : TEST_F(ArchiveInstrumenterTest, PassthroughForNonArchive) {
110 E : ArchiveInstrumenter inst(&IdentityInstrumenterFactory);
111 E : command_line_->AppendSwitchPath("input-image", test_dll_dll_);
112 :
113 E : EXPECT_TRUE(inst.ParseCommandLine(command_line_.get()));
114 E : EXPECT_TRUE(inst.Instrument());
115 E : EXPECT_EQ(1u, constructor_count);
116 E : EXPECT_EQ(1u, parse_count);
117 E : EXPECT_EQ(1u, instrument_count);
118 E : EXPECT_EQ(1u, input_images.size());
119 E : EXPECT_EQ(1u, output_images.size());
120 E : EXPECT_EQ(test_dll_dll_, *input_images.begin());
121 E : EXPECT_EQ(output_image_, *output_images.begin());
122 E : EXPECT_TRUE(base::PathExists(output_image_));
123 E : }
124 :
125 E : TEST_F(ArchiveInstrumenterTest, IteratesOverArchiveFiles) {
126 E : ArchiveInstrumenter inst(&IdentityInstrumenterFactory);
127 E : command_line_->AppendSwitchPath("input-image", zlib_lib_);
128 :
129 E : EXPECT_TRUE(inst.ParseCommandLine(command_line_.get()));
130 E : EXPECT_TRUE(inst.Instrument());
131 E : EXPECT_EQ(testing::kArchiveFileCount, constructor_count);
132 E : EXPECT_EQ(testing::kArchiveFileCount, parse_count);
133 E : EXPECT_EQ(testing::kArchiveFileCount, instrument_count);
134 E : EXPECT_EQ(testing::kArchiveFileCount, input_images.size());
135 E : EXPECT_EQ(testing::kArchiveFileCount, output_images.size());
136 E : EXPECT_EQ(0u, input_images.count(zlib_lib_));
137 E : EXPECT_EQ(0u, output_images.count(output_image_));
138 E : EXPECT_TRUE(base::PathExists(output_image_));
139 E : }
140 :
141 E : TEST_F(ArchiveInstrumenterTest, AsanInstrumentArchive) {
142 E : ArchiveInstrumenter inst(&AsanInstrumenterFactory);
143 E : command_line_->AppendSwitchPath("input-image", zlib_lib_);
144 :
145 E : EXPECT_TRUE(inst.ParseCommandLine(command_line_.get()));
146 E : EXPECT_TRUE(inst.Instrument());
147 E : EXPECT_TRUE(base::PathExists(output_image_));
148 E : }
149 :
150 : } // namespace instrumenters
151 : } // namespace instrument
|