[Sanitizers] Add support for -sanitize=memtag-stack

This sanitizer adds MTE (memory tagging extension) checks to stack variable accesses. It does not have a runtime library component, so some changes have been made to the driver logic to account for this cleanly.

Author: DanBlackwell

Sources/SwiftDriver/Driver/Driver.swift

@@ -2754,12 +2754,17 @@ extension Driver {
       // Support is determined by existence of the sanitizer library.
       // FIXME: Should we do this? This prevents cross-compiling with sanitizers
       //        enabled.
-      var sanitizerSupported = try toolchain.runtimeLibraryExists(
-        for: stableAbi ? .address_stable_abi : sanitizer,
-        targetInfo: targetInfo,
-        parsedOptions: &parsedOptions,
-        isShared: sanitizer != .fuzzer && !stableAbi
-      )
+      var sanitizerSupported = true
+
+      // memtag-stack sanitizer doesn't have a runtime library
+      if sanitizer.hasRuntimeLibrary {
+        sanitizerSupported = try toolchain.runtimeLibraryExists(
+          for: stableAbi ? .address_stable_abi : sanitizer,
+          targetInfo: targetInfo,
+          parsedOptions: &parsedOptions,
+          isShared: sanitizer != .fuzzer && !stableAbi
+        )
+      }
 
       if sanitizer == .thread {
         // TSAN is unavailable on Windows
@@ -2813,6 +2818,16 @@ extension Driver {
       )
     }
 
+    // Address and memtag sanitizers can not be enabled concurrently.
+    if set.contains(.memtag_stack) && set.contains(.address) {
+      diagnosticEngine.emit(
+        .error_argument_not_allowed_with(
+          arg: "-sanitize=memtag-stack",
+          other: "-sanitize=address"
+        )
+      )
+    }
+
     // Scudo can only be run with ubsan.
     if set.contains(.scudo) {
       let allowedSanitizers: Set<Sanitizer> = [.scudo, .undefinedBehavior]

Sources/SwiftDriver/Jobs/Toolchain+LinkerSupport.swift

@@ -65,17 +65,22 @@ extension Toolchain {
     return result
   }
 
+  /// Returns true if a runtime library exists for this sanitizer. Note: some
+  /// sanitizers don't have runtime libraries - ideally this function should
+  /// not be called on them - first check with `sanitizer.hasRuntimeLibrary`
   func runtimeLibraryExists(
     for sanitizer: Sanitizer,
     targetInfo: FrontendTargetInfo,
     parsedOptions: inout ParsedOptions,
     isShared: Bool
   ) throws -> Bool {
-    let runtimeName = try runtimeLibraryName(
+    guard let runtimeName = try runtimeLibraryName(
       for: sanitizer,
       targetTriple: targetInfo.target.triple,
       isShared: isShared
-    )
+    ) else {
+      return false
+    }
     let path = try clangLibraryPath(
       for: targetInfo,
       parsedOptions: &parsedOptions

Sources/SwiftDriver/Toolchains/DarwinToolchain.swift

@@ -168,9 +168,12 @@ public final class DarwinToolchain: Toolchain {
     for sanitizer: Sanitizer,
     targetTriple: Triple,
     isShared: Bool
-  ) throws -> String {
+  ) throws -> String? {
+    guard let libraryName = sanitizer.runtimeLibraryName else {
+      return nil
+    }
     return """
-    libclang_rt.\(sanitizer.libraryName)_\
+    libclang_rt.\(libraryName)_\
     \(targetTriple.darwinPlatform!.libraryNameSuffix)\
     \(isShared ? "_dynamic.dylib" : ".a")
     """

Sources/SwiftDriver/Toolchains/GenericUnixToolchain.swift

@@ -148,9 +148,9 @@ public final class GenericUnixToolchain: Toolchain {
     for sanitizer: Sanitizer,
     targetTriple: Triple,
     isShared: Bool
-  ) throws -> String {
+  ) throws -> String? {
     let environment = (targetTriple.environment == .android) ? "-android" : ""
-    return "libclang_rt.\(sanitizer.libraryName)-\(targetTriple.archName)\(environment).a"
+    return "libclang_rt.\(sanitizer.runtimeLibraryName)-\(targetTriple.archName)\(environment).a"
   }
 
   public func addPlatformSpecificCommonFrontendOptions(

Sources/SwiftDriver/Toolchains/Toolchain.swift

@@ -143,11 +143,12 @@ public protocol Toolchain {
     targetInfo: FrontendTargetInfo
   ) throws -> ResolvedTool
 
+  /// Returns the runtime library name for a given sanitizer (or nil if the sanitizer does not have a runtime library)
   func runtimeLibraryName(
     for sanitizer: Sanitizer,
     targetTriple: Triple,
     isShared: Bool
-  ) throws -> String
+  ) throws -> String?
 
   func platformSpecificInterpreterEnvironmentVariables(
     env: ProcessEnvironmentBlock,

Sources/SwiftDriver/Toolchains/WebAssemblyToolchain.swift

@@ -146,10 +146,10 @@ public final class WebAssemblyToolchain: Toolchain {
     for sanitizer: Sanitizer,
     targetTriple: Triple,
     isShared: Bool
-  ) throws -> String {
+  ) throws -> String? {
     switch sanitizer {
     case .address:
-      return "libclang_rt.\(sanitizer.libraryName)-\(targetTriple.archName).a"
+      return "libclang_rt.\(sanitizer.runtimeLibraryName!)-\(targetTriple.archName).a"
     default:
       throw Error.sanitizersUnsupportedForTarget(targetTriple.triple)
     }

Sources/SwiftDriver/Toolchains/WindowsToolchain.swift

@@ -135,14 +135,17 @@ extension WindowsToolchain.ToolchainValidationError {
   public var globalDebugPathRemapping: String? { nil }
 
   public func runtimeLibraryName(for sanitizer: Sanitizer, targetTriple: Triple,
-                                 isShared: Bool) throws -> String {
+                                 isShared: Bool) throws -> String? {
     // TODO(compnerd) handle shared linking
 
     // FIXME(compnerd) when should `clang_rt.ubsan_standalone_cxx` be used?
     if sanitizer == .undefinedBehavior {
       return "clang_rt.ubsan_standalone.lib"
     }
-    return "clang_rt.\(sanitizer.libraryName).lib"
+    if sanitizer == .memtag_stack {
+      throw ToolchainValidationError.unsupportedSanitizer(sanitizer)
+    }
+    return "clang_rt.\(sanitizer.runtimeLibraryName!).lib"
   }
 
   public func validateArgs(_ parsedOptions: inout ParsedOptions,

Sources/SwiftDriver/Utilities/Sanitizer.swift

@@ -33,15 +33,27 @@ public enum Sanitizer: String, Hashable {
   /// Scudo hardened allocator
   case scudo
 
-  /// The name inside the `compiler_rt` library path (e.g. libclang_rt.{name}.a)
-  var libraryName: String {
+  /// Memory tagging sanitizer (MemTag)
+  case memtag_stack = "memtag-stack"
+
+  /// Does this sanitizer have a runtime library
+  var hasRuntimeLibrary: Bool {
+    if self == .memtag_stack {
+      return false
+    }
+    return true
+  }
+
+  /// The name inside the `compiler_rt` runtime library path (e.g. libclang_rt.{name}.a)
+  var runtimeLibraryName: String? {
     switch self {
     case .address: return "asan"
     case .address_stable_abi: return "asan_abi"
     case .thread: return "tsan"
     case .undefinedBehavior: return "ubsan"
     case .fuzzer: return "fuzzer"
     case .scudo: return "scudo"
+    case .memtag_stack: return nil
     }
   }
 }

Tests/SwiftDriverTests/SwiftDriverTests.swift

@@ -3127,6 +3127,16 @@ final class SwiftDriverTests: XCTestCase {
       XCTAssertJobInvocationMatches(jobs[2], .flag("-fsanitize=undefined"))
     }
 
+    do {
+      // memory tagging stack sanitizer
+      var driver = try Driver(args: commonArgs + ["-sanitize=memtag-stack"])
+      let jobs = try driver.planBuild().removingAutolinkExtractJobs()
+
+      XCTAssertEqual(jobs.count, 3)
+      XCTAssertJobInvocationMatches(jobs[0], .flag("-sanitize=memtag-stack"))
+      // No runtime for memtag-stack - thus no linker arg required
+    }
+
     // FIXME: This test will fail when run on macOS, because the driver uses
     //        the existence of the runtime support libraries to determine if
     //        a sanitizer is supported. Until we allow cross-compiling with