Merge pull request #72870 from mikeash/libprespecialize-address-range-rejection

[Runtime] Don't attempt to look up prespecialized metadata involving pointers outside the shared cache.

Author: mikeash

stdlib/public/runtime/LibPrespecialized.cpp

@@ -11,6 +11,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "swift/Runtime/LibPrespecialized.h"
+#include "MetadataCache.h"
 #include "Private.h"
 #include "swift/Basic/Lazy.h"
 #include "swift/Runtime/EnvironmentVariables.h"
@@ -75,24 +76,89 @@ static const LibPrespecializedData<InProcess> *findLibPrespecialized() {
   return data;
 }
 
+struct LibPrespecializedState {
+  struct AddressRange {
+    uintptr_t start, end;
+
+    bool contains(const void *ptr) {
+      return start <= (uintptr_t)ptr && (uintptr_t)ptr < end;
+    }
+  };
+
+  bool loggingEnabled;
+  const LibPrespecializedData<InProcess> *data;
+  AddressRange sharedCacheRange{0, 0};
+  AddressRange metadataAllocatorInitialPoolRange{0, 0};
+
+  LibPrespecializedState() {
+    loggingEnabled =
+        runtime::environment::SWIFT_DEBUG_ENABLE_LIB_PRESPECIALIZED_LOGGING();
+    data = findLibPrespecialized();
+
+#if DYLD_GET_SWIFT_PRESPECIALIZED_DATA_DEFINED
+    size_t sharedCacheLength;
+    sharedCacheRange.start =
+        (uintptr_t)_dyld_get_shared_cache_range(&sharedCacheLength);
+    sharedCacheRange.end = sharedCacheRange.start + sharedCacheLength;
+
+    auto [initialPoolStart, initialPoolLength] =
+        MetadataAllocator::InitialPoolLocation();
+    metadataAllocatorInitialPoolRange.start = (uintptr_t)initialPoolStart;
+    metadataAllocatorInitialPoolRange.end =
+        metadataAllocatorInitialPoolRange.start + initialPoolLength;
+#endif
+  }
+};
+
+static Lazy<LibPrespecializedState> LibPrespecialized;
+
+#define LOG(fmt, ...)                                                          \
+  do {                                                                         \
+    if (SWIFT_UNLIKELY(prespecialized.loggingEnabled))                         \
+      fprintf(stderr, "Prespecializations library: " fmt "\n", __VA_ARGS__);   \
+  } while (0)
+
 const LibPrespecializedData<InProcess> *swift::getLibPrespecializedData() {
   return SWIFT_LAZY_CONSTANT(findLibPrespecialized());
 }
 
 // Returns true if the type has any arguments that aren't plain types (packs or
 // unknown kinds).
-static bool hasNonTypeGenericArguments(const TypeContextDescriptor *description) {
-  auto generics = description->getGenericContext();
-  if (!generics)
-    return false;
-
+static bool
+hasNonTypeGenericArguments(const TargetGenericContext<InProcess> *generics) {
   for (auto param : generics->getGenericParams())
     if (param.getKind() != GenericParamKind::Type)
       return true;
 
   return false;
 }
 
+static bool
+isPotentialPrespecializedPointer(LibPrespecializedState &prespecialized,
+                                 const void *pointer) {
+  // Prespecialized metadata descriptors and arguments are always in the shared
+  // cache. They're either statically emitted metadata, or they're
+  // prespecialized metadata. Anything that's dynamically allocated, or
+  // statically allocated outside the shared cache, is not a possible candidate.
+
+  // If we're loading a debug libprespecialized, we can't do these checks, so
+  // just say everything is a potential argument. Performance is not so
+  // important in that case.
+  if (!prespecialized.sharedCacheRange.contains(prespecialized.data))
+    return true;
+
+  // Anything outside the shared cache isn't a potential argument.
+  if (!prespecialized.sharedCacheRange.contains(pointer))
+    return false;
+
+  // Dynamically allocated metadata could be within the shared cache, in the
+  // initial metadata allocation pool. Reject anything in that region.
+  if (prespecialized.metadataAllocatorInitialPoolRange.contains(pointer))
+    return false;
+
+  return true;
+}
+
 static bool disableForValidation = false;
 
 void
@@ -114,24 +180,40 @@ swift::getLibPrespecializedMetadata(const TypeContextDescriptor *description,
         || disablePrespecializedMetadata.load(std::memory_order_acquire)))
     return nullptr;
 
