/*

 * Copyright (C) Max Dymond, <cmeister2@gmail.com>, et al.

 *

 * SPDX-License-Identifier: curl

 */

#include "proto_fuzzer/allocation_fault_injection.h"

#include <curl/curl.h>

#include <algorithm>

#include <array>

#include <atomic>

#include <cerrno>

#include <cstdlib>

#include <cstring>

#include <mutex>

#include "proto_fuzzer/scenario_limits.h"

namespace proto_fuzzer::allocation_fault_injection {

namespace {

/// Keep one thread's bounded failure plan and allocation counters without

/// allocating memory in the controller itself.

struct ThreadState {

  /// Sorted, distinct one-based ordinals; only ordinal_count entries are valid.

  std::array<std::uint32_t, scenario_limits::kMaxAllocationFailures> ordinals{};

  /// Number of valid ordinals retained in the fixed-capacity plan.

  std::size_t ordinal_count = 0;

  /// Index of the next configured ordinal that has not yet failed.

  std::size_t next_ordinal = 0;

  /// Positive-size allocation attempts observed while this thread is armed.

  std::uint64_t allocation_attempts = 0;

  /// Allocation attempts rejected at configured ordinals.

  std::uint64_t injected_failures = 0;

  /// Whether the plan is armed; suppression can still bypass its callbacks.

