folly/experimental/EnvUtil.cpp

   1 /*
   2  * Copyright 2017 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/EnvUtil.h>
  18
  19 #include <folly/String.h>
  20 #include <folly/portability/Stdlib.h>
  21 #include <folly/portability/Unistd.h>
  22
  23 using namespace folly;
  24 using namespace folly::experimental;
  25
  26 EnvironmentState EnvironmentState::fromCurrentEnvironment() {
  27   std::unordered_map<std::string, std::string> data;
  28   for (auto it = environ; it && *it; ++it) {
  29     std::string key, value;
  30     folly::StringPiece entry(*it);
  31     auto equalsPosition = entry.find('=');
  32     if (equalsPosition == entry.npos) {
  33       throw MalformedEnvironment{to<std::string>(
  34           "Environment contains an non key-value-pair string \"", entry, "\"")};
  35     }
  36     key = entry.subpiece(0, equalsPosition).toString();
  37     value = entry.subpiece(equalsPosition + 1).toString();
  38     if (data.count(key)) {
  39       throw MalformedEnvironment{to<std::string>(
  40           "Environment contains duplicate value for \"", key, "\"")};
  41     }
  42     data.emplace(std::move(key), std::move(value));
  43   }
  44   return EnvironmentState{std::move(data)};
  45 }
  46
  47 void EnvironmentState::setAsCurrentEnvironment() {
  48   PCHECK(0 == clearenv());
  49   for (const auto& kvp : env_) {
  50     PCHECK(0 == setenv(kvp.first.c_str(), kvp.second.c_str(), (int)true));
  51   }
  52 }
  53
  54 std::vector<std::string> EnvironmentState::toVector() const {
  55   std::vector<std::string> result;
  56   for (auto const& pair : env_) {
  57     result.emplace_back(to<std::string>(pair.first, "=", pair.second));
  58   }
  59   return result;
  60 }
  61
  62 std::unique_ptr<char*, void (*)(char**)> EnvironmentState::toPointerArray()
  63     const {
  64   size_t totalStringLength{};
  65   for (auto const& pair : env_) {
  66     totalStringLength += pair.first.size() + pair.second.size() +
  67         2 /* intermediate '=' and the terminating NUL */;
  68   }
  69   size_t allocationRequired =
  70       (totalStringLength / sizeof(char*) + 1) + env_.size() + 1;
  71   char** raw = new char*[allocationRequired];
  72   char** ptrBase = raw;
  73   char* stringBase = reinterpret_cast<char*>(&raw[env_.size() + 1]);
  74   char* const stringEnd = reinterpret_cast<char*>(&raw[allocationRequired]);
  75   for (auto const& pair : env_) {
  76     std::string const& key = pair.first;
  77     std::string const& value = pair.second;
  78     *ptrBase = stringBase;
  79     size_t lengthIncludingNullTerminator = key.size() + 1 + value.size() + 1;
  80     CHECK_GT(stringEnd - lengthIncludingNullTerminator, stringBase);
  81     memcpy(stringBase, key.c_str(), key.size());
  82     stringBase += key.size();
  83     *stringBase++ = '=';
  84     memcpy(stringBase, value.c_str(), value.size() + 1);
  85     stringBase += value.size() + 1;
  86     ++ptrBase;
  87   }
  88   *ptrBase = nullptr;
  89   CHECK_EQ(env_.size(), ptrBase - raw);
  90   return {raw, [](char** ptr) { delete[] ptr; }};
  91 }