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/msf/msf_writer.h"
16 :
17 : #include <algorithm>
18 : #include <vector>
19 :
20 : #include "gmock/gmock.h"
21 : #include "gtest/gtest.h"
22 : #include "syzygy/core/unittest_util.h"
23 : #include "syzygy/msf/msf_constants.h"
24 : #include "syzygy/msf/msf_data.h"
25 : #include "syzygy/msf/msf_reader.h"
26 : #include "syzygy/msf/unittest_util.h"
27 :
28 : namespace msf {
29 :
30 : namespace {
31 :
32 : uint32_t GetNumPages(uint32_t num_bytes) {
33 : return (num_bytes + msf::kMsfPageSize - 1) / msf::kMsfPageSize;
34 : }
35 :
36 : class TestMsfWriter : public MsfWriter {
37 : public:
38 E : TestMsfWriter() {
39 E : file_.reset(base::CreateAndOpenTemporaryFile(&path_));
40 E : EXPECT_TRUE(file_.get() != NULL);
41 E : }
42 :
43 E : ~TestMsfWriter() {
44 E : if (file_.get()) {
45 i : fclose(file_.get());
46 i : file_.reset();
47 : }
48 E : base::DeleteFile(path_, false);
49 E : }
50 :
51 E : base::ScopedFILE& file() { return file_; }
52 :
53 : using MsfWriter::AppendStream;
54 : using MsfWriter::WriteHeader;
55 :
56 : base::FilePath path_;
57 : };
58 :
59 : class TestMsfStream : public MsfStream {
60 : public:
61 E : TestMsfStream(uint32_t length, uint32_t mask)
62 E : : MsfStream(length), data_(length) {
63 E : uint32_t* data = reinterpret_cast<uint32_t*>(data_.data());
64 :
65 : // Just to make sure the data is non-repeating (so we can distinguish if it
66 : // has been correctly written or not) fill it with integers encoding their
67 : // own position in the stream.
68 E : for (size_t i = 0; i < data_.size() / sizeof(data[0]); ++i)
69 E : data[i] = i | mask;
70 E : }
71 :
72 E : bool ReadBytesAt(size_t pos, size_t count, void* dest) override {
73 E : DCHECK(dest != NULL);
74 :
75 E : if (count > length() - pos)
76 i : return false;
77 :
78 E : ::memcpy(dest, data_.data() + pos, count);
79 :
80 E : return true;
81 E : }
82 :
83 : private:
84 : std::vector<uint8_t> data_;
85 : };
86 :
87 : void EnsureMsfContentsAreIdentical(const MsfFile& msf_file,
88 : const MsfFile& msf_file_read) {
89 : ASSERT_EQ(msf_file.StreamCount(), msf_file_read.StreamCount());
90 :
91 : for (size_t i = 0; i < msf_file.StreamCount(); ++i) {
92 : MsfStream* stream = msf_file.GetStream(i).get();
93 : MsfStream* stream_read = msf_file_read.GetStream(i).get();
94 :
95 : ASSERT_TRUE(stream != NULL);
96 : ASSERT_TRUE(stream_read != NULL);
97 :
98 : ASSERT_EQ(stream->length(), stream_read->length());
99 :
100 : std::vector<uint8_t> data(stream->length());
101 : std::vector<uint8_t> data_read(stream_read->length());
102 : ASSERT_TRUE(stream->ReadBytesAt(0, stream->length(), &data.at(0)));
103 : ASSERT_TRUE(
104 : stream_read->ReadBytesAt(0, stream_read->length(), &data_read.at(0)));
105 :
106 : // We don't use ContainerEq because upon failure this generates a
107 : // ridiculously long and useless error message. We don't use memcmp because
108 : // it doesn't give any context as to where the failure occurs.
109 : for (size_t j = 0; j < data.size(); ++j)
110 : ASSERT_EQ(data[j], data_read[j]);
111 : }
112 : }
113 :
114 : } // namespace
115 :
116 : using msf::kMsfHeaderMagicString;
117 : using msf::kMsfPageSize;
118 : using msf::MsfHeader;
119 : using msf::MsfReader;
120 :
121 E : TEST(MsfWriterTest, AppendStream) {
122 E : TestMsfWriter writer;
123 :
124 E : testing::ScopedTempFile temp_file;
125 E : writer.file().reset(base::OpenFile(temp_file.path(), "wb"));
126 E : ASSERT_TRUE(writer.file().get() != NULL);
127 :
128 E : scoped_refptr<MsfStream> stream(new TestMsfStream(4 * kMsfPageSize, 0));
129 :
130 : // Test writing a stream that will force allocation of the free page map
131 : // pages.
132 E : std::vector<uint32_t> pages_written;
133 E : uint32_t page_count = 0;
134 E : EXPECT_TRUE(writer.AppendStream(stream.get(), &pages_written, &page_count));
135 E : writer.file().reset();
136 :
137 : // We expect pages_written to contain 4 pages, like the stream. However, we
138 : // expect page_count to have 2 more pages for the free page map.
139 E : uint32_t expected_pages_written[] = {0, 3, 4, 5};
140 E : EXPECT_THAT(pages_written,
141 E : ::testing::ElementsAreArray(expected_pages_written));
142 E : EXPECT_EQ(page_count, 6);
143 :
144 : // Build the expected stream contents. Two blank pages should have been
145 : // reserved by the append stream routine.
146 E : std::vector<uint8_t> expected_contents(6 * kMsfPageSize);
147 E : ASSERT_TRUE(stream->ReadBytesAt(0, kMsfPageSize, expected_contents.data()));
148 E : ASSERT_TRUE(stream->ReadBytesAt(kMsfPageSize, 3 * kMsfPageSize,
149 E : expected_contents.data() + 3 * kMsfPageSize));
150 :
151 E : std::vector<uint8_t> contents(6 * kMsfPageSize);
152 E : ASSERT_EQ(
153 : contents.size(),
154 : base::ReadFile(temp_file.path(), reinterpret_cast<char*>(contents.data()),
155 E : contents.size()));
156 :
157 E : EXPECT_THAT(contents, ::testing::ContainerEq(expected_contents));
158 E : }
159 :
160 E : TEST(MsfWriterTest, WriteHeader) {
161 E : TestMsfWriter writer;
162 :
163 E : testing::ScopedTempFile temp_file;
164 E : writer.file().reset(base::OpenFile(temp_file.path(), "wb"));
165 E : ASSERT_TRUE(writer.file().get() != NULL);
166 :
167 E : std::vector<uint32_t> root_directory_pages(kMsfMaxDirPages + 10, 1);
168 :
169 : // Try to write a root directorty that's too big and expect this to fail.
170 E : EXPECT_FALSE(writer.WriteHeader(root_directory_pages, 67 * 4, 438));
171 :
172 : // Now write a reasonable root directory size.
173 E : root_directory_pages.resize(1);
174 E : EXPECT_TRUE(writer.WriteHeader(root_directory_pages, 67 * 4, 438));
175 E : writer.file().reset();
176 :
177 : // Build the expected stream contents. Two blank pages should have been
178 : // reserved by the append stream routine.
179 E : std::vector<uint8_t> expected_contents(sizeof(MsfHeader));
180 E : MsfHeader* header = reinterpret_cast<MsfHeader*>(expected_contents.data());
181 E : ::memcpy(header->magic_string, kMsfHeaderMagicString,
182 : kMsfHeaderMagicStringSize);
183 E : header->page_size = kMsfPageSize;
184 E : header->free_page_map = 1;
185 E : header->num_pages = 438;
186 E : header->directory_size = 67 * 4;
187 E : header->root_pages[0] = 1;
188 :
189 E : std::vector<uint8_t> contents(sizeof(MsfHeader));
190 E : ASSERT_EQ(
191 : contents.size(),
192 : base::ReadFile(temp_file.path(), reinterpret_cast<char*>(contents.data()),
193 E : contents.size()));
194 :
195 E : EXPECT_THAT(contents, ::testing::ContainerEq(expected_contents));
196 E : }
197 :
198 E : TEST(MsfWriterTest, WriteMsfFile) {
199 E : MsfFile msf_file;
200 E : for (uint32_t i = 0; i < 4; ++i)
201 E : msf_file.AppendStream(new TestMsfStream(1 << (8 + i), (i << 24)));
202 :
203 : // Test that we can create an MSF file and then read it successfully.
204 E : testing::ScopedTempFile file;
205 : {
206 : // Create a scope so that the file gets closed.
207 E : TestMsfWriter writer;
208 E : EXPECT_TRUE(writer.Write(file.path(), msf_file));
209 E : }
210 :
211 E : MsfFile msf_file_read;
212 E : MsfReader reader;
213 E : EXPECT_TRUE(reader.Read(file.path(), &msf_file_read));
214 :
215 E : ASSERT_NO_FATAL_FAILURE(
216 E : testing::EnsureMsfContentsAreIdentical(msf_file, msf_file_read));
217 E : }
218 :
219 : } // namespace msf
|