[ReflectionMirror] Change swift-reflection-test.c to use addReflectionInfo on non-Macho platforms.

Author: mikeash

stdlib/private/SwiftReflectionTest/SwiftReflectionTest.swift

@@ -24,6 +24,7 @@ import Darwin
 let RequestInstanceKind = "k"
 let RequestInstanceAddress = "i"
 let RequestReflectionInfos = "r"
+let RequestImages = "m"
 let RequestReadBytes = "b"
 let RequestSymbolAddress = "s"
 let RequestStringLength = "l"
@@ -45,12 +46,98 @@ public enum InstanceKind : UInt8 {
   case Closure
 }
 
+/// Represents a section in a loaded image in this process.
+internal struct Section {
+  /// The absolute start address of the section's data in this address space.
+  let startAddress: UnsafeRawPointer
+
+  /// The size of the section in bytes.
+  let size: UInt
+}
+
+/// Holds the addresses and sizes of sections related to reflection
+internal struct ReflectionInfo : Sequence {
+  /// The name of the loaded image
+  internal let imageName: String
+
+  /// Reflection metadata sections
+  internal let fieldmd: Section?
+  internal let assocty: Section?
+  internal let builtin: Section?
+  internal let capture: Section?
+  internal let typeref: Section?
+  internal let reflstr: Section?
+
+  internal func makeIterator() -> AnyIterator<Section?> {
+    return AnyIterator([
+      fieldmd,
+      assocty,
+      builtin,
+      capture,
+      typeref,
+      reflstr
+    ].makeIterator())
+  }
+}
+
 #if arch(x86_64) || arch(arm64)
 typealias MachHeader = mach_header_64
 #else
 typealias MachHeader = mach_header
 #endif
 
+/// Get the location and size of a section in a binary.
+///
+/// - Parameter name: The name of the section
+/// - Parameter imageHeader: A pointer to the Mach header describing the
+///   image.
+/// - Returns: A `Section` containing the address and size, or `nil` if there
+///   is no section by the given name.
+internal func getSectionInfo(_ name: String,
+  _ imageHeader: UnsafePointer<MachHeader>) -> Section? {
+  debugLog("BEGIN \(#function)"); defer { debugLog("END \(#function)") }
+  var size: UInt = 0
+  let address = getsectiondata(imageHeader, "__TEXT", name, &size)
+  guard let nonNullAddress = address else { return nil }
+  guard size != 0 else { return nil }
+  return Section(startAddress: nonNullAddress, size: size)
+}
+
+/// Get the Swift Reflection section locations for a loaded image.
+///
+/// An image of interest must have the following sections in the __DATA
+/// segment:
+/// - __swift5_fieldmd
+/// - __swift5_assocty
+/// - __swift5_builtin
+/// - __swift5_capture
+/// - __swift5_typeref
+/// - __swift5_reflstr (optional, may have been stripped out)
+///
+/// - Parameter i: The index of the loaded image as reported by Dyld.
+/// - Returns: A `ReflectionInfo` containing the locations of all of the
+///   needed sections, or `nil` if the image doesn't contain all of them.
+internal func getReflectionInfoForImage(atIndex i: UInt32) -> ReflectionInfo? {
+  debugLog("BEGIN \(#function)"); defer { debugLog("END \(#function)") }
+  let header = unsafeBitCast(_dyld_get_image_header(i),
+    to: UnsafePointer<MachHeader>.self)
+
+  let imageName = _dyld_get_image_name(i)!
+  let fieldmd = getSectionInfo("__swift5_fieldmd", header)
+  let assocty = getSectionInfo("__swift5_assocty", header)
+  let builtin = getSectionInfo("__swift5_builtin", header)
+  let capture = getSectionInfo("__swift5_capture", header)
+  let typeref = getSectionInfo("__swift5_typeref", header)
+  let reflstr = getSectionInfo("__swift5_reflstr", header)
+  return ReflectionInfo(imageName: String(validatingUTF8: imageName)!,
+                        fieldmd: fieldmd,
+                        assocty: assocty,
+                        builtin: builtin,
+                        capture: capture,
+                        typeref: typeref,
+                        reflstr: reflstr)
+}
+
 /// Get the TEXT segment location and size for a loaded image.
 ///
 /// - Parameter i: The index of the loaded image as reported by Dyld.
