1 // Copyright 2022 The Pigweed Authors
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License"); you may not
4 // use this file except in compliance with the License. You may obtain a copy of
5 // the License at
6 //
7 // https://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, WITHOUT
11 // WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
12 // License for the specific language governing permissions and limitations under
13 // the License.
14 #include "pw_transfer/atomic_file_transfer_handler.h"
15
16 #include <cinttypes>
17 #include <filesystem>
18 #include <fstream>
19 #include <random>
20 #include <string>
21 #include <string_view>
22
23 #include "pw_random/xor_shift.h"
24 #include "pw_result/result.h"
25 #include "pw_status/status.h"
26 #include "pw_string/string_builder.h"
27 #include "pw_transfer/transfer.h"
28 #include "pw_transfer_private/filename_generator.h"
29 #include "pw_unit_test/framework.h"
30
31 namespace pw::transfer {
32
33 namespace {
34
35 // Copied from go/pw-src/+/main:pw_stream/std_file_stream_test.cc;l=75
36 class TempDir {
37 public:
TempDir(std::string_view prefix)38 TempDir(std::string_view prefix) : rng_(GetSeed()) {
39 temp_dir_ = std::filesystem::temp_directory_path();
40 temp_dir_ /= std::string(prefix) + GetRandomSuffix();
41 PW_ASSERT(std::filesystem::create_directory(temp_dir_));
42 }
43
~TempDir()44 ~TempDir() { PW_ASSERT(std::filesystem::remove_all(temp_dir_)); }
45
GetTempFileName()46 std::filesystem::path GetTempFileName() {
47 return temp_dir_ / GetRandomSuffix();
48 }
49
50 private:
GetRandomSuffix()51 std::string GetRandomSuffix() {
52 StringBuffer<9> random_suffix_str;
53 uint32_t random_suffix_int = 0;
54 rng_.GetInt(random_suffix_int);
55 PW_ASSERT(random_suffix_str.Format("%08" PRIx32, random_suffix_int).ok());
56 return std::string(random_suffix_str.view());
57 }
58
59 // Generate a 64-bit random from system entropy pool. This is used to seed a
60 // pseudo-random number generator for individual file names.
GetSeed()61 static uint64_t GetSeed() {
62 std::random_device sys_rand;
63 uint64_t seed = 0;
64 for (size_t seed_bytes = 0; seed_bytes < sizeof(seed);
65 seed_bytes += sizeof(std::random_device::result_type)) {
66 std::random_device::result_type val = sys_rand();
67 seed = seed << 8 * sizeof(std::random_device::result_type);
68 seed |= val;
69 }
70 return seed;
71 }
72
73 random::XorShiftStarRng64 rng_;
74 std::filesystem::path temp_dir_;
75 };
76
77 class AtomicFileTransferHandlerTest : public ::testing::Test {
78 public:
79 TempDir temp_dir_{"atomic_file_transfer_handler_test"};
80 std::string test_data_location_pass_ = temp_dir_.GetTempFileName();
81 std::string transfer_temp_file_ = GetTempFilePath(test_data_location_pass_);
82
83 protected:
84 static constexpr auto test_data_location_fail = "not/a/directory/no_data.txt";
85 static constexpr auto temp_file_content = "Temp File Success.";
86 static constexpr auto test_data_content = "Test File Success.";
87
WriteContentFile(std::string path,std::string value)88 bool WriteContentFile(std::string path, std::string value) {
89 std::ofstream file(path);
90 if (!file.is_open()) {
91 return false;
92 }
93 file << value;
94 return true;
95 }
96
ReadFile(std::string path)97 Result<std::string> ReadFile(std::string path) {
98 std::ifstream file(path);
99 if (!file.is_open()) {
100 return Status::NotFound();
101 }
102 std::string return_value;
103 std::getline(file, return_value);
104 return return_value;
105 }
106
ClearContent(std::string path)107 void ClearContent(std::string path) {
108 std::ofstream ofs(path, std::ofstream::out | std::ofstream::trunc);
109 }
110
check_finalize(Status status)111 void check_finalize(Status status) {
112 EXPECT_EQ(status, OkStatus());
113 // Temp file does not exist after finalize.
114 EXPECT_TRUE(!std::filesystem::exists(transfer_temp_file_));
115 // Test path does exist, file has been created.
116 EXPECT_TRUE(std::filesystem::exists(test_data_location_pass_));
117 // File content is the same as expected.
118 const auto file_content = ReadFile(test_data_location_pass_);
119 ASSERT_TRUE(file_content.ok());
120
121 EXPECT_EQ(file_content.value(), temp_file_content);
122 }
123
SetUp()124 void SetUp() override {
125 // Write content file and check correct.
126 ASSERT_TRUE(WriteContentFile(test_data_location_pass_, test_data_content));
127 const auto file_content_data = ReadFile(test_data_location_pass_);
128 ASSERT_TRUE(file_content_data.ok());
129 ASSERT_EQ(file_content_data.value(), test_data_content);
130
131 // Write temp file and check content is correct
132 ASSERT_TRUE(WriteContentFile(transfer_temp_file_, temp_file_content));
133 const auto file_content_tmp = ReadFile(transfer_temp_file_);
134 ASSERT_TRUE(file_content_tmp.ok());
135 ASSERT_EQ(file_content_tmp.value(), temp_file_content);
136 }
137
TearDown()138 void TearDown() override {
139 // Ensure temp file is deleted.
140 ASSERT_TRUE(!std::filesystem::exists(transfer_temp_file_) ||
141 std::filesystem::remove(transfer_temp_file_));
142 // Ensure test file is deleted.
143 ASSERT_TRUE(!std::filesystem::exists(test_data_location_pass_) ||
144 std::filesystem::remove(test_data_location_pass_));
145 }
146 };
147
TEST_F(AtomicFileTransferHandlerTest,PrepareReadPass)148 TEST_F(AtomicFileTransferHandlerTest, PrepareReadPass) {
149 AtomicFileTransferHandler test_handler{/*resource_id = */ 0,
150 test_data_location_pass_};
151 EXPECT_EQ(test_handler.PrepareRead(), OkStatus());
152 }
153
TEST_F(AtomicFileTransferHandlerTest,PrepareReadFail)154 TEST_F(AtomicFileTransferHandlerTest, PrepareReadFail) {
155 AtomicFileTransferHandler test_handler{/*resource_id = */ 0,
156 test_data_location_fail};
157 EXPECT_EQ(test_handler.PrepareRead(), Status::NotFound());
158 }
159
TEST_F(AtomicFileTransferHandlerTest,PrepareWritePass)160 TEST_F(AtomicFileTransferHandlerTest, PrepareWritePass) {
161 AtomicFileTransferHandler test_handler{/*resource_id = */ 0,
162 test_data_location_pass_};
163 // Open a file for write returns OkStatus.
164 EXPECT_EQ(test_handler.PrepareWrite(), OkStatus());
165 }
166
TEST_F(AtomicFileTransferHandlerTest,PrepareWriteFail)167 TEST_F(AtomicFileTransferHandlerTest, PrepareWriteFail) {
168 AtomicFileTransferHandler test_handler{/*resource_id = */ 0,
169 test_data_location_fail};
170 // Open a file with non existing path pass.
171 // No access to underlying stream
172 // so rely on the write during transfer to catch the error.
173 EXPECT_EQ(test_handler.PrepareWrite(), OkStatus());
174 }
175
TEST_F(AtomicFileTransferHandlerTest,FinalizeWriteRenameExisting)176 TEST_F(AtomicFileTransferHandlerTest, FinalizeWriteRenameExisting) {
177 ASSERT_TRUE(std::filesystem::exists(transfer_temp_file_));
178 ASSERT_TRUE(std::filesystem::exists(test_data_location_pass_));
179 AtomicFileTransferHandler test_handler{/*resource_id = */
180 0,
181 test_data_location_pass_};
182 // Prepare Write to open the stream. should be closed during Finalize.
183 ASSERT_EQ(test_handler.PrepareWrite(), OkStatus());
184 WriteContentFile(transfer_temp_file_, temp_file_content);
185 auto status = test_handler.FinalizeWrite(OkStatus());
186 check_finalize(status);
187 }
188
TEST_F(AtomicFileTransferHandlerTest,FinalizeWriteNoExistingFile)189 TEST_F(AtomicFileTransferHandlerTest, FinalizeWriteNoExistingFile) {
190 AtomicFileTransferHandler test_handler{/*resource_id = */
191 0,
192 test_data_location_pass_};
193 // Remove file test file and test creation.
194 ASSERT_TRUE(std::filesystem::remove(test_data_location_pass_));
195 ASSERT_EQ(test_handler.PrepareWrite(), OkStatus());
196 WriteContentFile(transfer_temp_file_, temp_file_content);
197 auto status = test_handler.FinalizeWrite(OkStatus());
198 check_finalize(status);
199 }
200
TEST_F(AtomicFileTransferHandlerTest,FinalizeWriteExpectErr)201 TEST_F(AtomicFileTransferHandlerTest, FinalizeWriteExpectErr) {
202 AtomicFileTransferHandler test_handler{/*resource_id = */
203 0,
204 test_data_location_pass_};
205 ASSERT_EQ(test_handler.PrepareWrite(), OkStatus());
206 // Simulate write fails, file is empty, No write here.
207 ClearContent(transfer_temp_file_);
208 ASSERT_TRUE(std::filesystem::is_empty(transfer_temp_file_));
209 ASSERT_TRUE(std::filesystem::exists(test_data_location_pass_));
210 EXPECT_EQ(test_handler.FinalizeWrite(Status::DataLoss()), Status::DataLoss());
211 }
212
213 } // namespace
214
215 } // namespace pw::transfer
216