codeCache.cpp

/*
 * Copyright The async-profiler authors
 * Copyright 2026, Datadog, Inc.
 * SPDX-License-Identifier: Apache-2.0
 */

#include "codeCache.h"
#include "dwarf.h"
#include "os.h"
#include "safeAccess.h"

#include <cassert>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>

char *NativeFunc::create(const char *name, short lib_index) {
  size_t size = allocSize(name);
  NativeFunc *f = (NativeFunc *)aligned_alloc(sizeof(NativeFunc*), size);
  f->_lib_index = lib_index;
  f->_mark = 0;
  // cppcheck-suppress memleak
  return strcpy(f->_name, name);
}

void NativeFunc::destroy(char *name) { free(from(name)); }

char NativeFunc::read_mark(const char* name) {
  if (name == nullptr) {
    return 0;
  }
  NativeFunc* func = from(name);
  if (!is_aligned(func, sizeof(func))) {
    return 0;
  }
  // Use SafeAccess to read the mark field in signal handler context
  // Read the first 4 bytes (lib_index + mark + reserved) and extract the mark byte
  int32_t prefix = SafeAccess::safeFetch32((int32_t*)func, 0);
  // Extract mark byte: shift right by 16 bits to skip lib_index (2 bytes), mask to 1 byte
  return (char)((prefix >> 16) & 0xFF);
}

CodeCache::CodeCache(const char *name, short lib_index,
                     const void *min_address, const void *max_address,
                     const char* image_base, bool imports_patchable) {
  _name = NativeFunc::create(name, -1);
  _lib_index = lib_index;
  _min_address = min_address;
  _max_address = max_address;
  _text_base = NULL;
  _image_base = image_base;

  _plt_offset = 0;
  _plt_size = 0;
  _debug_symbols = false;

  // Initialize build-id fields
  _build_id = nullptr;
  _build_id_len = 0;
  _load_bias = 0;
  _file_id = 0;
  _image_fingerprint = 0;

  memset(_imports, 0, sizeof(_imports));
  _imports_patchable = imports_patchable;

  _dwarf_table = NULL;
  _dwarf_table_length = 0;
  _default_frame = &FrameDesc::default_frame;

  _capacity = INITIAL_CODE_CACHE_CAPACITY;
  _count = 0;
  _blobs = new CodeBlob[_capacity];

  _memory_usage = (long long)_capacity * sizeof(CodeBlob) +
                  (long long)NativeFunc::allocSize(_name);
  _name_overhead = (long long)NativeFunc::nameOverhead(_name);

  _published.store(false, std::memory_order_relaxed);
}

void CodeCache::copyFrom(const CodeCache& other) {
  _name = NativeFunc::create(other._name, -1);
  _lib_index = other._lib_index;
  _min_address = other._min_address;
  _max_address = other._max_address;
  _text_base = other._text_base;
  _image_base = other._image_base;

  _plt_offset = other._plt_offset;
  _plt_size = other._plt_size;
  _debug_symbols = other._debug_symbols;

  // Copy build-id information
  _build_id_len = other._build_id_len;
  if (other._build_id != nullptr && other._build_id_len > 0) {
    size_t hex_str_len = strlen(other._build_id);
    _build_id = static_cast<char*>(malloc(hex_str_len + 1));
    if (_build_id != nullptr) {
      strcpy(_build_id, other._build_id);
    }
  } else {
    _build_id = nullptr;
  }
  _load_bias = other._load_bias;
  _file_id = other._file_id;
  _image_fingerprint = other._image_fingerprint;

  memset(_imports, 0, sizeof(_imports));
  _imports_patchable = other._imports_patchable;

  _dwarf_table_length = other._dwarf_table_length;
  if (_dwarf_table_length > 0) {
    _dwarf_table = (FrameDesc*)malloc(_dwarf_table_length * sizeof(FrameDesc));
    memcpy(_dwarf_table, other._dwarf_table,
           _dwarf_table_length * sizeof(FrameDesc));
  } else {
    _dwarf_table = nullptr;
  }
  _default_frame = other._default_frame;

  _capacity = other._capacity;
  _count = other._count;
  _blobs = new CodeBlob[_capacity];
  memcpy(_blobs, other._blobs, _count * sizeof(CodeBlob));
  // The memcpy above copied each CodeBlob's _name *pointer*, so without this the
  // two caches would share name allocations and both destructors would free
  // them (double-free / use-after-free). Give this copy its own name strings.
  for (int i = 0; i < _count; i++) {
    if (_blobs[i]._name != nullptr) {
      _blobs[i]._name = NativeFunc::create(other._blobs[i]._name, _lib_index);
    }
  }

  // Identical content to `other`, so its running total transfers as-is --
  // no need to recompute (and `other` may itself already be published, whose
  // _blobs this copy must not depend on after construction).
  _memory_usage = other._memory_usage;
  // Likewise for the overhead: the name strings above are fresh allocations,
  // but of exactly the same sizes, and overhead is a function of the requested
  // size, so `other`'s total describes this copy just as well.
  _name_overhead = other._name_overhead;

  // A copy is a fresh, not-yet-registered cache.
  _published.store(false, std::memory_order_relaxed);
}

