libc/src/stdlib/environ_internal.cpp

//===----------------------------------------------------------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
///
/// \file
/// Implementation of internal environment management utilities.
///
//===----------------------------------------------------------------------===//

#include "src/stdlib/environ_internal.h"
#include "config/app.h"
#include "src/__support/CPP/new.h"
#include "src/__support/CPP/string_view.h"
#include "src/__support/alloc-checker.h"
#include "src/__support/macros/config.h"
#include "src/string/memory_utils/inline_memcpy.h"
#ifdef LIBC_COPT_SUPPORT_ENVIRON
#include "src/unistd/environ.h"
#endif

namespace LIBC_NAMESPACE_DECL {
namespace internal {

// Minimum initial capacity for the environment array when first allocated.
// This avoids frequent reallocations for small environments.
constexpr size_t MIN_ENVIRON_CAPACITY = 32;

// Growth factor for environment array capacity when expanding.
// When capacity is exceeded, new_capacity = old_capacity *
// ENVIRON_GROWTH_FACTOR.
constexpr size_t ENVIRON_GROWTH_FACTOR = 2;

void EnvironmentManager::init_once() {
  if (initialized)
    return;

  // Count entries in the startup environ.
  char **env_ptr = reinterpret_cast<char **>(app.env_ptr);
  if (env_ptr) {
    size_t c = 0;
    for (char **env = env_ptr; *env != nullptr; env++)
      c++;
    count = c;
  }

  initialized = true;
}

EnvironmentManager &EnvironmentManager::get_instance() {
  static EnvironmentManager mgr;
  mgr.init_once();
  return mgr;
}

char **EnvironmentManager::get_array() {
  if (is_ours)
    return storage;
  return reinterpret_cast<char **>(app.env_ptr);
}

EnvironmentManager::iterator EnvironmentManager::begin() { return get_array(); }

EnvironmentManager::iterator EnvironmentManager::end() {
  return get_array() + count;
}

size_t EnvironmentManager::size() const { return count; }

char *EnvironmentManager::get(cpp::string_view name) {
  cpp::optional<size_t> idx = find_var(name);
  if (!idx)
    return nullptr;
  return get_array()[*idx] + name.size() + 1;
}

cpp::optional<size_t> EnvironmentManager::find_var(cpp::string_view name) {
  char **env_array = get_array();
  if (!env_array)
    return cpp::nullopt;

  for (size_t i = 0; i < count; i++) {
    cpp::string_view current(env_array[i]);
    if (current.starts_with(name) && current.size() > name.size() &&
        current[name.size()] == '=')
      return i;
  }

  return cpp::nullopt;
}

// Helper: allocate new storage and ownership arrays of the given capacity,
// copy the first `copy_count` entries from old_storage/old_ownership, and
// initialize the remaining ownership slots to default (not-owned).
// Returns nullopt on allocation failure; the old arrays are untouched.
cpp::optional<EnvironmentManager::AllocResult>
EnvironmentManager::alloc_and_copy(size_t new_capacity, char **old_storage,
                                   EnvStringOwnership *old_ownership,
                                   size_t copy_count) {
  AllocChecker ac;
  char **new_storage = new (ac) char *[new_capacity + 1];
  if (!ac)
    return cpp::nullopt;

  EnvStringOwnership *new_ownership =
      new (ac) EnvStringOwnership[new_capacity + 1];
  if (!ac) {
    delete[] new_storage;
    return cpp::nullopt;
  }

  for (size_t i = 0; i < copy_count; i++) {
    new_storage[i] = old_storage ? old_storage[i] : nullptr;
    new_ownership[i] = old_ownership ? old_ownership[i] : EnvStringOwnership();
  }
  new_storage[copy_count] = nullptr;

  return AllocResult{new_storage, new_ownership};
}

bool EnvironmentManager::ensure_capacity(size_t needed) {
  // If we're still using the startup environ (pointed to by app.env_ptr),
  // we must transition to our own managed storage. This allows us to
  // track ownership of strings and safely expand the array.
  if (!is_ours) {
    char **old_env = reinterpret_cast<char **>(app.env_ptr);

    // Allocate new array with room to grow.
    size_t new_capacity = needed < MIN_ENVIRON_CAPACITY
                              ? MIN_ENVIRON_CAPACITY
                              : needed * ENVIRON_GROWTH_FACTOR;

    auto result = alloc_and_copy(new_capacity, old_env, nullptr, count);
    if (!result)
      return false;

    auto [new_storage, new_ownership] = *result;
    storage = new_storage;
    ownership = new_ownership;
    capacity = new_capacity;
    is_ours = true;

    // Update the global environ pointer.
    app.env_ptr = reinterpret_cast<uintptr_t *>(storage);
#ifdef LIBC_COPT_SUPPORT_ENVIRON
    environ = storage;
#endif

    return true;
  }

  // We already own the environment array. Check if it's large enough.
  if (needed <= capacity)
    return true;

  // Grow capacity. We avoid realloc to ensure that failures don't leave the
  // manager in an inconsistent state.
  size_t new_capacity = needed * ENVIRON_GROWTH_FACTOR;

  auto result = alloc_and_copy(new_capacity, storage, ownership, count);
  if (!result)
    return false;

  delete[] storage;
  delete[] ownership;

  auto [new_storage, new_ownership] = *result;
  storage = new_storage;
  ownership = new_ownership;
  capacity = new_capacity;

  // Update the global environ pointer.
  app.env_ptr = reinterpret_cast<uintptr_t *>(storage);
#ifdef LIBC_COPT_SUPPORT_ENVIRON
  environ = storage;
#endif

  return true;
}

int EnvironmentManager::set(cpp::string_view name, cpp::string_view value,
                            bool overwrite) {
  cpp::optional<size_t> idx = find_var(name);

  // If the variable exists and we're not overwriting, do nothing.
  if (idx && !overwrite)
    return 0;

  // Ensure we have capacity. If the variable doesn't exist, we need one
  // more slot.
  size_t needed = idx ? count : count + 1;
  if (!ensure_capacity(needed))
    return -1;

  // Build the "name=value" string.
  size_t name_len = name.size();
  size_t value_len = value.size();
  size_t total_len = name_len + 1 + value_len + 1; // name + '=' + value + '\0'

  AllocChecker ac;
  char *new_string = new (ac) char[total_len];
  if (!ac)
    return -1;

  inline_memcpy(new_string, name.data(), name_len);
  new_string[name_len] = '=';
  inline_memcpy(new_string + name_len + 1, value.data(), value_len);
  new_string[name_len + 1 + value_len] = '\0';

  char **env_array = get_array();

  if (idx) {
    // Replace existing variable. Free old string if we own it.
    if (ownership[*idx].can_free())
      delete[] env_array[*idx];

    env_array[*idx] = new_string;
    ownership[*idx].allocated_by_us = true;
  } else {
    // Add new variable at the end.
    env_array[count] = new_string;
    ownership[count].allocated_by_us = true;
    count++;
    env_array[count] = nullptr; // Maintain null terminator.
  }

  return 0;
}

int EnvironmentManager::unset(cpp::string_view name) {
  cpp::optional<size_t> idx = find_var(name);
  if (!idx)
    return 0; // Variable not found; POSIX defines this as success.

  // Transition to managed storage so we can modify the array and track
  // ownership correctly.
  if (!ensure_capacity(count))
    return -1;

  char **env_array = get_array();

  // Loop to remove all instances of the variable (e.g. duplicates from execve).
  while (idx) {
    size_t i = *idx;
    if (ownership[i].can_free())
      delete[] env_array[i];

    // Compact: shift remaining entries left to fill the gap.
    // Shifting elements preserves the order of environment variables, which is
    // desirable to maintain consistency with getenv resolution order and
    // typical environ iteration behavior.
    for (size_t j = i; j < count - 1; j++) {
      env_array[j] = env_array[j + 1];
      ownership[j] = ownership[j + 1];
    }
    count--;
    env_array[count] = nullptr;

    idx = find_var(name);
  }

  return 0;
}

int EnvironmentManager::put(char *string) {
  cpp::string_view sv(string);
  size_t eq_pos = sv.find_first_of('=');

  // No '=' found: treat as unset (POSIX behavior).
  if (eq_pos == cpp::string_view::npos)
    return unset(sv);

  cpp::string_view name = sv.substr(0, eq_pos);

  cpp::optional<size_t> idx = find_var(name);

  size_t needed = idx ? count : count + 1;
  if (!ensure_capacity(needed))
    return -1;

  char **env_array = get_array();

  if (idx) {
    // Replace existing variable. Free old string if we own it.
    if (ownership[*idx].can_free())
      delete[] env_array[*idx];

    env_array[*idx] = string;
    ownership[*idx].allocated_by_us = false; // Caller owns this string.
  } else {
    // Add new variable at the end.
    env_array[count] = string;
    ownership[count].allocated_by_us = false;
    count++;
    env_array[count] = nullptr;
  }

  return 0;
}

} // namespace internal
} // namespace LIBC_NAMESPACE_DECL