1 // Copyright 2018 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 ///////////////////////////////////////////////////////////////////////////////
16
17 #include "tink/util/ostream_output_stream.h"
18
19 #include <algorithm>
20 #include <cstring>
21 #include <fstream>
22 #include <iostream>
23 #include <memory>
24 #include <ostream>
25 #include <string>
26 #include <utility>
27
28 #include "gtest/gtest.h"
29 #include "absl/memory/memory.h"
30 #include "absl/strings/str_cat.h"
31 #include "absl/strings/string_view.h"
32 #include "tink/internal/test_file_util.h"
33 #include "tink/subtle/random.h"
34 #include "tink/util/test_util.h"
35
36 namespace crypto {
37 namespace tink {
38 namespace {
39
40 // Creates a new test ostream which will write to the file 'filename'.
GetTestOstream(absl::string_view filename)41 std::unique_ptr<std::ostream> GetTestOstream(absl::string_view filename) {
42 std::string full_filename =
43 absl::StrCat(crypto::tink::test::TmpDir(), "/", filename);
44 auto test_ostream = absl::make_unique<std::ofstream>(
45 full_filename, std::ofstream::binary);
46 return std::move(test_ostream);
47 }
48
49 // Writes 'contents' to the specified 'output_stream', and closes the stream.
50 // Returns the status of output_stream->Close()-operation, or a non-OK status
51 // of a prior output_stream->Next()-operation, if any.
WriteToStream(util::OstreamOutputStream * output_stream,absl::string_view contents)52 util::Status WriteToStream(util::OstreamOutputStream* output_stream,
53 absl::string_view contents) {
54 void* buffer;
55 int pos = 0;
56 int remaining = contents.length();
57 int available_space = 0;
58 int available_bytes = 0;
59 while (remaining > 0) {
60 auto next_result = output_stream->Next(&buffer);
61 if (!next_result.ok()) return next_result.status();
62 available_space = next_result.value();
63 available_bytes = std::min(available_space, remaining);
64 memcpy(buffer, contents.data() + pos, available_bytes);
65 remaining -= available_bytes;
66 pos += available_bytes;
67 }
68 if (available_space > available_bytes) {
69 output_stream->BackUp(available_space - available_bytes);
70 }
71 return output_stream->Close();
72 }
73
74 class OstreamOutputStreamTest : public ::testing::Test {
75 };
76
TEST_F(OstreamOutputStreamTest,WritingStreams)77 TEST_F(OstreamOutputStreamTest, WritingStreams) {
78 for (size_t stream_size : {0, 10, 100, 1000, 10000, 100000, 1000000}) {
79 SCOPED_TRACE(absl::StrCat("stream_size = ", stream_size));
80 std::string stream_contents = subtle::Random::GetRandomBytes(stream_size);
81 std::string filename = absl::StrCat(
82 stream_size, internal::GetTestFileNamePrefix(), "_file.bin");
83 auto output = GetTestOstream(filename);
84 auto output_stream = absl::make_unique<util::OstreamOutputStream>(
85 std::move(output));
86 auto status = WriteToStream(output_stream.get(), stream_contents);
87 EXPECT_TRUE(status.ok()) << status;
88 std::string ostream_contents = test::ReadTestFile(filename);
89 EXPECT_EQ(stream_size, ostream_contents.size());
90 EXPECT_EQ(stream_contents, ostream_contents);
91 }
92 }
93
TEST_F(OstreamOutputStreamTest,CustomBufferSizes)94 TEST_F(OstreamOutputStreamTest, CustomBufferSizes) {
95 int stream_size = 1024 * 1024;
96 std::string stream_contents = subtle::Random::GetRandomBytes(stream_size);
97 for (int buffer_size : {1, 10, 100, 1000, 10000, 100000, 1000000}) {
98 SCOPED_TRACE(absl::StrCat("buffer_size = ", buffer_size));
99 std::string filename = absl::StrCat(
100 buffer_size, internal::GetTestFileNamePrefix(), "_file.bin");
101 auto output = GetTestOstream(filename);
102 auto output_stream = absl::make_unique<util::OstreamOutputStream>(
103 std::move(output), buffer_size);
104 void* buffer;
105 auto next_result = output_stream->Next(&buffer);
106 EXPECT_TRUE(next_result.ok()) << next_result.status();
107 EXPECT_EQ(buffer_size, next_result.value());
108 output_stream->BackUp(buffer_size);
109 auto status = WriteToStream(output_stream.get(), stream_contents);
110 EXPECT_TRUE(status.ok()) << status;
111 std::string ostream_contents = test::ReadTestFile(filename);
112 EXPECT_EQ(stream_size, ostream_contents.size());
113 EXPECT_EQ(stream_contents, ostream_contents);
114 }
115 }
116
TEST_F(OstreamOutputStreamTest,BackupAndPosition)117 TEST_F(OstreamOutputStreamTest, BackupAndPosition) {
118 int stream_size = 1024 * 1024;
119 int buffer_size = 1234;
120 void* buffer;
121 std::string stream_contents = subtle::Random::GetRandomBytes(stream_size);
122 std::string filename =
123 absl::StrCat(buffer_size, internal::GetTestFileNamePrefix(), "_file.bin");
124 auto output = GetTestOstream(filename);
125
126 // Prepare the stream and do the first call to Next().
127 auto output_stream = absl::make_unique<util::OstreamOutputStream>(
128 std::move(output), buffer_size);
129 EXPECT_EQ(0, output_stream->Position());
130 auto next_result = output_stream->Next(&buffer);
131 EXPECT_TRUE(next_result.ok()) << next_result.status();
132 EXPECT_EQ(buffer_size, next_result.value());
133 EXPECT_EQ(buffer_size, output_stream->Position());
134 std::memcpy(buffer, stream_contents.data(), buffer_size);
135
136 // BackUp several times, but in total fewer bytes than returned by Next().
137 int total_backup_size = 0;
138 for (int backup_size : {0, 1, 5, 0, 10, 100, -42, 400, 20, -100}) {
139 SCOPED_TRACE(absl::StrCat("backup_size = ", backup_size));
140 output_stream->BackUp(backup_size);
141 total_backup_size += std::max(backup_size, 0);
142 EXPECT_EQ(buffer_size - total_backup_size, output_stream->Position());
143 }
144 EXPECT_LT(total_backup_size, next_result.value());
145
146 // Call Next(), it should succeed.
147 next_result = output_stream->Next(&buffer);
148 EXPECT_TRUE(next_result.ok()) << next_result.status();
149
150 // BackUp() some bytes, again fewer than returned by Next().
151 total_backup_size = 0;
152 for (int backup_size : {0, 72, -94, 37, 82}) {
153 SCOPED_TRACE(absl::StrCat("backup_size = ", backup_size));
154 output_stream->BackUp(backup_size);
155 total_backup_size += std::max(0, backup_size);
156 EXPECT_EQ(buffer_size - total_backup_size, output_stream->Position());
157 }
158 EXPECT_LT(total_backup_size, next_result.value());
159
160 // Call Next(), it should succeed;
161 next_result = output_stream->Next(&buffer);
162 EXPECT_TRUE(next_result.ok()) << next_result.status();
163
164 // Call Next() again, it should return a full block.
165 auto prev_position = output_stream->Position();
166 next_result = output_stream->Next(&buffer);
167 EXPECT_TRUE(next_result.ok()) << next_result.status();
168 EXPECT_EQ(buffer_size, next_result.value());
169 EXPECT_EQ(prev_position + buffer_size, output_stream->Position());
170 std::memcpy(buffer, stream_contents.data() + buffer_size, buffer_size);
171
172 // BackUp a few times, with total over the returned buffer_size.
173 total_backup_size = 0;
174 for (int backup_size :
175 {0, 72, -100, buffer_size / 2, 200, -25, buffer_size / 2, 42}) {
176 SCOPED_TRACE(absl::StrCat("backup_size = ", backup_size));
177 output_stream->BackUp(backup_size);
178 total_backup_size = std::min(buffer_size,
179 total_backup_size + std::max(backup_size, 0));
180 EXPECT_EQ(prev_position + buffer_size - total_backup_size,
181 output_stream->Position());
182 }
183 EXPECT_EQ(total_backup_size, buffer_size);
184 EXPECT_EQ(prev_position, output_stream->Position());
185
186 // Call Next() again, it should return a full block.
187 next_result = output_stream->Next(&buffer);
188 EXPECT_TRUE(next_result.ok()) << next_result.status();
189 EXPECT_EQ(buffer_size, next_result.value());
190 EXPECT_EQ(prev_position + buffer_size, output_stream->Position());
191 std::memcpy(buffer, stream_contents.data() + buffer_size, buffer_size);
192
193 // Write the remaining stream contents to stream.
194 auto status = WriteToStream(
195 output_stream.get(), stream_contents.substr(output_stream->Position()));
196 EXPECT_TRUE(status.ok()) << status;
197 std::string ostream_contents = test::ReadTestFile(filename);
198 EXPECT_EQ(stream_size, ostream_contents.size());
199 EXPECT_EQ(stream_contents, ostream_contents);
200 }
201
202 } // namespace
203 } // namespace tink
204 } // namespace crypto
205