CodeCache::CodeCache(const CodeCache &other) {
  copyFrom(other);
}

CodeCache &CodeCache::operator=(const CodeCache &other) {
  if (&other == this) {
    return *this;
  }

  NativeFunc::destroy(_name);
  free(_dwarf_table);
  delete[] _blobs;
  free(_build_id);

  copyFrom(other);

  return *this;
}

CodeCache::~CodeCache() {
  for (int i = 0; i < _count; i++) {
    NativeFunc::destroy(_blobs[i]._name);
  }
  NativeFunc::destroy(_name);
  delete[] _blobs;
  free(_dwarf_table);
  free(_build_id);  // Free build-id memory
}

long long CodeCache::memoryUsage() const {
  // O(1): _memory_usage is updated incrementally by the constructor,
  // add()/commitAdd(), and setDwarfTable() -- the same pre-publication mutators that
  // used to be re-summed here on every call. Safe to read lock-free once
  // published (Libraries::native_libs) because those mutators never run
  // afterwards (asserted via _published), same invariant the old per-call
  // rescan relied on.
  //
  // Deliberately excludes the build-id string: the background library
  // refresher (Libraries::updateBuildIds) frees and replaces _build_id on
  // already-published caches under _build_id_lock, which dump does not hold,
  // so folding it into this running total would race. It is negligible
  // (~tens of bytes per library) next to the symbol tables.
  return _memory_usage;
}

void CodeCache::prepareAdd(const char *name, PreparedAdd &prepared) const {
  // Symbols are added while parsing a library, before it is registered into a
  // CodeCacheArray. Adding after publication would race memoryUsage() (which
  // reads _blobs/_count lock-free at dump time). Unpublished standalone caches
  // (e.g. _runtime_stubs) are exempt — they are never published.
  assert(!_published.load(std::memory_order_acquire) &&
         "add() on a published CodeCache races memoryUsage()");
  // A reused PreparedAdd would leak the allocations it still owns.
  assert(prepared.name == nullptr && prepared.grown == nullptr);
  char *name_copy = NativeFunc::create(name, _lib_index);
  prepared.name = name_copy;
  prepared.name_size = (long long)NativeFunc::allocSize(name);
  // Measured off the block just handed back, while the pointer is in hand: the
  // gauge is published as a total at add() time, with no pointer left to probe.
  prepared.name_overhead = (long long)NativeFunc::nameOverhead(name_copy);
  // Replace non-printable characters
  for (char *s = name_copy; *s != 0; s++) {
    if (*s < ' ')
      *s = '?';
  }

  prepared.capacity = _capacity;
  prepared.grown = nullptr;
  if (_count >= _capacity) {
    prepared.grown = new CodeBlob[_capacity * 2];
    memcpy(prepared.grown, _blobs, _count * sizeof(CodeBlob));
  }
}

CodeBlob *CodeCache::commitAdd(const void *start, int length,
                               PreparedAdd &prepared,
                               bool update_bounds) {
  assert(!_published.load(std::memory_order_acquire) &&
         "add() on a published CodeCache races memoryUsage()");
  // A mutation between prepareAdd() and commitAdd() would make the grown copy
  // stale or leave no free slot.
  assert(prepared.capacity == _capacity &&
         (prepared.grown != nullptr || _count < _capacity));
  CodeBlob *retired = nullptr;
  if (prepared.grown != nullptr) {
    // Delta is exactly the current (pre-doubling) capacity's worth of blobs.
    _memory_usage += (long long)_capacity * sizeof(CodeBlob);
    _capacity *= 2;
    retired = _blobs;
    _blobs = prepared.grown;
  }
  _memory_usage += prepared.name_size;
  _name_overhead += prepared.name_overhead;

  const void *end = (const char *)start + length;
  _blobs[_count]._start = start;
  _blobs[_count]._end = end;
  _blobs[_count]._name = prepared.name;
  _count++;
  // The cache owns both allocations now; the PreparedAdd destructor must not
  // free them.
  prepared.name = nullptr;
  prepared.grown = nullptr;

  if (update_bounds) {
    updateBounds(start, end);
  }
  return retired;
}

void CodeCache::add(const void *start, int length, const char *name,
                    bool update_bounds) {
  PreparedAdd prepared;
  prepareAdd(name, prepared);
  delete[] commitAdd(start, length, prepared, update_bounds);
}

