1 | /* |
2 | * Copyright 2012-present Facebook, Inc. |
3 | * |
4 | * Licensed under the Apache License, Version 2.0 (the "License"); |
5 | * you may not use this file except in compliance with the License. |
6 | * You may obtain a copy of the License at |
7 | * |
8 | * http://www.apache.org/licenses/LICENSE-2.0 |
9 | * |
10 | * Unless required by applicable law or agreed to in writing, software |
11 | * distributed under the License is distributed on an "AS IS" BASIS, |
12 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
13 | * See the License for the specific language governing permissions and |
14 | * limitations under the License. |
15 | */ |
16 | |
17 | #include <folly/experimental/TestUtil.h> |
18 | |
19 | #include <sys/stat.h> |
20 | #include <sys/types.h> |
21 | |
22 | #include <boost/regex.hpp> |
23 | |
24 | #include <folly/Exception.h> |
25 | #include <folly/File.h> |
26 | #include <folly/FileUtil.h> |
27 | #include <folly/Memory.h> |
28 | #include <folly/String.h> |
29 | #include <folly/portability/Fcntl.h> |
30 | |
31 | #ifdef _WIN32 |
32 | #include <crtdbg.h> // @manual |
33 | #endif |
34 | |
35 | namespace folly { |
36 | namespace test { |
37 | |
38 | namespace { |
39 | |
40 | fs::path generateUniquePath(fs::path path, StringPiece namePrefix) { |
41 | if (path.empty()) { |
42 | path = fs::temp_directory_path(); |
43 | } |
44 | if (namePrefix.empty()) { |
45 | path /= fs::unique_path(); |
46 | } else { |
47 | path /= |
48 | fs::unique_path(to<std::string>(namePrefix, ".%%%%-%%%%-%%%%-%%%%" )); |
49 | } |
50 | return path; |
51 | } |
52 | |
53 | } // namespace |
54 | |
55 | TemporaryFile::TemporaryFile( |
56 | StringPiece namePrefix, |
57 | fs::path dir, |
58 | Scope scope, |
59 | bool closeOnDestruction) |
60 | : scope_(scope), |
61 | closeOnDestruction_(closeOnDestruction), |
62 | fd_(-1), |
63 | path_(generateUniquePath(std::move(dir), namePrefix)) { |
64 | fd_ = open(path_.string().c_str(), O_RDWR | O_CREAT | O_EXCL, 0666); |
65 | checkUnixError(fd_, "open failed" ); |
66 | |
67 | if (scope_ == Scope::UNLINK_IMMEDIATELY) { |
68 | boost::system::error_code ec; |
69 | fs::remove(path_, ec); |
70 | if (ec) { |
71 | LOG(WARNING) << "unlink on construction failed: " << ec; |
72 | } else { |
73 | path_.clear(); |
74 | } |
75 | } |
76 | } |
77 | |
78 | void TemporaryFile::close() { |
79 | if (::close(fd_) == -1) { |
80 | PLOG(ERROR) << "close failed" ; |
81 | } |
82 | fd_ = -1; |
83 | } |
84 | |
85 | const fs::path& TemporaryFile::path() const { |
86 | CHECK(scope_ != Scope::UNLINK_IMMEDIATELY); |
87 | DCHECK(!path_.empty()); |
88 | return path_; |
89 | } |
90 | |
91 | void TemporaryFile::reset() { |
92 | if (fd_ != -1 && closeOnDestruction_) { |
93 | if (::close(fd_) == -1) { |
94 | PLOG(ERROR) << "close failed (fd = " << fd_ << "): " ; |
95 | } |
96 | } |
97 | |
98 | // If we previously failed to unlink() (UNLINK_IMMEDIATELY), we'll |
99 | // try again here. |
100 | if (scope_ != Scope::PERMANENT && !path_.empty()) { |
101 | boost::system::error_code ec; |
102 | fs::remove(path_, ec); |
103 | if (ec) { |
104 | LOG(WARNING) << "unlink on destruction failed: " << ec; |
105 | } |
106 | } |
107 | } |
108 | |
109 | TemporaryFile::~TemporaryFile() { |
110 | reset(); |
111 | } |
112 | |
113 | TemporaryDirectory::TemporaryDirectory( |
114 | StringPiece namePrefix, |
115 | fs::path dir, |
116 | Scope scope) |
117 | : scope_(scope), |
118 | path_(std::make_unique<fs::path>( |
119 | generateUniquePath(std::move(dir), namePrefix))) { |
120 | fs::create_directory(path()); |
121 | } |
122 | |
123 | TemporaryDirectory::~TemporaryDirectory() { |
124 | if (scope_ == Scope::DELETE_ON_DESTRUCTION && path_ != nullptr) { |
125 | boost::system::error_code ec; |
126 | fs::remove_all(path(), ec); |
127 | if (ec) { |
128 | LOG(WARNING) << "recursive delete on destruction failed: " << ec; |
129 | } |
130 | } |
131 | } |
132 | |
133 | ChangeToTempDir::ChangeToTempDir() { |
134 | orig_ = fs::current_path(); |
135 | fs::current_path(path()); |
136 | } |
137 | |
138 | ChangeToTempDir::~ChangeToTempDir() { |
139 | if (!orig_.empty()) { |
140 | fs::current_path(orig_); |
141 | } |
142 | } |
143 | |
144 | namespace detail { |
145 | |
146 | SavedState disableInvalidParameters() { |
147 | #ifdef _WIN32 |
148 | SavedState ret; |
149 | ret.previousThreadLocalHandler = |
150 | _set_thread_local_invalid_parameter_handler([](const wchar_t*, |
151 | const wchar_t*, |
152 | const wchar_t*, |
153 | unsigned int, |
154 | uintptr_t) {}); |
155 | ret.previousCrtReportMode = _CrtSetReportMode(_CRT_ASSERT, 0); |
156 | return ret; |
157 | #else |
158 | return SavedState(); |
159 | #endif |
160 | } |
161 | |
162 | #ifdef _WIN32 |
163 | void enableInvalidParameters(SavedState state) { |
164 | _set_thread_local_invalid_parameter_handler( |
165 | (_invalid_parameter_handler)state.previousThreadLocalHandler); |
166 | _CrtSetReportMode(_CRT_ASSERT, state.previousCrtReportMode); |
167 | } |
168 | #else |
169 | void enableInvalidParameters(SavedState) {} |
170 | #endif |
171 | |
172 | bool hasPCREPatternMatch(StringPiece pattern, StringPiece target) { |
173 | return boost::regex_match( |
174 | target.begin(), |
175 | target.end(), |
176 | boost::regex(pattern.begin(), pattern.end())); |
177 | } |
178 | |
179 | bool hasNoPCREPatternMatch(StringPiece pattern, StringPiece target) { |
180 | return !hasPCREPatternMatch(pattern, target); |
181 | } |
182 | |
183 | } // namespace detail |
184 | |
185 | CaptureFD::CaptureFD(int fd, ChunkCob chunk_cob) |
186 | : chunkCob_(std::move(chunk_cob)), fd_(fd), readOffset_(0) { |
187 | oldFDCopy_ = dup(fd_); |
188 | PCHECK(oldFDCopy_ != -1) << "Could not copy FD " << fd_; |
189 | |
190 | int file_fd = open(file_.path().string().c_str(), O_WRONLY | O_CREAT, 0600); |
191 | PCHECK(dup2(file_fd, fd_) != -1) |
192 | << "Could not replace FD " << fd_ << " with " << file_fd; |
193 | PCHECK(close(file_fd) != -1) << "Could not close " << file_fd; |
194 | } |
195 | |
196 | void CaptureFD::release() { |
197 | if (oldFDCopy_ != fd_) { |
198 | readIncremental(); // Feed chunkCob_ |
199 | PCHECK(dup2(oldFDCopy_, fd_) != -1) |
200 | << "Could not restore old FD " << oldFDCopy_ << " into " << fd_; |
201 | PCHECK(close(oldFDCopy_) != -1) << "Could not close " << oldFDCopy_; |
202 | oldFDCopy_ = fd_; // Make this call idempotent |
203 | } |
204 | } |
205 | |
206 | CaptureFD::~CaptureFD() { |
207 | release(); |
208 | } |
209 | |
210 | std::string CaptureFD::read() const { |
211 | std::string contents; |
212 | std::string filename = file_.path().string(); |
213 | PCHECK(folly::readFile(filename.c_str(), contents)); |
214 | return contents; |
215 | } |
216 | |
217 | std::string CaptureFD::readIncremental() { |
218 | std::string filename = file_.path().string(); |
219 | // Yes, I know that I could just keep the file open instead. So sue me. |
220 | folly::File f(openNoInt(filename.c_str(), O_RDONLY), true); |
221 | auto size = size_t(lseek(f.fd(), 0, SEEK_END) - readOffset_); |
222 | std::unique_ptr<char[]> buf(new char[size]); |
223 | auto bytes_read = folly::preadFull(f.fd(), buf.get(), size, readOffset_); |
224 | PCHECK(ssize_t(size) == bytes_read); |
225 | readOffset_ += off_t(size); |
226 | chunkCob_(StringPiece(buf.get(), buf.get() + size)); |
227 | return std::string(buf.get(), size); |
228 | } |
229 | |
230 | } // namespace test |
231 | } // namespace folly |
232 | |