folly/experimental/TestUtil.cpp

   1 /*
   2  * Copyright 2015 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
  30 namespace folly {
  31 namespace test {
  32
  33 namespace {
  34
  35 fs::path generateUniquePath(fs::path path, StringPiece namePrefix) {
  36   if (path.empty()) {
  37     path = fs::temp_directory_path();
  38   }
  39   if (namePrefix.empty()) {
  40     path /= fs::unique_path();
  41   } else {
  42     path /= fs::unique_path(
  43         to<std::string>(namePrefix, ".%%%%-%%%%-%%%%-%%%%"));
  44   }
  45   return path;
  46 }
  47
  48 }  // namespace
  49
  50 TemporaryFile::TemporaryFile(StringPiece namePrefix,
  51                              fs::path dir,
  52                              Scope scope,
  53                              bool closeOnDestruction)
  54   : scope_(scope),
  55     closeOnDestruction_(closeOnDestruction),
  56     fd_(-1),
  57     path_(generateUniquePath(std::move(dir), namePrefix)) {
  58   fd_ = open(path_.string().c_str(), O_RDWR | O_CREAT | O_EXCL, 0666);
  59   checkUnixError(fd_, "open failed");
  60
  61   if (scope_ == Scope::UNLINK_IMMEDIATELY) {
  62     boost::system::error_code ec;
  63     fs::remove(path_, ec);
  64     if (ec) {
  65       LOG(WARNING) << "unlink on construction failed: " << ec;
  66     } else {
  67       path_.clear();
  68     }
  69   }
  70 }
  71
  72 const fs::path& TemporaryFile::path() const {
  73   CHECK(scope_ != Scope::UNLINK_IMMEDIATELY);
  74   DCHECK(!path_.empty());
  75   return path_;
  76 }
  77
  78 TemporaryFile::~TemporaryFile() {
  79   if (fd_ != -1 && closeOnDestruction_) {
  80     if (close(fd_) == -1) {
  81       PLOG(ERROR) << "close failed";
  82     }
  83   }
  84
  85   // If we previously failed to unlink() (UNLINK_IMMEDIATELY), we'll
  86   // try again here.
  87   if (scope_ != Scope::PERMANENT && !path_.empty()) {
  88     boost::system::error_code ec;
  89     fs::remove(path_, ec);
  90     if (ec) {
  91       LOG(WARNING) << "unlink on destruction failed: " << ec;
  92     }
  93   }
  94 }
  95
  96 TemporaryDirectory::TemporaryDirectory(StringPiece namePrefix,
  97                                        fs::path dir,
  98                                        Scope scope)
  99   : scope_(scope),
 100     path_(generateUniquePath(std::move(dir), namePrefix)) {
 101   fs::create_directory(path_);
 102 }
 103
 104 TemporaryDirectory::~TemporaryDirectory() {
 105   if (scope_ == Scope::DELETE_ON_DESTRUCTION) {
 106     boost::system::error_code ec;
 107     fs::remove_all(path_, ec);
 108     if (ec) {
 109       LOG(WARNING) << "recursive delete on destruction failed: " << ec;
 110     }
 111   }
 112 }
 113
 114 ChangeToTempDir::ChangeToTempDir() : initialPath_(fs::current_path()) {
 115   std::string p = dir_.path().string();
 116   ::chdir(p.c_str());
 117 }
 118
 119 ChangeToTempDir::~ChangeToTempDir() {
 120   std::string p = initialPath_.string();
 121   ::chdir(p.c_str());
 122 }
 123
 124 namespace detail {
 125
 126 bool hasPCREPatternMatch(StringPiece pattern, StringPiece target) {
 127   return boost::regex_match(
 128     target.begin(),
 129     target.end(),
 130     boost::regex(pattern.begin(), pattern.end())
 131   );
 132 }
 133
 134 bool hasNoPCREPatternMatch(StringPiece pattern, StringPiece target) {
 135   return !hasPCREPatternMatch(pattern, target);
 136 }
 137
 138 }  // namespace detail
 139
 140 CaptureFD::CaptureFD(int fd) : fd_(fd), readOffset_(0) {
 141   oldFDCopy_ = dup(fd_);
 142   PCHECK(oldFDCopy_ != -1) << "Could not copy FD " << fd_;
 143
 144   int file_fd = open(file_.path().string().c_str(), O_WRONLY|O_CREAT, 0600);
 145   PCHECK(dup2(file_fd, fd_) != -1) << "Could not replace FD " << fd_
 146     << " with " << file_fd;
 147   PCHECK(close(file_fd) != -1) << "Could not close " << file_fd;
 148 }
 149
 150 void CaptureFD::release() {
 151   if (oldFDCopy_ != fd_) {
 152     PCHECK(dup2(oldFDCopy_, fd_) != -1) << "Could not restore old FD "
 153       << oldFDCopy_ << " into " << fd_;
 154     PCHECK(close(oldFDCopy_) != -1) << "Could not close " << oldFDCopy_;
 155     oldFDCopy_ = fd_;  // Make this call idempotent
 156   }
 157 }
 158
 159 CaptureFD::~CaptureFD() {
 160   release();
 161 }
 162
 163 std::string CaptureFD::read() {
 164   std::string contents;
 165   std::string filename = file_.path().string();
 166   PCHECK(folly::readFile(filename.c_str(), contents));
 167   return contents;
 168 }
 169
 170 std::string CaptureFD::readIncremental() {
 171   std::string filename = file_.path().string();
 172   // Yes, I know that I could just keep the file open instead. So sue me.
 173   folly::File f(openNoInt(filename.c_str(), O_RDONLY), true);
 174   auto size = lseek(f.fd(), 0, SEEK_END) - readOffset_;
 175   std::unique_ptr<char[]> buf(new char[size]);
 176   auto bytes_read = folly::preadFull(f.fd(), buf.get(), size, readOffset_);
 177   PCHECK(size == bytes_read);
 178   readOffset_ += size;
 179   return std::string(buf.get(), size);
 180 }
 181
 182 }  // namespace test
 183 }  // namespace folly