src/process_utils.h

// 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.

#ifndef NINJA_PROCESS_UTILS_H
#define NINJA_PROCESS_UTILS_H

#include <memory>
#include <string>
#include <unordered_map>
#include <vector>

#ifdef _WIN32
#include <windows.h>
#else
#include <sys/types.h>
#endif

/// Convenience class to temporarily set the value of an environment variable
/// in the current process.
class ScopedEnvironmentVariable {
 public:
  /// Constructor sets the value of a given environment variable.
  /// Set |value| to nullptr to remove the variable from the environment.
  /// Note that on Windows, setting |value| to an empty string will also
  /// remove the value, while it will set it to the empty string on Posix.
  ScopedEnvironmentVariable(const char* varname, const char* value);

  /// Destructor restores the previous value.
  ~ScopedEnvironmentVariable();

  /// Allow move operations.
  ScopedEnvironmentVariable(ScopedEnvironmentVariable&&) noexcept;
  ScopedEnvironmentVariable& operator=(ScopedEnvironmentVariable&&) noexcept;

 private:
  int SetEnv(const char* varname, const char* value);
  void UnsetEnv(const char* varname);

  std::string varname_;
  std::string prev_value_;

  // On Posix, it is possible to set an environment variable to the empty value
  // while on Windows, this simply unsets / removes the variable. This flag is
  // only used on Posix to distinguish between these two cases.
  bool has_prev_value_ = false;
};

/// A convenience wrapper for a set of environment variables.
class EnvironmentBlock {
 public:
  /// Default constructor creates an empty set.
  EnvironmentBlock();

  /// Create new instance from current process environment.
  static EnvironmentBlock CreateFromCurrentEnvironment();

  /// Merge with another instance. This only gets variables from |other|
  /// which were not set
  void MergeWith(const EnvironmentBlock& other);

  /// Add a given variable to the set, with its associated value.
  /// If the variable is already in the set, replace its value.
  void Insert(const char* varname, const char* value);

  /// Remove |varname| and its associated value from the set.
  void Remove(const char* varname);

  /// Return the value associated with |varname| in the set, or nullptr
  /// if it is not in it.
  const char* Get(const char* varname) const;

  /// Convert this instance to an encoded string, useful to send it
  /// to another process.
  std::string ToEncodedString() const;

  /// Decode a new instance from an encoded string. In case of
  /// failure, return an empty environment block, then set |*err|
  /// with an error message. On success |*err| is cleared to ease
  /// error handling on the client side.
  static EnvironmentBlock FromEncodedString(const std::string& encoded,
                                            std::string* err);

  /// Convert instance to string for debugging and testing only.
  std::string AsString() const;

  /// Compare for equality.
  bool operator==(const EnvironmentBlock& other) const;

  /// Compare for inequality.
  bool operator!=(const EnvironmentBlock& other) const {
    return !(*this == other);
  }

#ifdef _WIN32
  /// Return the current instance as a Win32 environment block, which is
  /// a sequence of zero-terminated C strings, followed by a single zero
  /// (as in `FOO=foo<0>BAR=bar<0><0>`, this is the value expected by
  /// the `lpEnvironment` parameter of CreateProcessA.
  ///
  /// NOTE: The returned address belongs to the instance, but its content
  /// may become invalid after calling any other methods.
  char* AsAnsiEnvironmentBlock() const;

  /// Same as AsWin32AnsiEnvironmentBlock(), but uses Unicode characters
  /// instead. This is the format expected by CreateProcessW.
  ///
  /// NOTE: The returned address belongs to the instance, but its content
  /// may become invalid after calling any other methods.
  wchar_t* AsUnicodeEnvironmentBlock() const;
#else   // !_WIN32
  /// A pointer to an array of zero-terminated C string pointers, followed
  /// by a NULL pointer. This is the value expected by the execve() function.
  ///
  /// NOTE: The returned address belongs to the instance, but its content
  /// may become invalid after calling any other methods.
  char** AsExecEnvironmentBlock() const;
#endif  // !_WIN32
 private:
  struct Definitions {
    std::vector<std::string> list;
    size_t total_size = 0;
  };
  Definitions GetDefinitions() const;

  using MapType = std::unordered_map<std::string, std::string>;
  MapType map_;
  mutable std::string block_;
#ifndef _WIN32
  mutable std::vector<char*> env_;
#endif  // !_WIN32
};

#endif  // NINJA_PROCESS_UTILS_H