  bool armed = false;

};

/// Isolate each calling thread's failure plan and counters.

thread_local ThreadState state;

/// Install and verify libcurl's process-wide allocator callbacks only once.

std::once_flag initialize_once;

/// Remember whether the initialization probe reached the installed callbacks.

bool callbacks_active = false;

/// Remember whether this module owns one successful curl global initialization.

bool curl_initialized = false;

/// Transfer that one initialization reference to at most one lifecycle owner.

std::atomic<bool> curl_initialization_claimed{false};

/// Fold a positive ordinal into the bounded one-based allocation range.

/// @param ordinal Requested allocation attempt number; zero disables the entry.

/// @return Zero for a disabled entry, otherwise a value in

/// 1..scenario_limits::kMaxAllocationOrdinal.

std::uint32_t CanonicalOrdinal(std::uint32_t ordinal) noexcept {

  if (ordinal == 0) {

    return 0;

  }

  const std::uint32_t maximum = static_cast<std::uint32_t>(scenario_limits::kMaxAllocationOrdinal);

  return 1U + ((ordinal - 1U) % maximum);

}

/// Insert an ordinal into the calling thread's sorted, unique plan.

/// When the plan is full, retain the smallest ordinals by dropping its largest

/// entry if the new value sorts before it. Zero and duplicates are ignored.

/// @param ordinal Canonical allocation attempt number to retain.

void InsertOrdinal(std::uint32_t ordinal) noexcept {

  if (ordinal == 0) {

    return;

  }

  const auto begin = state.ordinals.begin();

  const auto end = begin + static_cast<std::ptrdiff_t>(state.ordinal_count);

  const auto position = std::lower_bound(begin, end, ordinal);

  if (position != end && *position == ordinal) {

    return;

  }

  const std::size_t index = static_cast<std::size_t>(position - begin);

  if (index >= state.ordinals.size()) {

    return;

  }

  if (state.ordinal_count < state.ordinals.size()) {

    ++state.ordinal_count;

  }

  for (std::size_t destination = state.ordinal_count - 1; destination > index; --destination) {

    state.ordinals[destination] = state.ordinals[destination - 1];

  }

  state.ordinals[index] = ordinal;

}

/// Count an armed positive-size allocation and reject a selected attempt.

/// Each configured ordinal fails once; rejected calls advance the plan and

/// failure counter and set errno to ENOMEM. Zero-size and disarmed calls leave

/// the controller state unchanged.

/// @param positive_size Whether the allocation requests nonzero storage.

/// @return true when the caller should reject this allocation.

bool ShouldFail(bool positive_size) noexcept {

  if (!positive_size || !state.armed || internal::FailuresSuppressed()) {

    return false;

  }

  ++state.allocation_attempts;

  if (state.next_ordinal < state.ordinal_count && state.allocation_attempts == state.ordinals[state.next_ordinal]) {

    ++state.next_ordinal;

    ++state.injected_failures;

    errno = ENOMEM;

    return true;

  }

  return false;

}

}  // namespace

/// Install the allocator callbacks before any other libcurl global

/// initialization. The callbacks remain installed for the process lifetime;

/// individual inputs only change thread-local state. A small allocation

/// verifies that a pre-existing global initialization did not silently keep

/// the default callbacks.

/// @return true when libcurl invoked the installed callbacks.

bool InitializeCurl() noexcept {

  std::call_once(initialize_once, [] {

    if (curl_global_init_mem(CURL_GLOBAL_ALL, internal::Malloc, internal::Free, internal::Realloc, internal::Strdup,

                             internal::Calloc) != CURLE_OK) {

      return;

    }

    curl_initialized = true;

    state.armed = true;

    char* probe = curl_escape("allocator callback probe", 0);

    callbacks_active = state.allocation_attempts != 0;

    state.armed = false;

    curl_free(probe);

    state = ThreadState{};

  });

  return callbacks_active;

}

/// Reinstall the same allocator callbacks after a deliberate, balanced

/// curl_global_cleanup(). Call only when no curl handles or workers remain.

/// @return true when libcurl accepted the callbacks.

bool ReinitializeCurl() noexcept {

  return curl_global_init_mem(CURL_GLOBAL_ALL, internal::Malloc, internal::Free, internal::Realloc, internal::Strdup,

                              internal::Calloc) == CURLE_OK;

}

/// Report whether InitializeCurl owns a successful global initialization.

/// This is distinct from its return value: allocation callbacks can be

/// unavailable even though curl_global_init_mem itself succeeded.

bool CurlInitializationSucceeded() noexcept {

  (void)InitializeCurl();

  return curl_initialized;

}

/// Claim ownership of InitializeCurl's retained global reference exactly once.

/// A failed initialization has no reference to claim, and a reference remains

/// claimed after cleanup because initialize_once cannot create it again.

/// @return true only for the caller that acquires the retained reference.

bool ClaimCurlInitializationReference() noexcept {

  (void)InitializeCurl();

  if (!curl_initialized) {

    return false;

  }

  bool expected = false;

  return curl_initialization_claimed.compare_exchange_strong(expected, true, std::memory_order_relaxed,

                                                             std::memory_order_relaxed);

}

/// Store a bounded, sorted, unique allocation plan for the calling thread.

/// Values outside the useful range are folded into it. Retaining the smallest

/// ordinals makes oversized compatibility inputs deterministic without heap

/// allocation in the controller itself.

/// @param ordinals One-based positive allocation attempt numbers.

/// @param count Number of values available at ordinals.

void Configure(const std::uint32_t* ordinals, std::size_t count) noexcept {

  state = ThreadState{};

  if (ordinals == nullptr) {

    return;

  }

  for (std::size_t index = 0; index < count; ++index) {

    InsertOrdinal(CanonicalOrdinal(ordinals[index]));

  }

}

/// Begin counting positive-size allocation attempts on the calling thread.

/// Each configured ordinal can fail once. A later ordinal is evaluated on the

/// execution path already changed by every earlier injected failure.

void Arm() noexcept {

  state.next_ordinal = 0;

  state.allocation_attempts = 0;

  state.injected_failures = 0;

  state.armed = state.ordinal_count != 0;

}

/// Stop fault injection while retaining counters for diagnostics and tests.

void Disarm() noexcept { state.armed = false; }

/// Clear all thread-local controller state and invalidate live suppression

/// guards. Fuzzer entrypoints call this before parsing so a standalone-runner

/// siglongjmp cannot leak an armed plan or suppression into the next input.

void Reset() noexcept {

  state = ThreadState{};

  internal::ResetFailureSuppression();

}

/// Read the calling thread's current plan and counters.

/// @return A value snapshot that remains valid after later controller calls.

Statistics GetStatistics() noexcept {

  return {state.allocation_attempts, state.injected_failures, state.ordinal_count,

          state.armed && !internal::FailuresSuppressed()};

}

namespace internal {

/// Allocate with libc unless this positive-size attempt is selected to fail.

void* Malloc(std::size_t size) noexcept { return ShouldFail(size != 0) ? nullptr : std::malloc(size); }

/// Release an allocation. Frees are never counted or rejected.

void Free(void* pointer) noexcept { std::free(pointer); }

/// Resize with libc unless this positive-size attempt is selected to fail.

/// Rejecting the call before libc preserves ownership of the original block.

void* Realloc(void* pointer, std::size_t size) noexcept {

  return ShouldFail(size != 0) ? nullptr : std::realloc(pointer, size);

}

/// Duplicate one string while counting the operation as one allocation.

char* Strdup(const char* string) noexcept {

  if (string == nullptr) {

    return nullptr;

  }

  const std::size_t length = std::strlen(string) + 1;

  if (ShouldFail(length != 0)) {

    return nullptr;

  }

  char* copy = static_cast<char*>(std::malloc(length));

  if (copy != nullptr) {

    std::memcpy(copy, string, length);

  }

  return copy;

}

/// Allocate zero-filled storage unless this positive-size request is selected

/// to fail. The request counts once rather than once per element.

void* Calloc(std::size_t count, std::size_t size) noexcept {

  return ShouldFail(count != 0 && size != 0) ? nullptr : std::calloc(count, size);

}

}  // namespace internal

}  // namespace proto_fuzzer::allocation_fault_injection