@@ -107,6 +194,24 @@ internal func readUInt() -> UInt {
 /// Send all known `ReflectionInfo`s for all images loaded in the current
 /// process.
 internal func sendReflectionInfos() {
+  debugLog("BEGIN \(#function)"); defer { debugLog("END \(#function)") }
+  let infos = (0..<_dyld_image_count()).compactMap(getReflectionInfoForImage)
+
+  var numInfos = infos.count
+  debugLog("\(numInfos) reflection info bundles.")
+  precondition(numInfos >= 1)
+  sendBytes(from: &numInfos, count: MemoryLayout<UInt>.size)
+  for info in infos {
+    debugLog("Sending info for \(info.imageName)")
+    for section in info {
+      sendValue(section?.startAddress)
+      sendValue(section?.size ?? 0)
+    }
+  }
+}
+
+/// Send all loadedimages loaded in the current process.
+internal func sendImages() {
   debugLog("BEGIN \(#function)"); defer { debugLog("END \(#function)") }
   let infos = (0..<_dyld_image_count()).map(getAddressInfoForImage)
 
@@ -200,6 +305,8 @@ internal func reflect(instanceAddress: UInt, kind: InstanceKind) {
       sendValue(instanceAddress)
     case String(validatingUTF8: RequestReflectionInfos)!:
       sendReflectionInfos()
+    case String(validatingUTF8: RequestImages)!:
+      sendImages()
     case String(validatingUTF8: RequestReadBytes)!:
       sendBytes()
     case String(validatingUTF8: RequestSymbolAddress)!:

tools/swift-reflection-test/messages.h

@@ -13,6 +13,7 @@
 static const char *REQUEST_INSTANCE_KIND = "k\n";
 static const char *REQUEST_INSTANCE_ADDRESS = "i\n";
 static const char *REQUEST_REFLECTION_INFO = "r\n";
+static const char *REQUEST_IMAGES = "m\n";
 static const char *REQUEST_READ_BYTES = "b\n";
 static const char *REQUEST_SYMBOL_ADDRESS = "s\n";
 static const char *REQUEST_STRING_LENGTH = "l\n";

tools/swift-reflection-test/swift-reflection-test.c

@@ -34,6 +34,23 @@ typedef struct PipeMemoryReader {
   int from_child[2];
 } PipeMemoryReader;
 
+typedef struct RemoteSection {
+  uintptr_t StartAddress;
+  uintptr_t Size;
+  uintptr_t EndAddress;
+} RemoteSection;
+
+typedef struct RemoteReflectionInfo {
+  RemoteSection fieldmd;
+  RemoteSection assocty;
+  RemoteSection builtin;
+  RemoteSection capture;
+  RemoteSection typeref;
+  RemoteSection reflstr;
+  uintptr_t StartAddress;
+  size_t TotalSize;
+} RemoteReflectionInfo;
+
 static void errorAndExit(const char *message) {
   fprintf(stderr, "%s\n", message);
   abort();
@@ -44,6 +61,77 @@ static void errnoAndExit(const char *message) {
   abort();
 }
 
+static swift_reflection_section_t
+makeLocalSection(void *Buffer, RemoteSection Section,
+                 RemoteReflectionInfo Info) {
+  if (Section.Size == 0) {
+    swift_reflection_section_t LS = {NULL, NULL};
+    return LS;
+  }
+
+  uintptr_t Base
+    = (uintptr_t)Buffer + Section.StartAddress - Info.StartAddress;
+  swift_reflection_section_t LS = {
+    (void *)Base,
+    (void *)(Base + Section.Size)
+  };
+  return LS;
+}
+
+static
+uintptr_t getStartAddress(const RemoteSection Sections[], size_t Count) {
+  uintptr_t Start = 0;
+  for (size_t i = 0; i < Count; ++i) {
+    if (Sections[i].StartAddress != 0) {
+      if (Start != 0)
+        Start = MIN(Start, Sections[i].StartAddress);
+      else
+        Start = Sections[i].StartAddress;
+    }
+  }
+  return Start;
+}
+
+static
+uintptr_t getEndAddress(const RemoteSection Sections[], size_t Count) {
+  uintptr_t End = 0;
+  for (size_t i = 0; i < Count; ++i) {
+    if (Sections[i].StartAddress != 0)
+      End = MAX(End, Sections[i].EndAddress);
+  }
+  return End;
+}
+
+static
+RemoteReflectionInfo makeRemoteReflectionInfo(RemoteSection fieldmd,
+                                              RemoteSection assocty,
+                                              RemoteSection builtin,
+                                              RemoteSection capture,
+                                              RemoteSection typeref,
+                                              RemoteSection reflstr) {
+  RemoteReflectionInfo Info = {
+    fieldmd,
+    assocty,
+    builtin,
+    capture,
+    typeref,
+    reflstr,
+    0,
+    0
+  };
+
+  const RemoteSection Sections[6] = {
+    fieldmd, assocty, builtin, capture, typeref, reflstr
+  };
+
+  Info.StartAddress = getStartAddress(Sections, 6);
+
+  uintptr_t EndAddress = getEndAddress(Sections, 6);
+  Info.TotalSize = EndAddress - Info.StartAddress;
+
+  return Info;
+}
+
 static const size_t ReadEnd = 0;
 static const size_t WriteEnd = 1;
 
@@ -165,11 +253,24 @@ PipeMemoryReader createPipeMemoryReader() {
   return Reader;
 }
 
+static
+RemoteSection makeRemoteSection(const PipeMemoryReader *Reader) {
+  uintptr_t Start;
+  size_t Size;
+
+  PipeMemoryReader_collectBytesFromPipe(Reader, &Start, sizeof(Start));
+  PipeMemoryReader_collectBytesFromPipe(Reader, &Size, sizeof(Size));
+
+  RemoteSection RS = {Start, Size, Start + Size};
+  return RS;
+}
+
+#if defined(__APPLE__) && defined(__MACH__)
 static void
-PipeMemoryReader_receiveReflectionInfo(SwiftReflectionContextRef RC,
+PipeMemoryReader_receiveImages(SwiftReflectionContextRef RC,
                                        const PipeMemoryReader *Reader) {
   int WriteFD = PipeMemoryReader_getParentWriteFD(Reader);
-  write(WriteFD, REQUEST_REFLECTION_INFO, 2);
+  write(WriteFD, REQUEST_IMAGES, 2);
   size_t NumReflectionInfos;
   PipeMemoryReader_collectBytesFromPipe(Reader, &NumReflectionInfos,
                                         sizeof(NumReflectionInfos));
@@ -188,6 +289,64 @@ PipeMemoryReader_receiveReflectionInfo(SwiftReflectionContextRef RC,
 
   free(Images);
 }
+#endif
+
+static void
+PipeMemoryReader_receiveReflectionInfo(SwiftReflectionContextRef RC,
+                                       const PipeMemoryReader *Reader) {
+  int WriteFD = PipeMemoryReader_getParentWriteFD(Reader);
+  write(WriteFD, REQUEST_REFLECTION_INFO, 2);
+  size_t NumReflectionInfos;
+  PipeMemoryReader_collectBytesFromPipe(Reader, &NumReflectionInfos,
+                                        sizeof(NumReflectionInfos));
+
+  if (NumReflectionInfos == 0)
+    return;
+
+  RemoteReflectionInfo *RemoteInfos = calloc(NumReflectionInfos,
+                                             sizeof(RemoteReflectionInfo));
+  if (RemoteInfos == NULL)
+    errnoAndExit("malloc failed");
+
+  for (size_t i = 0; i < NumReflectionInfos; ++i) {
+    RemoteInfos[i] = makeRemoteReflectionInfo(
+      makeRemoteSection(Reader),
+      makeRemoteSection(Reader),
+      makeRemoteSection(Reader),
+      makeRemoteSection(Reader),
+      makeRemoteSection(Reader),
+      makeRemoteSection(Reader));
+  }
+
+  // Now pull in the remote sections into our address space.
+
+  for (size_t i = 0; i < NumReflectionInfos; ++i) {
+    RemoteReflectionInfo RemoteInfo = RemoteInfos[i];
+
+    void *Buffer = malloc(RemoteInfo.TotalSize);
+
+    int Success = PipeMemoryReader_readBytes((void *)Reader,
+                                             RemoteInfo.StartAddress,
+                                             Buffer,
+                                             RemoteInfo.TotalSize);
+    if (!Success)
+      errorAndExit("Couldn't read reflection information");
+
+    swift_reflection_info_t Info = {
+      {makeLocalSection(Buffer, RemoteInfo.fieldmd, RemoteInfo), 0},
+      {makeLocalSection(Buffer, RemoteInfo.assocty, RemoteInfo), 0},
+      {makeLocalSection(Buffer, RemoteInfo.builtin, RemoteInfo), 0},
+      {makeLocalSection(Buffer, RemoteInfo.capture, RemoteInfo), 0},
+      {makeLocalSection(Buffer, RemoteInfo.typeref, RemoteInfo), 0},
+      {makeLocalSection(Buffer, RemoteInfo.reflstr, RemoteInfo), 0},
+      /*LocalStartAddress*/ (uintptr_t) Buffer,
+      /*RemoteStartAddress*/ RemoteInfo.StartAddress,
+    };
+    swift_reflection_addReflectionInfo(RC, Info);
+  }
+
+  free(RemoteInfos);
+}
 
 uint64_t PipeMemoryReader_getStringLength(void *Context, swift_addr_t Address) {
   const PipeMemoryReader *Reader = (const PipeMemoryReader *)Context;
@@ -291,7 +450,11 @@ int doDumpHeapInstance(const char *BinaryFilename) {
       if (PointerSize != sizeof(uintptr_t))
         errorAndExit("Child process had unexpected architecture");
 
+#if defined(__APPLE__) && defined(__MACH__)
+      PipeMemoryReader_receiveImages(RC, &Pipe);
+#else
       PipeMemoryReader_receiveReflectionInfo(RC, &Pipe);
+#endif
 
       while (1) {
         InstanceKind Kind = PipeMemoryReader_receiveInstanceKind(&Pipe);