void CodeCache::updateBounds(const void *start, const void *end) {
  if (start < _min_address)
    _min_address = start;
  if (end > _max_address)
    _max_address = end;
}

void CodeCache::sort() {
  if (_count == 0)
    return;

  qsort(_blobs, _count, sizeof(CodeBlob), CodeBlob::comparator);

  if (_min_address == NO_MIN_ADDRESS)
    _min_address = _blobs[0]._start;
  if (_max_address == NO_MAX_ADDRESS)
    _max_address = _blobs[_count - 1]._end;
}

CodeBlob *CodeCache::findBlob(const char *name) {
  for (int i = 0; i < _count; i++) {
    const char *blob_name = _blobs[i]._name;
    if (blob_name != NULL && strcmp(blob_name, name) == 0) {
      return &_blobs[i];
    }
  }
  return NULL;
}

CodeBlob *CodeCache::findBlobByAddress(const void *address) {
  for (int i = 0; i < _count; i++) {
    if (address >= _blobs[i]._start && address < _blobs[i]._end) {
      return &_blobs[i];
    }
  }
  return NULL;
}

const void *CodeCache::binarySearch(const void *address, const char **name) {
  int low = 0;
  int high = _count - 1;

  while (low <= high) {
    int mid = (unsigned int)(low + high) >> 1;
    if (_blobs[mid]._end <= address) {
      low = mid + 1;
    } else if (_blobs[mid]._start > address) {
      high = mid - 1;
    } else {
      if (name != NULL) {
        *name = _blobs[mid]._name;
      }
      return _blobs[mid]._start;
    }
  }

  // Symbols with zero size can be valid functions: e.g. ASM entry points or
  // kernel code. Also, in some cases (endless loop) the return address may
  // point beyond the function.
  if (low > 0 && (_blobs[low - 1]._start == _blobs[low - 1]._end ||
                  _blobs[low - 1]._end == address)) {

    if (name != NULL) {
      *name = _blobs[low - 1]._name;
    }

    return _blobs[low - 1]._start;
  }
  return _name;
}

void CodeCache::dump() {
#ifdef TRACE
  fprintf(stdout, "Dumping symbols for %s:\n+-\n", _name);
  for (int i = 0; i < _count; i++) {
    fprintf(stdout, "%d. %s\n", i, _blobs[i]._name);
  }
  fprintf(stdout, "+-\n");
#endif // TRACE
}

const void *CodeCache::findSymbol(const char *name) {
  CodeBlob *blob = findBlob(name);
  return blob == NULL ? NULL : blob->_start;
}

const void *CodeCache::findSymbolByPrefix(const char *prefix) {
  return findSymbolByPrefix(prefix, strlen(prefix));
}

const void *CodeCache::findSymbolByPrefix(const char *prefix, int prefix_len) {
  for (int i = 0; i < _count; i++) {
    const char *blob_name = _blobs[i]._name;
    if (blob_name != NULL && strncmp(blob_name, prefix, prefix_len) == 0) {
      return _blobs[i]._start;
    }
  }
  return NULL;
}

void CodeCache::findSymbolsByPrefix(std::vector<const char *> &prefixes,
                                    std::vector<const void *> &symbols) {
  std::vector<int> prefix_lengths;
  prefix_lengths.reserve(prefixes.size());
  for (const char *prefix : prefixes) {
    prefix_lengths.push_back(strlen(prefix));
  }
  for (int i = 0; i < _count; i++) {
    const char *blob_name = _blobs[i]._name;
    if (blob_name != NULL) {
      for (size_t i = 0; i < prefixes.size(); i++) {
        if (strncmp(blob_name, prefixes[i], prefix_lengths[i]) == 0) {
          symbols.push_back(_blobs[i]._start);
        }
      }
    }
  }
}

void CodeCache::saveImport(ImportId id, void** entry) {
    for (int ty = 0; ty < NUM_IMPORT_TYPES; ty++) {
        if (_imports[id][ty] == nullptr) {
            _imports[id][ty] = entry;
            return;
        }
    }
}

