[RemoteMirror] Put a limit on pointer-chasing loops to avoid infinite loops. Report Slab's Next pointer as a separate chunk.

When iterating over linked lists in the target process, bugs or corrupt data might cause us to enter an infinite loop. Limit these loops to a fixed number of iterations so that we fail more gracefully. These limits should be set to a level that no healthy program will approach.

Report Slab's Next pointer as a separate chunk in asyncTaskSlabAllocations. This lets tools see it as a pointer and see the connection to the next slab.

rdar://91289920
(cherry picked from commit 81a6a71)

Author: mikeash

include/swift/Reflection/ReflectionContext.h

@@ -1377,8 +1377,15 @@ class ReflectionContext
       return std::string("failure reading allocation pool contents");
     auto Pool = reinterpret_cast<const PoolRange *>(PoolBytes.get());
 
+    // Limit how many iterations of this loop we'll do, to avoid potential
+    // infinite loops when reading bad data. Limit to 1 million iterations. In
+    // normal operation, each pool allocation is 16kB, so that would be ~16GB of
+    // metadata which is far more than any normal program should have.
+    unsigned LoopCount = 0;
+    unsigned LoopLimit = 1000000;
+
     auto TrailerPtr = Pool->Begin + Pool->Remaining;
-    while (TrailerPtr) {
+    while (TrailerPtr && LoopCount++ < LoopLimit) {
       auto TrailerBytes = getReader()
         .readBytes(RemoteAddress(TrailerPtr), sizeof(PoolTrailer));
       if (!TrailerBytes)
@@ -1426,8 +1433,15 @@ class ReflectionContext
     if (!BacktraceListNextPtr)
       return llvm::None;
 
+    // Limit how many iterations of this loop we'll do, to avoid potential
+    // infinite loops when reading bad data. Limit to 1 billion iterations. In
+    // normal operation, a program shouldn't have anywhere near 1 billion
+    // metadata allocations.
+    unsigned LoopCount = 0;
+    unsigned LoopLimit = 1000000000;
+
     auto BacktraceListNext = BacktraceListNextPtr->getResolvedAddress();
-    while (BacktraceListNext) {
+    while (BacktraceListNext && LoopCount++ < LoopLimit) {
       auto HeaderBytes = getReader().readBytes(
           RemoteAddress(BacktraceListNext),
           sizeof(MetadataAllocationBacktraceHeader<Runtime>));
@@ -1473,30 +1487,40 @@ class ReflectionContext
     // provide the whole thing as one big chunk.
     size_t HeaderSize =
         llvm::alignTo(sizeof(*Slab), llvm::Align(MaximumAlignment));
-    AsyncTaskAllocationChunk Chunk;
 
-    Chunk.Start = SlabPtr + HeaderSize;
-    Chunk.Length = Slab->CurrentOffset;
-    Chunk.Kind = AsyncTaskAllocationChunk::ChunkKind::Unknown;
+    AsyncTaskAllocationChunk AllocatedSpaceChunk;
+    AllocatedSpaceChunk.Start = SlabPtr + HeaderSize;
+    AllocatedSpaceChunk.Length = Slab->CurrentOffset;
+    AllocatedSpaceChunk.Kind = AsyncTaskAllocationChunk::ChunkKind::Unknown;
+
+    // Provide a second chunk just for the Next pointer, so the client knows
+    // that there's an allocation there.
+    AsyncTaskAllocationChunk NextPtrChunk;
+    NextPtrChunk.Start =
+        SlabPtr + offsetof(typename StackAllocator::Slab, Next);
+    NextPtrChunk.Length = sizeof(Slab->Next);
+    NextPtrChunk.Kind = AsyncTaskAllocationChunk::ChunkKind::RawPointer;
 
     // Total slab size is the slab's capacity plus the header.
     StoredPointer SlabSize = Slab->Capacity + HeaderSize;
 
-    return {llvm::None, {Slab->Next, SlabSize, {Chunk}}};
+    return {llvm::None,
+            {Slab->Next, SlabSize, {NextPtrChunk, AllocatedSpaceChunk}}};
   }
 
   std::pair<llvm::Optional<std::string>, AsyncTaskInfo>
-  asyncTaskInfo(StoredPointer AsyncTaskPtr) {
+  asyncTaskInfo(StoredPointer AsyncTaskPtr, unsigned ChildTaskLimit,
+                unsigned AsyncBacktraceLimit) {
     loadTargetPointers();
 
     if (supportsPriorityEscalation)
       return asyncTaskInfo<
           AsyncTask<Runtime, ActiveTaskStatusWithEscalation<Runtime>>>(
-          AsyncTaskPtr);
+          AsyncTaskPtr, ChildTaskLimit, AsyncBacktraceLimit);
     else
       return asyncTaskInfo<
           AsyncTask<Runtime, ActiveTaskStatusWithoutEscalation<Runtime>>>(
-          AsyncTaskPtr);
+          AsyncTaskPtr, ChildTaskLimit, AsyncBacktraceLimit);
   }
 
   std::pair<llvm::Optional<std::string>, ActorInfo>
@@ -1588,7 +1612,8 @@ class ReflectionContext
 
   template <typename AsyncTaskType>
   std::pair<llvm::Optional<std::string>, AsyncTaskInfo>
-  asyncTaskInfo(StoredPointer AsyncTaskPtr) {
+  asyncTaskInfo(StoredPointer AsyncTaskPtr, unsigned ChildTaskLimit,
+                unsigned AsyncBacktraceLimit) {
     auto AsyncTaskObj = readObj<AsyncTaskType>(AsyncTaskPtr);
     if (!AsyncTaskObj)
       return {std::string("failure reading async task"), {}};
@@ -1620,8 +1645,9 @@ class ReflectionContext
     Info.RunJob = getRunJob(AsyncTaskObj.get());
 
     // Find all child tasks.
+    unsigned ChildTaskLoopCount = 0;
     auto RecordPtr = AsyncTaskObj->PrivateStorage.Status.Record;
-    while (RecordPtr) {
+    while (RecordPtr && ChildTaskLoopCount++ < ChildTaskLimit) {
       auto RecordObj = readObj<TaskStatusRecord<Runtime>>(RecordPtr);
       if (!RecordObj)
         break;
@@ -1661,7 +1687,8 @@ class ReflectionContext
     // Walk the async backtrace.
     if (Info.HasIsRunning && !Info.IsRunning) {
       auto ResumeContext = AsyncTaskObj->ResumeContextAndReserved[0];
-      while (ResumeContext) {
+      unsigned AsyncBacktraceLoopCount = 0;
+      while (ResumeContext && AsyncBacktraceLoopCount++ < AsyncBacktraceLimit) {
         auto ResumeContextObj = readObj<AsyncContext<Runtime>>(ResumeContext);
         if (!ResumeContextObj)
           break;

stdlib/public/SwiftRemoteMirror/SwiftRemoteMirror.cpp

@@ -54,7 +54,7 @@ struct SwiftReflectionContext {
     auto Reader = std::make_shared<CMemoryReader>(impl);
     nativeContext = new NativeReflectionContext(Reader);
   }
-  
+
   ~SwiftReflectionContext() {
     freeTemporaryAllocation();
     delete nativeContext;
@@ -203,9 +203,9 @@ ReflectionSection<Iterator> sectionFromInfo(const swift_reflection_info_t &Info,
   auto RemoteSectionStart = (uint64_t)(uintptr_t)Section.section.Begin
     - Info.LocalStartAddress
     + Info.RemoteStartAddress;
-  
+
   auto Start = RemoteRef<void>(RemoteSectionStart, Section.section.Begin);
-  
+
   return ReflectionSection<Iterator>(Start,
              (uintptr_t)Section.section.End - (uintptr_t)Section.section.Begin);
 }
@@ -225,7 +225,7 @@ void
 swift_reflection_addReflectionInfo(SwiftReflectionContextRef ContextRef,
                                    swift_reflection_info_t Info) {
   auto Context = ContextRef->nativeContext;
-  
+
   // The `offset` fields must be zero.
   if (Info.field.offset != 0
       || Info.associated_types.offset != 0
@@ -246,7 +246,7 @@ swift_reflection_addReflectionInfo(SwiftReflectionContextRef ContextRef,
     sectionFromInfo<const void *>(Info, Info.reflection_strings),
     ReflectionSection<const void *>(nullptr, 0),
     ReflectionSection<MultiPayloadEnumDescriptorIterator>(0, 0)};
-  
+
   Context->addReflectionInfo(ContextInfo);
 }
 
@@ -821,10 +821,23 @@ const char *swift_reflection_iterateMetadataAllocationBacktraces(
 swift_async_task_slab_return_t
 swift_reflection_asyncTaskSlabPointer(SwiftReflectionContextRef ContextRef,
                                       swift_reflection_ptr_t AsyncTaskPtr) {
-  auto Info = swift_reflection_asyncTaskInfo(ContextRef, AsyncTaskPtr);
+  auto Context = ContextRef->nativeContext;
+
+  // We only care about the AllocatorSlabPtr field. Disable child task and async
+  // backtrace iteration to save wasted work.
+  unsigned ChildTaskLimit = 0;
+  unsigned AsyncBacktraceLimit = 0;
+
+  llvm::Optional<std::string> Error;
+  NativeReflectionContext::AsyncTaskInfo TaskInfo;
+  std::tie(Error, TaskInfo) =
+      Context->asyncTaskInfo(AsyncTaskPtr, ChildTaskLimit, AsyncBacktraceLimit);
+
   swift_async_task_slab_return_t Result = {};
-  Result.Error = Info.Error;
-  Result.SlabPtr = Info.AllocatorSlabPtr;
+  if (Error) {
+    Result.Error = returnableCString(ContextRef, Error);
+  }
+  Result.SlabPtr = TaskInfo.AllocatorSlabPtr;
   return Result;
 }
 
@@ -866,9 +879,16 @@ swift_async_task_info_t
 swift_reflection_asyncTaskInfo(SwiftReflectionContextRef ContextRef,
                                swift_reflection_ptr_t AsyncTaskPtr) {
   auto Context = ContextRef->nativeContext;
+
+  // Limit the child task and async backtrace iteration to semi-reasonable
+  // numbers to avoid doing excessive work on bad data.
+  unsigned ChildTaskLimit = 1000000;
+  unsigned AsyncBacktraceLimit = 1000;
+
   llvm::Optional<std::string> Error;
   NativeReflectionContext::AsyncTaskInfo TaskInfo;
-  std::tie(Error, TaskInfo) = Context->asyncTaskInfo(AsyncTaskPtr);
+  std::tie(Error, TaskInfo) =
+      Context->asyncTaskInfo(AsyncTaskPtr, ChildTaskLimit, AsyncBacktraceLimit);
 
   swift_async_task_info_t Result = {};
   if (Error) {