//===--- ImageInspectionMachO.cpp - Mach-O image inspection ---------------===//
//
// This source file is part of the Swift.org open source project
//
// Copyright (c) 2014 - 2017 Apple Inc. and the Swift project authors
// Licensed under Apache License v2.0 with Runtime Library Exception
//
// See https://swift.org/LICENSE.txt for license information
// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
//
//===----------------------------------------------------------------------===//
///
/// \file
///
/// This file includes routines that interact with dyld on Mach-O-based
/// platforms to extract runtime metadata embedded in images generated by the
/// Swift compiler.
///
//===----------------------------------------------------------------------===//

#if defined(__APPLE__) && defined(__MACH__)

#include "ImageInspection.h"
#include <mach-o/dyld.h>
#include <mach-o/getsect.h>
#include <objc/runtime.h>
#include <assert.h>
#include <dlfcn.h>

using namespace swift;

namespace {
/// The Mach-O section name for the section containing protocol descriptor
/// references. This lives within SEG_TEXT.
constexpr const char ProtocolsSection[] = "__swift5_protos";
/// The Mach-O section name for the section containing protocol conformances.
/// This lives within SEG_TEXT.
constexpr const char ProtocolConformancesSection[] = "__swift5_proto";
/// The Mach-O section name for the section containing type references.
/// This lives within SEG_TEXT.
constexpr const char TypeMetadataRecordSection[] = "__swift5_types";
/// The Mach-O section name for the section containing dynamic replacements.
/// This lives within SEG_TEXT.
constexpr const char DynamicReplacementSection[] = "__swift5_replace";
constexpr const char DynamicReplacementSomeSection[] = "__swift5_replac2";

constexpr const char TextSegment[] = SEG_TEXT;

#if __POINTER_WIDTH__ == 64
using mach_header_platform = mach_header_64;
#else
using mach_header_platform = mach_header;
#endif

extern "C" void *_NSGetMachExecuteHeader();

template <const char *SEGMENT_NAME, const char *SECTION_NAME,
         void CONSUME_BLOCK(const void *start, uintptr_t size)>
void addImageCallback(const mach_header *mh) {
#if __POINTER_WIDTH__ == 64
  assert(mh->magic == MH_MAGIC_64 && "loaded non-64-bit image?!");
#endif
  
  // Look for a __swift5_proto section.
  unsigned long size;
  const uint8_t *section =
  getsectiondata(reinterpret_cast<const mach_header_platform *>(mh),
                 SEGMENT_NAME, SECTION_NAME,
                 &size);
  
  if (!section)
    return;
  
  CONSUME_BLOCK(section, size);
}
template <const char *SEGMENT_NAME, const char *SECTION_NAME,
         void CONSUME_BLOCK(const void *start, uintptr_t size)>
void addImageCallback(const mach_header *mh, intptr_t vmaddr_slide) {
  addImageCallback<SEGMENT_NAME, SECTION_NAME, CONSUME_BLOCK>(mh);
}

template <const char *SEGMENT_NAME, const char *SECTION_NAME,
          const char *SEGMENT_NAME2, const char *SECTION_NAME2,
          void CONSUME_BLOCK(const void *start, uintptr_t size,
                             const void *start2, uintptr_t size2)>
void addImageCallback2Sections(const mach_header *mh) {
#if __POINTER_WIDTH__ == 64
  assert(mh->magic == MH_MAGIC_64 && "loaded non-64-bit image?!");
#endif

  // Look for a section.
  unsigned long size;
  const uint8_t *section =
  getsectiondata(reinterpret_cast<const mach_header_platform *>(mh),
                 SEGMENT_NAME, SECTION_NAME,
                 &size);

  if (!section)
    return;

  // Look for another section.
  unsigned long size2;
  const uint8_t *section2 =
  getsectiondata(reinterpret_cast<const mach_header_platform *>(mh),
                 SEGMENT_NAME2, SECTION_NAME2,
                 &size2);
  if (!section2)
    size2 = 0;

  CONSUME_BLOCK(section, size, section2, size2);
}
template <const char *SEGMENT_NAME, const char *SECTION_NAME,
          const char *SEGMENT_NAME2, const char *SECTION_NAME2,
          void CONSUME_BLOCK(const void *start, uintptr_t size,
                             const void *start2, uintptr_t size2)>
void addImageCallback2Sections(const mach_header *mh, intptr_t vmaddr_slide) {
  addImageCallback2Sections<SEGMENT_NAME, SECTION_NAME,
                            SEGMENT_NAME2, SECTION_NAME2,
                            CONSUME_BLOCK>(mh);
}

} // end anonymous namespace

#if OBJC_ADDLOADIMAGEFUNC_DEFINED
#define REGISTER_FUNC(...)                                               \
  if (__builtin_available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *)) { \
    objc_addLoadImageFunc(__VA_ARGS__);                                  \
  } else {                                                               \
    _dyld_register_func_for_add_image(__VA_ARGS__);                      \
  }
#else
#define REGISTER_FUNC(...) _dyld_register_func_for_add_image(__VA_ARGS__)
#endif

void swift::initializeProtocolLookup() {
  REGISTER_FUNC(
    addImageCallback<TextSegment, ProtocolsSection,
                     addImageProtocolsBlockCallback>);
}

void swift::initializeProtocolConformanceLookup() {
  REGISTER_FUNC(
    addImageCallback<TextSegment, ProtocolConformancesSection,
                     addImageProtocolConformanceBlockCallback>);
}
void swift::initializeTypeMetadataRecordLookup() {
  REGISTER_FUNC(
    addImageCallback<TextSegment, TypeMetadataRecordSection,
                     addImageTypeMetadataRecordBlockCallback>);
}

void swift::initializeDynamicReplacementLookup() {
  REGISTER_FUNC(
      addImageCallback2Sections<TextSegment, DynamicReplacementSection,
                                TextSegment, DynamicReplacementSomeSection,
                                addImageDynamicReplacementBlockCallback>);
}

int swift::lookupSymbol(const void *address, SymbolInfo *info) {
  Dl_info dlinfo;
  if (dladdr(address, &dlinfo) == 0) {
    return 0;
  }

  info->fileName = dlinfo.dli_fname;
  info->baseAddress = dlinfo.dli_fbase;
  info->symbolName.reset(dlinfo.dli_sname);
  info->symbolAddress = dlinfo.dli_saddr;
  return 1;
}

void *swift::lookupSection(const char *segment, const char *section, size_t *outSize) {
  unsigned long size;
  auto *executableHeader = static_cast<mach_header_platform *>(_NSGetMachExecuteHeader());
  uint8_t *data = getsectiondata(executableHeader, segment, section, &size);
  if (outSize != nullptr && data != nullptr)
    *outSize = size;
  return static_cast<void *>(data);
}

#endif // defined(__APPLE__) && defined(__MACH__)