allocation_fault_injection.cc

/*
 * 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 {

struct ThreadState {
  std::array<std::uint32_t, scenario_limits::kMaxAllocationFailures> ordinals{};
  std::size_t ordinal_count = 0;
  std::size_t next_ordinal = 0;
  std::uint64_t allocation_attempts = 0;
  std::uint64_t injected_failures = 0;
  bool armed = false;
};

thread_local ThreadState state;
std::once_flag initialize_once;
bool callbacks_active = false;
bool curl_initialized = false;
std::atomic<bool> curl_initialization_claimed{false};

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);
}

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;
}

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

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;
}

bool ReinitializeCurl() noexcept {
  return curl_global_init_mem(CURL_GLOBAL_ALL, internal::Malloc, internal::Free, internal::Realloc, internal::Strdup,
                              internal::Calloc) == CURLE_OK;
}

bool CurlInitializationSucceeded() noexcept {
  (void)InitializeCurl();
  return curl_initialized;
}

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);
}

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]));
  }
}

void Arm() noexcept {
  state.next_ordinal = 0;
  state.allocation_attempts = 0;
  state.injected_failures = 0;
  state.armed = state.ordinal_count != 0;
}

void ArmNextAllocationFailure() noexcept {
  state = ThreadState{};
  state.ordinals[0] = 1;
  state.ordinal_count = 1;
  state.armed = true;
}

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

void Reset() noexcept {
  state = ThreadState{};
  internal::ResetFailureSuppression();
}

Statistics GetStatistics() noexcept {
  return {state.allocation_attempts, state.injected_failures, state.ordinal_count,
          state.armed && !internal::FailuresSuppressed()};
}
 
namespace internal {

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

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

void* Realloc(void* pointer, std::size_t size) noexcept {
  return ShouldFail(size != 0) ? nullptr : std::realloc(pointer, size);
}

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;
}

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