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 : #include "syzygy/pe/serialization.h"
16 :
17 : #include "gtest/gtest.h"
18 : #include "syzygy/block_graph/unittest_util.h"
19 : #include "syzygy/core/unittest_util.h"
20 : #include "syzygy/pe/decomposer.h"
21 : #include "syzygy/pe/image_layout.h"
22 : #include "syzygy/pe/pe_file.h"
23 : #include "syzygy/pe/unittest_util.h"
24 :
25 : namespace pe {
26 :
27 : namespace {
28 :
29 : using block_graph::BlockGraph;
30 : using block_graph::BlockGraphSerializer;
31 :
32 : // Compares two image-layouts for equality.
33 E : bool ImageLayoutsEqual(const ImageLayout& il1, const ImageLayout& il2) {
34 E : if (il1.sections != il2.sections)
35 i : return false;
36 :
37 E : if (il1.blocks.size() != il2.blocks.size())
38 i : return false;
39 :
40 : typedef block_graph::BlockGraph::AddressSpace::RangeMapConstIter ConstIt;
41 E : ConstIt it1 = il1.blocks.begin();
42 E : ConstIt it2 = il2.blocks.begin();
43 E : for (; it1 != il1.blocks.end(); ++it1, ++it2) {
44 E : if (it1->first != it2->first)
45 i : return false;
46 E : if (it1->second->id() != it2->second->id())
47 i : return false;
48 E : }
49 :
50 E : return true;
51 E : }
52 :
53 : class SerializationTest : public testing::PELibUnitTest {
54 : public:
55 E : SerializationTest() : image_layout_(&block_graph_) { }
56 E : virtual void SetUp() override {}
57 :
58 E : void InitPEFile() {
59 : base::FilePath image_path(
60 E : testing::GetExeRelativePath(testing::kTestDllName));
61 E : ASSERT_TRUE(pe_file_.Init(image_path));
62 E : }
63 :
64 E : void InitDecomposition() {
65 E : ASSERT_NO_FATAL_FAILURE(InitPEFile());
66 E : Decomposer decomposer(pe_file_);
67 E : ASSERT_TRUE(decomposer.Decompose(&image_layout_));
68 E : }
69 :
70 E : void InitOutArchive() {
71 E : v_.clear();
72 E : os_.reset(core::CreateByteOutStream(std::back_inserter(v_)));
73 : oa_.reset(new core::NativeBinaryOutArchive(
74 E : os_.get()));
75 E : }
76 :
77 E : void InitInArchive() {
78 E : is_.reset(core::CreateByteInStream(v_.begin(), v_.end()));
79 : ia_.reset(new core::NativeBinaryInArchive(
80 E : is_.get()));
81 E : }
82 :
83 E : void Serialize(BlockGraphSerializer::Attributes attributes) {
84 E : ASSERT_TRUE(SaveBlockGraphAndImageLayout(pe_file_,
85 : attributes,
86 : image_layout_,
87 : oa_.get()));
88 E : }
89 :
90 : void TestRoundTrip(BlockGraphSerializer::Attributes attributes,
91 E : bool search_for_pe_file) {
92 E : ASSERT_NO_FATAL_FAILURE(InitDecomposition());
93 E : ASSERT_NO_FATAL_FAILURE(InitOutArchive());
94 E : ASSERT_NO_FATAL_FAILURE(Serialize(attributes));
95 :
96 E : ASSERT_NO_FATAL_FAILURE(InitInArchive());
97 : BlockGraphSerializer::Attributes attributes2;
98 E : PEFile pe_file;
99 E : BlockGraph block_graph;
100 E : ImageLayout image_layout(&block_graph);
101 :
102 E : if (search_for_pe_file) {
103 E : ASSERT_TRUE(LoadBlockGraphAndImageLayout(&pe_file,
104 : &attributes2,
105 : &image_layout,
106 : ia_.get()));
107 E : } else {
108 E : ASSERT_TRUE(LoadBlockGraphAndImageLayout(pe_file_,
109 : &attributes2,
110 : &image_layout,
111 : ia_.get()));
112 : }
113 :
114 E : ASSERT_EQ(attributes, attributes2);
115 :
116 E : BlockGraphSerializer bgs;
117 E : bgs.set_data_mode(BlockGraphSerializer::OUTPUT_NO_DATA);
118 E : bgs.set_attributes(attributes);
119 E : ASSERT_TRUE(testing::BlockGraphsEqual(block_graph_, block_graph, bgs));
120 E : ASSERT_TRUE(ImageLayoutsEqual(image_layout_, image_layout));
121 E : }
122 :
123 : // Decomposition information.
124 : PEFile pe_file_;
125 : BlockGraph block_graph_;
126 : ImageLayout image_layout_;
127 :
128 : // Streams and archives.
129 : std::vector<uint8> v_;
130 : scoped_ptr<core::OutStream> os_;
131 : scoped_ptr<core::InStream> is_;
132 : scoped_ptr<core::OutArchive> oa_;
133 : scoped_ptr<core::InArchive> ia_;
134 : };
135 :
136 : } // namespace
137 :
138 E : TEST_F(SerializationTest, TestDllRoundTripFull) {
139 : ASSERT_NO_FATAL_FAILURE(
140 E : TestRoundTrip(BlockGraphSerializer::DEFAULT_ATTRIBUTES, true));
141 E : }
142 :
143 E : TEST_F(SerializationTest, TestDllRoundTripNoStrings) {
144 : ASSERT_NO_FATAL_FAILURE(
145 E : TestRoundTrip(BlockGraphSerializer::OMIT_STRINGS, false));
146 E : }
147 :
148 E : TEST_F(SerializationTest, FailsForInvalidVersion) {
149 E : ASSERT_NO_FATAL_FAILURE(InitOutArchive());
150 E : ASSERT_NO_FATAL_FAILURE(InitDecomposition());
151 E : ASSERT_NO_FATAL_FAILURE(Serialize(0));
152 E : ASSERT_NO_FATAL_FAILURE(InitInArchive());
153 :
154 : // Change the version.
155 E : v_[0] += 1;
156 :
157 E : PEFile pe_file;
158 E : BlockGraph block_graph;
159 E : ImageLayout image_layout(&block_graph);
160 : ASSERT_FALSE(LoadBlockGraphAndImageLayout(
161 E : &pe_file, NULL, &image_layout, ia_.get()));
162 E : }
163 :
164 : // TODO(chrisha): Check in a serialized stream, and ensure that it can still be
165 : // deserialized. As we evolve stream versions, keep doing this. This will be
166 : // done once decompose.exe has been updated to use the new serialization
167 : // engine.
168 :
169 : } // namespace pe
|