[Reflection] Fix undersized reads in MetadataReader::readContextDescriptor.

Determining a descriptor's size requires reading its contents, but reading its contents from out of process requires knowing its size.

Build up the size incrementally by walking over the TrailingObjects. Take advantage of the fact that each trailing object's presence/count depends only on data that comes before it. This allows us to read prefixes that we gradually expand until we've covered the whole thing.

Add calls to TrailingObjects to allow iterating over the prefix sizes, and modify readContextDescriptor to use them. This replaces the old code which attempted to determine the descriptor size in an ad-hoc fashion that didn't always get it right.

rdar://146006006

Author: mikeash

include/swift/ABI/GenericContext.h

@@ -711,6 +711,14 @@ class TrailingGenericContextObjects<TargetSelf<Runtime>,
             getGenericValueDescriptors().data()};
   }
 
+  static size_t trailingTypeCount() {
+    return TrailingObjects::trailingTypeCount();
+  }
+
+  size_t sizeWithTrailingTypeCount(size_t n) const {
+    return TrailingObjects::sizeWithTrailingTypeCount(n);
+  }
+
 protected:
   size_t numTrailingObjects(OverloadToken<GenericContextHeaderType>) const {
     return asSelf()->isGeneric() ? 1 : 0;

include/swift/ABI/Metadata.h

@@ -3168,7 +3168,7 @@ struct swift_ptrauth_struct_context_descriptor(AnonymousContextDescriptor)
   using TrailingGenericContextObjects::numTrailingObjects;
 
   size_t numTrailingObjects(OverloadToken<MangledContextName>) const {
-    return this->hasMangledNam() ? 1 : 0;
+    return this->hasMangledName() ? 1 : 0;
   }
 
 public:

include/swift/ABI/TrailingObjects.h

@@ -207,6 +207,21 @@ class TrailingObjectsImpl<Align, BaseTy, TopTrailingObj, PrevTy, NextTy,
             sizeof(NextTy) * Count1,
         MoreCounts...);
   }
+
+  // Helper function for TrailingObjects::sizeWithTrailingTypeCount. This
+  // recurses to superclasses until n reaches zero, then computes the size of
+  // the object up to that point.
+  static size_t sizeWithTrailingTypeCountImpl(const BaseTy *Obj, size_t n) {
+    if (n > 0)
+      return ParentType::sizeWithTrailingTypeCountImpl(Obj, n - 1);
+
+    auto *Ptr = getTrailingObjectsImpl(
+        Obj, TrailingObjectsBase::OverloadToken<NextTy>());
+    auto Count = TopTrailingObj::callNumTrailingObjects(
+        Obj, TrailingObjectsBase::OverloadToken<NextTy>());
+    auto *End = Ptr + Count;
+    return (const char *)End - (const char *)Obj;
+  }
 };
 
 // The base case of the TrailingObjectsImpl inheritance recursion,
@@ -225,6 +240,10 @@ class TrailingObjectsImpl<Align, BaseTy, TopTrailingObj, PrevTy>
   }
 
   template <bool CheckAlignment> static void verifyTrailingObjectsAlignment() {}
+
+  static size_t sizeWithTrailingTypeCountImpl(const BaseTy *Obj, size_t n) {
+    return 0;
+  }
 };
 
 } // end namespace trailing_objects_internal
@@ -386,6 +405,20 @@ class swift_ptrauth_struct_derived(BaseTy) TrailingObjects
 
     BaseTy *const p;
   };
+
+  // Get the number of trailing types in this TrailingObjects specialization.
+  static size_t trailingTypeCount() { return sizeof...(TrailingTys); }
+
+  // Get the size of the object including trailing objects through index N. This
+  // allows working out the size of a TrailingObjects subclass incrementally,
+  // by calling this repeatedly starting from 0. This is needed for remote
+  // inspection, which needs to figure out how much memory to read just from the
+  // contents of the object. It can repeatedly read a prefix until it has the
+  // whole thing.
+  size_t sizeWithTrailingTypeCount(size_t n) const {
+    return ParentType::sizeWithTrailingTypeCountImpl(
+        static_cast<const BaseTy *>(this), n);
+  }
 };
 
 } // end namespace ABI

include/swift/Remote/MetadataReader.h

@@ -18,6 +18,7 @@
 #define SWIFT_REMOTE_METADATAREADER_H
 
 
+#include "swift/ABI/Metadata.h"
 #include "swift/Runtime/Metadata.h"
 #include "swift/Remote/MemoryReader.h"
 #include "swift/Demangling/Demangler.h"
