Test creation of RequestContext inside a fiber task
[folly.git] / folly / experimental / TestUtil.cpp
1 /*
2  * Copyright 2016 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/types.h>
20 #include <sys/stat.h>
21 #include <fcntl.h>
22 #include <unistd.h>
23
24 #include <boost/regex.hpp>
25 #include <folly/Conv.h>
26 #include <folly/Exception.h>
27 #include <folly/File.h>
28 #include <folly/FileUtil.h>
29 #include <folly/String.h>
30 #include <folly/portability/Environment.h>
31
32 namespace folly {
33 namespace test {
34
35 namespace {
36
37 fs::path generateUniquePath(fs::path path, StringPiece namePrefix) {
38   if (path.empty()) {
39     path = fs::temp_directory_path();
40   }
41   if (namePrefix.empty()) {
42     path /= fs::unique_path();
43   } else {
44     path /= fs::unique_path(
45         to<std::string>(namePrefix, ".%%%%-%%%%-%%%%-%%%%"));
46   }
47   return path;
48 }
49
50 }  // namespace
51
52 TemporaryFile::TemporaryFile(StringPiece namePrefix,
53                              fs::path dir,
54                              Scope scope,
55                              bool closeOnDestruction)
56   : scope_(scope),
57     closeOnDestruction_(closeOnDestruction),
58     fd_(-1),
59     path_(generateUniquePath(std::move(dir), namePrefix)) {
60   fd_ = open(path_.string().c_str(), O_RDWR | O_CREAT | O_EXCL, 0666);
61   checkUnixError(fd_, "open failed");
62
63   if (scope_ == Scope::UNLINK_IMMEDIATELY) {
64     boost::system::error_code ec;
65     fs::remove(path_, ec);
66     if (ec) {
67       LOG(WARNING) << "unlink on construction failed: " << ec;
68     } else {
69       path_.clear();
70     }
71   }
72 }
73
74 const fs::path& TemporaryFile::path() const {
75   CHECK(scope_ != Scope::UNLINK_IMMEDIATELY);
76   DCHECK(!path_.empty());
77   return path_;
78 }
79
80 TemporaryFile::~TemporaryFile() {
81   if (fd_ != -1 && closeOnDestruction_) {
82     if (close(fd_) == -1) {
83       PLOG(ERROR) << "close failed";
84     }
85   }
86
87   // If we previously failed to unlink() (UNLINK_IMMEDIATELY), we'll
88   // try again here.
89   if (scope_ != Scope::PERMANENT && !path_.empty()) {
90     boost::system::error_code ec;
91     fs::remove(path_, ec);
92     if (ec) {
93       LOG(WARNING) << "unlink on destruction failed: " << ec;
94     }
95   }
96 }
97
98 TemporaryDirectory::TemporaryDirectory(StringPiece namePrefix,
99                                        fs::path dir,
100                                        Scope scope)
101   : scope_(scope),
102     path_(generateUniquePath(std::move(dir), namePrefix)) {
103   fs::create_directory(path_);
104 }
105
106 TemporaryDirectory::~TemporaryDirectory() {
107   if (scope_ == Scope::DELETE_ON_DESTRUCTION) {
108     boost::system::error_code ec;
109     fs::remove_all(path_, ec);
110     if (ec) {
111       LOG(WARNING) << "recursive delete on destruction failed: " << ec;
112     }
113   }
114 }
115
116 ChangeToTempDir::ChangeToTempDir() : initialPath_(fs::current_path()) {
117   std::string p = dir_.path().string();
118   ::chdir(p.c_str());
119 }
120
121 ChangeToTempDir::~ChangeToTempDir() {
122   std::string p = initialPath_.string();
123   ::chdir(p.c_str());
124 }
125
126 namespace detail {
127
128 bool hasPCREPatternMatch(StringPiece pattern, StringPiece target) {
129   return boost::regex_match(
130     target.begin(),
131     target.end(),
132     boost::regex(pattern.begin(), pattern.end())
133   );
134 }
135
136 bool hasNoPCREPatternMatch(StringPiece pattern, StringPiece target) {
137   return !hasPCREPatternMatch(pattern, target);
138 }
139
140 }  // namespace detail
141
142 CaptureFD::CaptureFD(int fd, ChunkCob chunk_cob)
143     : chunkCob_(std::move(chunk_cob)), fd_(fd), readOffset_(0) {
144   oldFDCopy_ = dup(fd_);
145   PCHECK(oldFDCopy_ != -1) << "Could not copy FD " << fd_;
146
147   int file_fd = open(file_.path().string().c_str(), O_WRONLY|O_CREAT, 0600);
148   PCHECK(dup2(file_fd, fd_) != -1) << "Could not replace FD " << fd_
149     << " with " << file_fd;
150   PCHECK(close(file_fd) != -1) << "Could not close " << file_fd;
151 }
152
153 void CaptureFD::release() {
154   if (oldFDCopy_ != fd_) {
155     readIncremental();  // Feed chunkCob_
156     PCHECK(dup2(oldFDCopy_, fd_) != -1) << "Could not restore old FD "
157       << oldFDCopy_ << " into " << fd_;
158     PCHECK(close(oldFDCopy_) != -1) << "Could not close " << oldFDCopy_;
159     oldFDCopy_ = fd_;  // Make this call idempotent
160   }
161 }
162
163 CaptureFD::~CaptureFD() {
164   release();
165 }
166
167 std::string CaptureFD::read() const {
168   std::string contents;
169   std::string filename = file_.path().string();
170   PCHECK(folly::readFile(filename.c_str(), contents));
171   return contents;
172 }
173
174 std::string CaptureFD::readIncremental() {
175   std::string filename = file_.path().string();
176   // Yes, I know that I could just keep the file open instead. So sue me.
177   folly::File f(openNoInt(filename.c_str(), O_RDONLY), true);
178   auto size = lseek(f.fd(), 0, SEEK_END) - readOffset_;
179   std::unique_ptr<char[]> buf(new char[size]);
180   auto bytes_read = folly::preadFull(f.fd(), buf.get(), size, readOffset_);
181   PCHECK(size == bytes_read);
182   readOffset_ += size;
183   chunkCob_(StringPiece(buf.get(), buf.get() + size));
184   return std::string(buf.get(), size);
185 }
186
187 static std::map<std::string, std::string> getEnvVarMap() {
188   std::map<std::string, std::string> data;
189   for (auto it = environ; *it != nullptr; ++it) {
190     std::string key, value;
191     split("=", *it, key, value);
192     if (key.empty()) {
193       continue;
194     }
195     CHECK(!data.count(key)) << "already contains: " << key;
196     data.emplace(move(key), move(value));
197   }
198   return data;
199 }
200
201 EnvVarSaver::EnvVarSaver() {
202   saved_ = getEnvVarMap();
203 }
204
205 EnvVarSaver::~EnvVarSaver() {
206   for (const auto& kvp : getEnvVarMap()) {
207     if (saved_.count(kvp.first)) {
208       continue;
209     }
210     PCHECK(0 == unsetenv(kvp.first.c_str()));
211   }
212   for (const auto& kvp : saved_) {
213     PCHECK(0 == setenv(kvp.first.c_str(), kvp.second.c_str(), (int)true));
214   }
215 }
216
217 }  // namespace test
218 }  // namespace folly