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 : #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 E : FilePath image_path(testing::GetExeRelativePath(kDllName));
60 E : ASSERT_TRUE(pe_file_.Init(image_path));
61 E : }
62 :
63 E : void InitDecomposition() {
64 E : ASSERT_NO_FATAL_FAILURE(InitPEFile());
65 E : Decomposer decomposer(pe_file_);
66 E : ASSERT_TRUE(decomposer.Decompose(&image_layout_));
67 E : }
68 :
69 E : void InitOutArchive() {
70 E : v_.clear();
71 E : os_.reset(core::CreateByteOutStream(std::back_inserter(v_)));
72 : oa_.reset(new core::NativeBinaryOutArchive(
73 E : os_.get()));
74 E : }
75 :
76 E : void InitInArchive() {
77 E : is_.reset(core::CreateByteInStream(v_.begin(), v_.end()));
78 : ia_.reset(new core::NativeBinaryInArchive(
79 E : is_.get()));
80 E : }
81 :
82 E : void Serialize(BlockGraphSerializer::Attributes attributes) {
83 E : ASSERT_TRUE(SaveBlockGraphAndImageLayout(pe_file_,
84 : attributes,
85 : image_layout_,
86 : oa_.get()));
87 E : }
88 :
89 : void TestRoundTrip(BlockGraphSerializer::Attributes attributes,
90 E : bool search_for_pe_file) {
91 E : ASSERT_NO_FATAL_FAILURE(InitDecomposition());
92 E : ASSERT_NO_FATAL_FAILURE(InitOutArchive());
93 E : ASSERT_NO_FATAL_FAILURE(Serialize(attributes));
94 :
95 E : ASSERT_NO_FATAL_FAILURE(InitInArchive());
96 : BlockGraphSerializer::Attributes attributes2;
97 E : PEFile pe_file;
98 E : BlockGraph block_graph;
99 E : ImageLayout image_layout(&block_graph);
100 :
101 E : if (search_for_pe_file) {
102 E : ASSERT_TRUE(LoadBlockGraphAndImageLayout(&pe_file,
103 : &attributes2,
104 : &image_layout,
105 : ia_.get()));
106 E : } else {
107 E : ASSERT_TRUE(LoadBlockGraphAndImageLayout(pe_file_,
108 : &attributes2,
109 : &image_layout,
110 : ia_.get()));
111 : }
112 :
113 E : ASSERT_EQ(attributes, attributes2);
114 :
115 E : BlockGraphSerializer bgs;
116 E : bgs.set_data_mode(BlockGraphSerializer::OUTPUT_NO_DATA);
117 E : bgs.set_attributes(attributes);
118 E : ASSERT_TRUE(testing::BlockGraphsEqual(block_graph_, block_graph, bgs));
119 E : ASSERT_TRUE(ImageLayoutsEqual(image_layout_, image_layout));
120 E : }
121 :
122 : // Decomposition information.
123 : PEFile pe_file_;
124 : BlockGraph block_graph_;
125 : ImageLayout image_layout_;
126 :
127 : // Streams and archives.
128 : std::vector<uint8> v_;
129 : scoped_ptr<core::OutStream> os_;
130 : scoped_ptr<core::InStream> is_;
131 : scoped_ptr<core::OutArchive> oa_;
132 : scoped_ptr<core::InArchive> ia_;
133 : };
134 :
135 : } // namespace
136 :
137 E : TEST_F(SerializationTest, TestDllRoundTripFull) {
138 : ASSERT_NO_FATAL_FAILURE(
139 E : TestRoundTrip(BlockGraphSerializer::DEFAULT_ATTRIBUTES, true));
140 E : }
141 :
142 E : TEST_F(SerializationTest, TestDllRoundTripNoStrings) {
143 : ASSERT_NO_FATAL_FAILURE(
144 E : TestRoundTrip(BlockGraphSerializer::OMIT_STRINGS, false));
145 E : }
146 :
147 E : TEST_F(SerializationTest, FailsForInvalidVersion) {
148 E : ASSERT_NO_FATAL_FAILURE(InitOutArchive());
149 E : ASSERT_NO_FATAL_FAILURE(InitDecomposition());
150 E : ASSERT_NO_FATAL_FAILURE(Serialize(0));
151 E : ASSERT_NO_FATAL_FAILURE(InitInArchive());
152 :
153 : // Change the version.
154 E : v_[0] += 1;
155 :
156 E : PEFile pe_file;
157 E : BlockGraph block_graph;
158 E : ImageLayout image_layout(&block_graph);
159 : ASSERT_FALSE(LoadBlockGraphAndImageLayout(
160 E : &pe_file, NULL, &image_layout, ia_.get()));
161 E : }
162 :
163 : // TODO(chrisha): Check in a serialized stream, and ensure that it can still be
164 : // deserialized. As we evolve stream versions, keep doing this. This will be
165 : // done once decompose.exe has been updated to use the new serialization
166 : // engine.
167 :
168 : } // namespace pe
|