@@ -1451,135 +1452,107 @@ class MetadataReader {
     if (address == 0)
       return nullptr;
 
+    auto remoteAddress = RemoteAddress(address);
+    auto ptr = Reader->readBytes(remoteAddress,
+                                 sizeof(TargetContextDescriptor<Runtime>));
+    if (!ptr)
+      return nullptr;
+
     auto cached = ContextDescriptorCache.find(address);
     if (cached != ContextDescriptorCache.end())
       return ContextDescriptorRef(
           address, reinterpret_cast<const TargetContextDescriptor<Runtime> *>(
                        cached->second.get()));
 
-    // Read the flags to figure out how much space we should read.
-    ContextDescriptorFlags flags;
-    if (!Reader->readBytes(RemoteAddress(address), (uint8_t*)&flags,
-                           sizeof(flags)))
-      return nullptr;
-    
-    TypeContextDescriptorFlags typeFlags(flags.getKindSpecificFlags());
-    uint64_t baseSize = 0;
-    uint64_t genericHeaderSize = sizeof(GenericContextDescriptorHeader);
-    uint64_t metadataInitSize = 0;
-    bool hasVTable = false;
-
-    auto readMetadataInitSize = [&]() -> unsigned {
-      switch (typeFlags.getMetadataInitialization()) {
-      case TypeContextDescriptorFlags::NoMetadataInitialization:
-        return 0;
-      case TypeContextDescriptorFlags::SingletonMetadataInitialization:
-        // FIXME: classes
-        return sizeof(TargetSingletonMetadataInitialization<Runtime>);
-      case TypeContextDescriptorFlags::ForeignMetadataInitialization:
-        return sizeof(TargetForeignMetadataInitialization<Runtime>);
-      }
-      return 0;
-    };
-
-    switch (flags.getKind()) {
+    bool success = false;
+    switch (
+        reinterpret_cast<const TargetContextDescriptor<Runtime> *>(ptr.get())
+            ->getKind()) {
     case ContextDescriptorKind::Module:
-      baseSize = sizeof(TargetModuleContextDescriptor<Runtime>);
+      ptr = Reader->readBytes(remoteAddress,
+                              sizeof(TargetModuleContextDescriptor<Runtime>));
+      success = ptr != nullptr;
       break;
-    // TODO: Should we include trailing generic arguments in this load?
     case ContextDescriptorKind::Extension:
-      baseSize = sizeof(TargetExtensionContextDescriptor<Runtime>);
+      success =
+          readFullContextDescriptor<TargetExtensionContextDescriptor<Runtime>>(
+              remoteAddress, ptr);
       break;
     case ContextDescriptorKind::Anonymous:
-      baseSize = sizeof(TargetAnonymousContextDescriptor<Runtime>);
-      if (AnonymousContextDescriptorFlags(flags.getKindSpecificFlags())
-            .hasMangledName()) {
-        metadataInitSize = sizeof(TargetMangledContextName<Runtime>);
-      }
+      success =
+          readFullContextDescriptor<TargetAnonymousContextDescriptor<Runtime>>(
+              remoteAddress, ptr);
       break;
     case ContextDescriptorKind::Class:
-      baseSize = sizeof(TargetClassDescriptor<Runtime>);
-      genericHeaderSize = sizeof(TypeGenericContextDescriptorHeader);
-      hasVTable = typeFlags.class_hasVTable();
-      metadataInitSize = readMetadataInitSize();
+      success = readFullContextDescriptor<TargetClassDescriptor<Runtime>>(
+          remoteAddress, ptr);
       break;
     case ContextDescriptorKind::Enum:
-      baseSize = sizeof(TargetEnumDescriptor<Runtime>);
-      genericHeaderSize = sizeof(TypeGenericContextDescriptorHeader);
-      metadataInitSize = readMetadataInitSize();
+      success = readFullContextDescriptor<TargetEnumDescriptor<Runtime>>(
+          remoteAddress, ptr);
       break;
     case ContextDescriptorKind::Struct:
-      baseSize = sizeof(TargetStructDescriptor<Runtime>);
-      genericHeaderSize = sizeof(TypeGenericContextDescriptorHeader);
-      metadataInitSize = readMetadataInitSize();
+      success = readFullContextDescriptor<TargetStructDescriptor<Runtime>>(
+          remoteAddress, ptr);
       break;
     case ContextDescriptorKind::Protocol:
-      baseSize = sizeof(TargetProtocolDescriptor<Runtime>);
+      success = readFullContextDescriptor<TargetProtocolDescriptor<Runtime>>(
+          remoteAddress, ptr);
       break;
     case ContextDescriptorKind::OpaqueType:
-      baseSize = sizeof(TargetOpaqueTypeDescriptor<Runtime>);
-      metadataInitSize =
-        sizeof(typename Runtime::template RelativeDirectPointer<const char>)
-          * flags.getKindSpecificFlags();
+      success = readFullContextDescriptor<TargetOpaqueTypeDescriptor<Runtime>>(
+          remoteAddress, ptr);
       break;
     default:
       // We don't know about this kind of context.
       return nullptr;
     }
-
-    // Determine the full size of the descriptor. This is reimplementing a fair
-    // bit of TrailingObjects but for out-of-process; maybe there's a way to
-    // factor the layout stuff out...
-    uint64_t genericsSize = 0;
-    if (flags.isGeneric()) {
-      GenericContextDescriptorHeader header;
-      auto headerAddr = address
-        + baseSize
-        + genericHeaderSize
-        - sizeof(header);
-      
-      if (!Reader->readBytes(RemoteAddress(headerAddr),
-                             (uint8_t*)&header, sizeof(header)))
-        return nullptr;
-      
-      genericsSize = genericHeaderSize
-        + (header.NumParams + 3u & ~3u)
-        + header.NumRequirements
-          * sizeof(TargetGenericRequirementDescriptor<Runtime>);
-    }
-
-    uint64_t vtableSize = 0;
-    if (hasVTable) {
-      TargetVTableDescriptorHeader<Runtime> header;
-      auto headerAddr = address
-        + baseSize
-        + genericsSize
-        + metadataInitSize;
-      
-      if (!Reader->readBytes(RemoteAddress(headerAddr),
-                             (uint8_t*)&header, sizeof(header)))
-        return nullptr;
-
-      vtableSize = sizeof(header)
-        + header.VTableSize * sizeof(TargetMethodDescriptor<Runtime>);
-    }
-
-    uint64_t size = baseSize + genericsSize + metadataInitSize + vtableSize;
-    if (size > MaxMetadataSize)
-      return nullptr;
-    auto readResult = Reader->readBytes(RemoteAddress(address), size);
-    if (!readResult)
+    if (!success)
       return nullptr;
 
-    auto descriptor =
-        reinterpret_cast<const TargetContextDescriptor<Runtime> *>(
-            readResult.get());
-
-    ContextDescriptorCache.insert(
-        std::make_pair(address, std::move(readResult)));
+    auto *descriptor =
+        reinterpret_cast<const TargetContextDescriptor<Runtime> *>(ptr.get());
+    ContextDescriptorCache.insert(std::make_pair(address, std::move(ptr)));
     return ContextDescriptorRef(address, descriptor);
   }
-  
+
+  template <typename DescriptorTy>
+  bool readFullContextDescriptor(RemoteAddress address,
+                                 MemoryReader::ReadBytesResult &ptr) {
+    // Read the full base descriptor if it's bigger than what we have so far.
+    if (sizeof(DescriptorTy) > sizeof(TargetContextDescriptor<Runtime>)) {
+      ptr = Reader->readObj<DescriptorTy>(address);
+      if (!ptr)
+        return false;
+    }
+
+    // We don't know how much memory we need to read to get all the trailing
+    // objects, but we need to read the memory to figure out how much memory we
+    // need to read. Handle this by reading incrementally.
+    //
+    // We rely on the fact that each trailing object's count depends only on
+    // that comes before it. If we've read the first N trailing objects, then we
+    // can safely compute the size with N+1 trailing objects. If that size is
+    // bigger than what we've read so far, re-read the descriptor with the new
+    // size. Once we've walked through all the trailing objects, we've read
+    // everything.
+
+    size_t sizeSoFar = sizeof(DescriptorTy);
+
+    for (size_t i = 0; i < DescriptorTy::trailingTypeCount(); i++) {
+      const DescriptorTy *descriptorSoFar =
+          reinterpret_cast<const DescriptorTy *>(ptr.get());
+      size_t thisSize = descriptorSoFar->sizeWithTrailingTypeCount(i);
+      if (thisSize > sizeSoFar) {
+        ptr = Reader->readBytes(address, thisSize);
+        if (!ptr)
+          return false;
+        sizeSoFar = thisSize;
+      }
+    }
+    return true;
+  }
+
   /// Demangle the entity represented by a symbolic reference to a given symbol name.
   Demangle::NodePointer
   buildContextManglingForSymbol(StringRef symbol, Demangler &dem) {