src/process_utils.cc

// Copyright 2023 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "process_utils.h"

#include <assert.h>

#include <algorithm>

#include "ipc_utils.h"
#include "util.h"

///////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////
/////
/////  ScopedEnvironmentVariable
/////

ScopedEnvironmentVariable::ScopedEnvironmentVariable(const char* varname,
                                                     const char* value)
    : varname_(varname) {
  const char* env = getenv(varname);
  if (env) {
    has_prev_value_ = true;
    prev_value_ = env;
  }
#ifdef _WIN32
  // On Windows, setting the value to the empty string removes the variable.
  // The following check is to ensure that UnsetEnv() is always called when
  // trying to set an undefined |varname| to an empty |value|.
  if (value && value[0]) {
#else
  if (value) {
#endif
    if (SetEnv(varname_.c_str(), value) < 0)
      ErrnoFatal("setenv");
  } else {
    UnsetEnv(varname_.c_str());
  }
}

ScopedEnvironmentVariable::~ScopedEnvironmentVariable() {
  if (has_prev_value_) {
    SetEnv(varname_.c_str(), prev_value_.c_str());
  } else {
    UnsetEnv(varname_.c_str());
  }
}

ScopedEnvironmentVariable::ScopedEnvironmentVariable(
    ScopedEnvironmentVariable&&) noexcept = default;
ScopedEnvironmentVariable& ScopedEnvironmentVariable::operator=(
    ScopedEnvironmentVariable&&) noexcept = default;

int ScopedEnvironmentVariable::SetEnv(const char* varname, const char* value) {
#ifdef _WIN32
  // There is no setenv() on Win32!?
  // _putenv_s() returns EINVAL in case of error, but also
  // sets errno, so convert errors to -1 here.
  return (_putenv_s(varname, value) == 0) ? 0 : -1;
#else   // !_WIN32
  return ::setenv(varname, value, 1);
#endif  // !_WIN32
}

void ScopedEnvironmentVariable::UnsetEnv(const char* varname) {
#ifdef _WIN32
  _putenv_s(varname, "");
#else   // !_WIN32
  ::unsetenv(varname);
#endif  // !_WIN32
}

///////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////
/////
/////  EnvironmentBlock
/////

EnvironmentBlock::EnvironmentBlock() = default;

extern char** environ;

// static
EnvironmentBlock EnvironmentBlock::CreateFromCurrentEnvironment() {
  EnvironmentBlock result;

  for (const char* const* env = environ; *env; ++env) {
    std::string definition = *env;
    size_t pos = definition.find('=');
    if (pos == std::string::npos) {
      // Should not happen, ignore it.
      continue;
    }

    std::string varname = definition.substr(0, pos);
    std::string varvalue = definition.substr(pos + 1);

#ifdef _WIN32
    // Empty values are not possible on Win32, so skip them on this platform
    // if we ever find them in the environ block.
    if (varvalue.empty())
      continue;
#endif
    result.map_[varname] = varvalue;
  }
  return result;
}

void EnvironmentBlock::MergeWith(const EnvironmentBlock& other) {
  for (const auto& pair : other.map_)
    map_.emplace(pair.first, pair.second);
}

void EnvironmentBlock::Insert(const char* varname, const char* value) {
  if (!value)
    value = "";
  map_[std::string(varname)] = value;
}

void EnvironmentBlock::Remove(const char* varname) {
  map_.erase(std::string(varname));
}

const char* EnvironmentBlock::Get(const char* varname) const {
  auto it = map_.find(std::string(varname));
  if (it == map_.end())
    return nullptr;
  return it->second.c_str();
}

std::string EnvironmentBlock::ToEncodedString() const {
  auto definitions = GetDefinitions();
  WireEncoder encoder;
  encoder.Write(map_.size());
  for (const auto& pair : map_) {
    encoder.Write(pair.first);
    encoder.Write(pair.second);
  }
  return encoder.TakeResult();
}

// static
EnvironmentBlock EnvironmentBlock::FromEncodedString(const std::string& encoded,
                                                     std::string* error) {
  error->clear();
  EnvironmentBlock::MapType map;
  WireDecoder decoder(encoded);
  size_t count = 0;
  decoder.Read(count);
  for (; count > 0; --count) {
    std::string varname, varvalue;
    decoder.Read(varname);
    decoder.Read(varvalue);
    map.emplace(varname, std::move(varvalue));
  }
  EnvironmentBlock result;
  if (!decoder.has_error()) {
    result.map_ = std::move(map);
  } else {
    *error = "Truncated encoded EnvironmentBlock string";
  }
  return result;
}

std::string EnvironmentBlock::AsString() const {
  auto defs = GetDefinitions();
  std::string result = StringFormat("Environment(%zu)[\n", defs.list.size());
  for (const auto& def : defs.list)
    result += StringFormat("  %s\n", def.c_str());
  result += "]";
  return result;
}

bool EnvironmentBlock::operator==(const EnvironmentBlock& other) const {
  auto defs = GetDefinitions();
  auto other_defs = other.GetDefinitions();
  if (defs.total_size != other_defs.total_size)
    return false;
  if (defs.list.size() != other_defs.list.size())
    return false;
  auto other_it = other_defs.list.begin();
  for (const auto& def : defs.list) {
    if (def != *other_it)
      return false;
    ++other_it;
  }
  return true;
}

EnvironmentBlock::Definitions EnvironmentBlock::GetDefinitions() const {
  Definitions result;
  for (const auto& pair : map_) {
    result.list.push_back(pair.first + "=" + pair.second);
    result.total_size += result.list.back().size() + 1;
  }
  std::sort(result.list.begin(), result.list.end());
  return result;
}

#ifdef _WIN32
char* EnvironmentBlock::AsAnsiEnvironmentBlock() const {
  auto definitions = GetDefinitions();

  block_.clear();
  block_.reserve(definitions.total_size + 1);
  for (const auto& def : definitions.list) {
    block_.append(def.c_str(), def.size() + 1);
  }
  block_ += '\0';

  return &block_[0];
}

wchar_t* EnvironmentBlock::AsUnicodeEnvironmentBlock() const {
  auto definitions = GetDefinitions();

  block_.clear();
  block_.reserve(definitions.total_size * 2 + 2);
  for (const auto& def : definitions.list) {
    std::wstring wdef = ConvertToUnicodeString(def);
    block_.append(reinterpret_cast<const char*>(wdef.c_str()),
                  wdef.size() * 2 + 2);
  }
  block_ += '\0';
  block_ += '\0';

  return reinterpret_cast<wchar_t*>(&block_[0]);
}

#else   // !_WIN32
char** EnvironmentBlock::AsExecEnvironmentBlock() const {
  auto definitions = GetDefinitions();
  block_.clear();
  block_.reserve(definitions.total_size);

  env_.clear();
  env_.reserve(definitions.list.size() + 1);
  for (const auto& def : definitions.list) {
    char* dst = const_cast<char*>(block_.data()) + block_.size();
    env_.push_back(dst);
    block_.append(def.c_str(), def.size() + 1);
  }
  assert(block_.size() == definitions.total_size);
  env_.push_back(nullptr);
  return &env_.front();
}
#endif  // !_WIN32

// static
std::vector<StringPiece> SplitCommaOrSpaceSeparatedList(StringPiece input) {
  std::vector<StringPiece> result;
  const char* p = input.str_;
  const char* end = p + input.len_;

  while (p < end) {
    // Find next variable name length.
    const char* q = p;
    while (q < end && *q != ' ' && *q != ',')
      q++;

    size_t len = q - p;
    if (len > 0)
      result.emplace_back(p, len);
    if (q == end)
      break;

    p = q + 1;
  }
  return result;
}