-  auto *data = getLibPrespecializedData();
+  auto &prespecialized = LibPrespecialized.get();
+
+  auto *data = prespecialized.data;
   if (!data)
     return nullptr;
 
+  auto *generics = description->getGenericContext();
+  if (!generics)
+    return nullptr;
+
   // We don't support types with pack parameters yet (and especially not types
   // with unknown parameter kinds) so don't even try to look those up.
-  if (hasNonTypeGenericArguments(description))
+  if (hasNonTypeGenericArguments(generics))
     return nullptr;
 
+  if (!isPotentialPrespecializedPointer(prespecialized, description)) {
+    LOG("Rejecting descriptor %p, not in the shared cache",
+        (const void *)description);
+    return nullptr;
+  }
+
+  auto numKeyArguments = generics->getGenericContextHeader().NumKeyArguments;
+  for (unsigned i = 0; i < numKeyArguments; i++) {
+    if (!isPotentialPrespecializedPointer(prespecialized, arguments[i])) {
+      LOG("Rejecting argument %u %p to descriptor %p, not in the shared cache",
+          i, arguments[i], (const void *)description);
+      return nullptr;
+    }
+  }
+
   Demangler dem;
   auto mangleNode = _buildDemanglingForGenericType(description, arguments, dem);
   if (!mangleNode) {
-    if (SWIFT_UNLIKELY(runtime::environment::
-                           SWIFT_DEBUG_ENABLE_LIB_PRESPECIALIZED_LOGGING()))
-      fprintf(stderr,
-              "Prespecializations library: failed to build demangling with "
-              "descriptor %p.\n",
-              description);
+    LOG("failed to build demangling with descriptor %p.", description);
     return nullptr;
   }
 
@@ -159,10 +241,7 @@ swift::getLibPrespecializedMetadata(const TypeContextDescriptor *description,
   auto *metadataMap = data->getMetadataMap();
   auto *element = metadataMap->find(key.data(), key.size());
   auto *result = element ? element->value : nullptr;
-  if (SWIFT_UNLIKELY(runtime::environment::
-                         SWIFT_DEBUG_ENABLE_LIB_PRESPECIALIZED_LOGGING()))
-    fprintf(stderr, "Prespecializations library: found %p for key '%.*s'.\n",
-            result, (int)key.size(), key.data());
+  LOG("found %p for key '%.*s'.", result, (int)key.size(), key.data());
   return result;
 }
 

stdlib/public/runtime/Metadata.cpp

@@ -7420,6 +7420,10 @@ static swift::atomic<PoolRange>
 AllocationPool{PoolRange{InitialAllocationPool.Pool,
                          sizeof(InitialAllocationPool.Pool)}};
 
+std::tuple<const void *, size_t> MetadataAllocator::InitialPoolLocation() {
+  return {InitialAllocationPool.Pool, sizeof(InitialAllocationPool.Pool)};
+}
+
 bool swift::_swift_debug_metadataAllocationIterationEnabled = false;
 const void * const swift::_swift_debug_allocationPoolPointer = &AllocationPool;
 std::atomic<const void *> swift::_swift_debug_metadataAllocationBacktraceList;

stdlib/public/runtime/MetadataCache.h

@@ -25,6 +25,7 @@
 #include <atomic>
 #include <condition_variable>
 #include <optional>
+#include <tuple>
 
 #ifndef SWIFT_DEBUG_RUNTIME
 #define SWIFT_DEBUG_RUNTIME 0
@@ -46,6 +47,11 @@ class MetadataAllocator : public llvm::AllocatorBase<MetadataAllocator> {
 
   void Reset() {}
 
+  /// Get the location of the allocator's initial statically allocated pool.
+  /// The return values are start and size. If there is no statically allocated
+  /// pool, the return values are NULL, 0.
+  static std::tuple<const void *, size_t> InitialPoolLocation();
+
   SWIFT_RETURNS_NONNULL SWIFT_NODISCARD
   void *Allocate(size_t size, size_t alignment);
   using AllocatorBase<MetadataAllocator>::Allocate;
@@ -67,6 +73,9 @@ class MetadataAllocator : public llvm::AllocatorBase<MetadataAllocator> {
 class MetadataAllocator {
 public:
   MetadataAllocator(uint16_t tag) {}
+  static std::tuple<const void *, size_t> InitialPoolLocation() {
+    return {nullptr, 0};
+  }
   SWIFT_RETURNS_NONNULL SWIFT_NODISCARD
   void *Allocate(size_t size, size_t alignment) {
     if (alignment < sizeof(void*)) alignment = sizeof(void*);