Merge pull request #79773 from mikeash/metadatareader-better-descriptor-sizing

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

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) {