void CodeCache::addImport(void **entry, const char *name) {
    switch (name[0]) {
      case 'a':
          if (strcmp(name, "aligned_alloc") == 0) {
              saveImport(im_aligned_alloc, entry);
          }
          break;
      case 'c':
          if (strcmp(name, "calloc") == 0) {
              saveImport(im_calloc, entry);
          }
          break;
      case 'd':
          if (strcmp(name, "dlopen") == 0) {
              saveImport(im_dlopen, entry);
          }
          break;
      case 'f':
          if (strcmp(name, "free") == 0) {
              saveImport(im_free, entry);
          }
          break;
      case 'm':
          if (strcmp(name, "malloc") == 0) {
              saveImport(im_malloc, entry);
          }
          break;
      case 'p':
          if (strcmp(name, "pthread_create") == 0) {
              saveImport(im_pthread_create, entry);
          } else if (strcmp(name, "pthread_exit") == 0) {
              saveImport(im_pthread_exit, entry);
          } else if (strcmp(name, "pthread_setspecific") == 0) {
              saveImport(im_pthread_setspecific, entry);
          } else if (strcmp(name, "poll") == 0) {
              saveImport(im_poll, entry);
          } else if (strcmp(name, "posix_memalign") == 0) {
              saveImport(im_posix_memalign, entry);
          }
          break;
      case 'r':
          if (strcmp(name, "realloc") == 0) {
              saveImport(im_realloc, entry);
          } else if (strcmp(name, "recv") == 0) {
              saveImport(im_recv, entry);
          } else if (strcmp(name, "read") == 0) {
              saveImport(im_read, entry);
          }
          break;
      case 's':
          if (strcmp(name, "send") == 0) {
              saveImport(im_send, entry);
          } else if (strcmp(name, "sigaction") == 0) {
              saveImport(im_sigaction, entry);
          }
          break;
      case 'w':
          if (strcmp(name, "write") == 0) {
              saveImport(im_write, entry);
          }
          break;
    }
}

void **CodeCache::findImport(ImportId id) {
  if (!_imports_patchable) {
    makeImportsPatchable();
    _imports_patchable = true;
  }
  return _imports[id][PRIMARY];
}

void CodeCache::patchImport(ImportId id, void *hook_func) {
  if (!_imports_patchable) {
        makeImportsPatchable();
        _imports_patchable = true;
    }

    for (int ty = 0; ty < NUM_IMPORT_TYPES; ty++) {void **entry = _imports[id][ty];
  if (entry != NULL) {
    *entry = hook_func;
  }}
}

void CodeCache::makeImportsPatchable() {
  void **min_import = (void **)-1;
  void **max_import = NULL;
  for (int i = 0; i < NUM_IMPORTS; i++) {
    for (int j = 0; j < NUM_IMPORT_TYPES; j++) {
            void** entry = _imports[i][j];
            if (entry == NULL) continue;
            if (entry < min_import)
      min_import = entry;
    if (entry > max_import)
      max_import = entry;
        }
  }

  if (max_import != NULL) {
    uintptr_t patch_start = (uintptr_t)min_import & ~OS::page_mask;
    uintptr_t patch_end = (uintptr_t)max_import & ~OS::page_mask;
    mprotect((void *)patch_start, patch_end - patch_start + OS::page_size,
             PROT_READ | PROT_WRITE);
  }
}

void CodeCache::setDwarfTable(FrameDesc *table, int length, const FrameDesc &default_frame) {
  // Set during library parsing, before publication. memoryUsage() reads
  // _dwarf_table_length lock-free at dump time, so this must not run afterwards.
  assert(!_published.load(std::memory_order_acquire) &&
         "setDwarfTable() on a published CodeCache races memoryUsage()");
  // The parser (SFrameParser/DwarfParser) already shrinks its capacity-doubled
  // buffer to exactly `length` entries before returning it from table(), so
  // memoryUsage()'s length-based formula matches the real allocation with no
  // trim step needed here.
  _memory_usage += (long long)length * sizeof(FrameDesc);
  _dwarf_table = table;
  _dwarf_table_length = length;
  _default_frame = &default_frame;
}

FrameDesc CodeCache::findFrameDesc(const void *pc) {
  if (_dwarf_table == NULL || _dwarf_table_length == 0) {
    return *_default_frame;
  }

  u32 target_loc = (const char *)pc - _text_base;
  int low = 0;
  int high = _dwarf_table_length - 1;

  while (low <= high) {
    int mid = (unsigned int)(low + high) >> 1;
    if (_dwarf_table[mid].loc < target_loc) {
      low = mid + 1;
    } else if (_dwarf_table[mid].loc > target_loc) {
      high = mid - 1;
    } else {
      return _dwarf_table[mid];
    }
  }

  if (low > 0) {
    return _dwarf_table[low - 1];
  } else if (target_loc - _plt_offset < _plt_size) {
    return FrameDesc::empty_frame;
  } else {
    return *_default_frame;
  }
}

void CodeCache::setBuildId(const char* build_id, size_t build_id_len) {
  // Free existing build-id if any
  free(_build_id);
  _build_id = nullptr;
  _build_id_len = 0;

  if (build_id != nullptr && build_id_len > 0) {
    // build_id is a hex string, allocate based on actual string length
    size_t hex_str_len = strlen(build_id);
    _build_id = static_cast<char*>(malloc(hex_str_len + 1));

    if (_build_id != nullptr) {
      // Copy the hex string
      strcpy(_build_id, build_id);
      // Store the original byte length (not hex string length)
      _build_id_len = build_id_len;
    }
  }
}