[Sanitizers] Add Driver/Frontend option to enable sanitizer instrumentation that supports error recovery

include/swift/AST/DiagnosticsFrontend.def

@@ -71,6 +71,10 @@ ERROR(error_option_requires_sanitizer, none,
       "option '%0' requires a sanitizer to be enabled. Use -sanitize= to "
       "enable a sanitizer", (StringRef))
 
+WARNING(warning_option_requires_specific_sanitizer, none,
+      "option '%0' has no effect when '%1' sanitizer is disabled. Use -sanitize=%1 to "
+      "enable the sanitizer", (StringRef, StringRef))
+
 ERROR(error_option_missing_required_argument, none,
       "option '%0' is missing a required argument (%1)", (StringRef, StringRef))
 

include/swift/AST/IRGenOptions.h

@@ -106,6 +106,9 @@ class IRGenOptions {
   /// Which sanitizer is turned on.
   OptionSet<SanitizerKind> Sanitizers;
 
+  /// Which sanitizer(s) have recovery instrumentation enabled.
+  OptionSet<SanitizerKind> SanitizersWithRecoveryInstrumentation;
+
   /// Path prefixes that should be rewritten in debug info.
   PathRemapper DebugPrefixMap;
 
@@ -239,6 +242,7 @@ class IRGenOptions {
       : DWARFVersion(2), OutputKind(IRGenOutputKind::LLVMAssembly),
         Verify(true), OptMode(OptimizationMode::NotSet),
         Sanitizers(OptionSet<SanitizerKind>()),
+        SanitizersWithRecoveryInstrumentation(OptionSet<SanitizerKind>()),
         DebugInfoLevel(IRGenDebugInfoLevel::None),
         DebugInfoFormat(IRGenDebugInfoFormat::None),
         DisableClangModuleSkeletonCUs(false), UseJIT(false),

include/swift/Option/Options.td

@@ -889,6 +889,15 @@ def sanitize_EQ : CommaJoined<["-"], "sanitize=">,
   Flags<[FrontendOption, NoInteractiveOption]>, MetaVarName<"<check>">,
   HelpText<"Turn on runtime checks for erroneous behavior.">;
 
+def sanitize_recover_EQ
+    : CommaJoined<["-"], "sanitize-recover=">,
+      Flags<[FrontendOption, NoInteractiveOption]>,
+      MetaVarName<"<check>">,
+      HelpText<"Specify which sanitizer runtime checks (see -sanitize=) will "
+               "generate instrumentation that allows error recovery. Listed "
+               "checks should be comma separated. Default behavior is to not "
+               "allow error recovery.">;
+
 def sanitize_coverage_EQ : CommaJoined<["-"], "sanitize-coverage=">,
   Flags<[FrontendOption, NoInteractiveOption]>,
   MetaVarName<"<type>">,

include/swift/Option/SanitizerOptions.h

@@ -36,6 +36,10 @@ OptionSet<SanitizerKind> parseSanitizerArgValues(
     const llvm::Triple &Triple, DiagnosticEngine &Diag,
     llvm::function_ref<bool(llvm::StringRef, bool)> sanitizerRuntimeLibExists);
 
+OptionSet<SanitizerKind> parseSanitizerRecoverArgValues(
+    const llvm::opt::Arg *A, const OptionSet<SanitizerKind> &enabledSanitizers,
+    DiagnosticEngine &Diags, bool emitWarnings);
+
 /// Parses a -sanitize-coverage= argument's value.
 llvm::SanitizerCoverageOptions parseSanitizerCoverageArgValue(
         const llvm::opt::Arg *A,

lib/Driver/Driver.cpp

@@ -1636,6 +1636,14 @@ void Driver::buildOutputInfo(const ToolChain &TC, const DerivedArgList &Args,
           return TC.sanitizerRuntimeLibExists(Args, sanitizerName, shared);
         });
 
+  if (const Arg *A = Args.getLastArg(options::OPT_sanitize_recover_EQ)) {
+    // Just validate the args. The frontend will parse these again and actually
+    // use them. To avoid emitting warnings multiple times we surpress warnings
+    // here but not in the frontend.
+    (void)parseSanitizerRecoverArgValues(A, OI.SelectedSanitizers, Diags,
+                                         /*emitWarnings=*/false);
+  }
+
   if (const Arg *A = Args.getLastArg(options::OPT_sanitize_coverage_EQ)) {
 
     // Check that the sanitizer coverage flags are supported if supplied.

lib/Driver/ToolChains.cpp

@@ -214,6 +214,7 @@ static void addCommonFrontendArgs(const ToolChain &TC, const OutputInfo &OI,
   inputArgs.AddLastArg(arguments, options::OPT_profile_coverage_mapping);
   inputArgs.AddLastArg(arguments, options::OPT_warnings_as_errors);
   inputArgs.AddLastArg(arguments, options::OPT_sanitize_EQ);
+  inputArgs.AddLastArg(arguments, options::OPT_sanitize_recover_EQ);
   inputArgs.AddLastArg(arguments, options::OPT_sanitize_coverage_EQ);
   inputArgs.AddLastArg(arguments, options::OPT_static);
   inputArgs.AddLastArg(arguments, options::OPT_swift_version);

lib/Frontend/CompilerInvocation.cpp

@@ -877,6 +877,12 @@ static bool ParseSILArgs(SILOptions &Opts, ArgList &Args,
     IRGenOpts.Sanitizers = Opts.Sanitizers;
   }
 
+  if (const Arg *A = Args.getLastArg(options::OPT_sanitize_recover_EQ)) {
+    IRGenOpts.SanitizersWithRecoveryInstrumentation =
+        parseSanitizerRecoverArgValues(A, Opts.Sanitizers, Diags,
+                                       /*emitWarnings=*/true);
+  }
+
   if (auto A = Args.getLastArg(OPT_enable_verify_exclusivity,
                                OPT_disable_verify_exclusivity)) {
     Opts.VerifyExclusivity

lib/IRGen/IRGen.cpp

@@ -126,8 +126,14 @@ static void addSwiftMergeFunctionsPass(const PassManagerBuilder &Builder,
 
 static void addAddressSanitizerPasses(const PassManagerBuilder &Builder,
                                       legacy::PassManagerBase &PM) {
-  PM.add(createAddressSanitizerFunctionPass());
-  PM.add(createModuleAddressSanitizerLegacyPassPass());
+  auto &BuilderWrapper =
+      static_cast<const PassManagerBuilderWrapper &>(Builder);
+  auto recover =
+      bool(BuilderWrapper.IRGOpts.SanitizersWithRecoveryInstrumentation &
+           SanitizerKind::Address);
+  PM.add(createAddressSanitizerFunctionPass(/*CompileKernel=*/false, recover));
+  PM.add(createModuleAddressSanitizerLegacyPassPass(/*CompileKernel=*/false,
+                                                    recover));
 }
 
 static void addThreadSanitizerPass(const PassManagerBuilder &Builder,

lib/Option/SanitizerOptions.cpp

@@ -185,6 +185,49 @@ OptionSet<SanitizerKind> swift::parseSanitizerArgValues(
   return sanitizerSet;
 }
 
+OptionSet<SanitizerKind> swift::parseSanitizerRecoverArgValues(
+    const llvm::opt::Arg *A, const OptionSet<SanitizerKind> &enabledSanitizers,
+    DiagnosticEngine &Diags, bool emitWarnings) {
+  OptionSet<SanitizerKind> sanitizerRecoverSet;
+
+  // Find the sanitizer kind.
+  for (const char *arg : A->getValues()) {
+    Optional<SanitizerKind> optKind = parse(arg);
+
+    // Unrecognized sanitizer option
+    if (!optKind.hasValue()) {
+      Diags.diagnose(SourceLoc(), diag::error_unsupported_option_argument,
+                     A->getOption().getPrefixedName(), arg);
+      continue;
+    }
+    SanitizerKind kind = optKind.getValue();
+
+    // Only support ASan for now.
+    if (kind != SanitizerKind::Address) {
+      Diags.diagnose(SourceLoc(), diag::error_unsupported_option_argument,
+                     A->getOption().getPrefixedName(), arg);
+      continue;
+    }
+
+    // Check that the sanitizer is enabled.
+    if (!(enabledSanitizers & kind)) {
+      SmallString<128> b;
+      if (emitWarnings) {
+        Diags.diagnose(SourceLoc(),
+                       diag::warning_option_requires_specific_sanitizer,
+                       (A->getOption().getPrefixedName() + toStringRef(kind))
+                           .toStringRef(b),
+                       toStringRef(kind));
+      }
+      continue;
+    }
+
+    sanitizerRecoverSet |= kind;
+  }
+
+  return sanitizerRecoverSet;
+}
+
 std::string swift::getSanitizerList(const OptionSet<SanitizerKind> &Set) {
   std::string list;
   #define SANITIZER(_, kind, name, file) \

test/Driver/sanitize_recover.swift

@@ -0,0 +1,18 @@
+// RUN: not %swiftc_driver -driver-print-jobs -sanitize=address -sanitize-recover=foo %s 2>&1 | %FileCheck -check-prefix=SAN_RECOVER_INVALID_ARG %s
+// RUN: not %swiftc_driver -driver-print-jobs -sanitize=address -sanitize-recover=thread %s 2>&1 | %FileCheck -check-prefix=SAN_RECOVER_UNSUPPORTED_ARG %s
+// RUN: %swiftc_driver -v -sanitize-recover=address %s -o %t 2>&1 | %FileCheck -check-prefix=SAN_RECOVER_MISSING_INSTRUMENTATION_OPTION %s
+// RUN: %swiftc_driver -driver-print-jobs -sanitize=address -sanitize-recover=address %s 2>&1 | %FileCheck -check-prefix=ASAN_WITH_RECOVER  %s
+// RUN: %swiftc_driver -driver-print-jobs -sanitize=address %s 2>&1 | %FileCheck -check-prefix=ASAN_WITHOUT_RECOVER --implicit-check-not='-sanitize-recover=address' %s
+
+// SAN_RECOVER_INVALID_ARG: unsupported argument 'foo' to option '-sanitize-recover='
+// SAN_RECOVER_UNSUPPORTED_ARG: unsupported argument 'thread' to option '-sanitize-recover='
+
+// SAN_RECOVER_MISSING_INSTRUMENTATION_OPTION: warning: option '-sanitize-recover=address' has no effect when 'address' sanitizer is disabled. Use -sanitize=address to enable the sanitizer
+// SAN_RECOVER_MISSING_INSTRUMENTATION_OPTION-NOT: warning: option '-sanitize-recover=address' has no effect when 'address' sanitizer is disabled. Use -sanitize=address to enable the sanitizer
+
+// ASAN_WITH_RECOVER: swift
+// ASAN_WITH_RECOVER-DAG: -sanitize=address
+// ASAN_WITH_RECOVER-DAG: -sanitize-recover=address
+
+// ASAN_WITHOUT_RECOVER: swift
+// ASAN_WITHOUT_RECOVER: -sanitize=address

test/IRGen/address_sanitizer_recover.swift

@@ -0,0 +1,16 @@
+// RUN: %target-swift-frontend -emit-ir -sanitize=address -sanitize-recover=address %s | %FileCheck %s -check-prefix=ASAN_RECOVER
+// RUN: %target-swift-frontend -emit-ir -sanitize=address  %s | %FileCheck %s -check-prefix=ASAN_NO_RECOVER
+// RUN: %target-swift-frontend -emit-ir -sanitize-recover=address %s | %FileCheck %s -check-prefix=NO_ASAN_RECOVER
+
+// ASAN_RECOVER: declare void @__asan_loadN_noabort(
+// ASAN_NO_RECOVER: declare void @__asan_loadN(
+
+let size:Int = 128;
+let x = UnsafeMutablePointer<UInt8>.allocate(capacity: size)
+x.initialize(repeating: 0, count: size)
+x.advanced(by: 0).pointee = 5;
+print("Read first element:\(x.advanced(by: 0).pointee)")
+x.deallocate();
+
+// There should be no ASan instrumentation in this case.
+// NO_ASAN_RECOVER-NOT: declare void @__asan_load

test/Sanitizers/asan_interface.h

@@ -0,0 +1,4 @@
+// This file is a swift bridging header to ASan's interface. It exists so
+// we don't need to worry about where the header lives and instead let Clang
+// figure out where the header lives.
+#include "sanitizer/asan_interface.h"

test/Sanitizers/asan_recover.swift

@@ -0,0 +1,98 @@
+// REQUIRES: executable_test
+// REQUIRES: asan_runtime
+// UNSUPPORTED: windows
+
+// Check with recovery instrumentation and runtime option to continue execution.
+// RUN: %target-swiftc_driver %s -target %sanitizers-target-triple -g -sanitize=address -sanitize-recover=address -import-objc-header %S/asan_interface.h -emit-ir -o %t.asan_recover.ll
+// RUN: %FileCheck -check-prefix=CHECK-IR -input-file=%t.asan_recover.ll %s
+// RUN: %target-swiftc_driver %s -target %sanitizers-target-triple -g -sanitize=address -sanitize-recover=address -import-objc-header %S/asan_interface.h -o %t_asan_recover
+// RUN: %target-codesign %t_asan_recover
+// RUN: env %env-ASAN_OPTIONS=halt_on_error=0 %target-run %t_asan_recover > %t_asan_recover.stdout 2> %t_asan_recover.stderr
+// RUN: %FileCheck --check-prefixes=CHECK-COMMON-STDERR,CHECK-RECOVER-STDERR -input-file=%t_asan_recover.stderr %s
+// RUN: %FileCheck --check-prefixes=CHECK-COMMON-STDOUT,CHECK-RECOVER-STDOUT -input-file=%t_asan_recover.stdout %s
+
+// Check with recovery instrumentation but without runtime option to continue execution.
+// RUN: not env %env-ASAN_OPTIONS=abort_on_error=0,halt_on_error=1 %target-run %t_asan_recover > %t_asan_no_runtime_recover.stdout 2> %t_asan_no_runtime_recover.stderr
+// RUN: %FileCheck --check-prefixes=CHECK-COMMON-STDERR -input-file=%t_asan_no_runtime_recover.stderr %s
+// RUN: %FileCheck --check-prefixes=CHECK-COMMON-STDOUT,CHECK-NO-RECOVER-STDOUT -input-file=%t_asan_no_runtime_recover.stdout %s
+
+// Check that without recovery instrumentation and runtime option to continue execution that error recovery does not happen.
+// RUN: %target-swiftc_driver %s -target %sanitizers-target-triple -g -sanitize=address -import-objc-header %S/asan_interface.h -o %t_asan_no_recover
+// RUN: %target-codesign %t_asan_no_recover
+// RUN: not env %env-ASAN_OPTIONS=abort_on_error=0,halt_on_error=0 %target-run %t_asan_no_recover > %t_asan_no_recover.stdout 2> %t_asan_no_recover.stderr
+// RUN: %FileCheck --check-prefixes=CHECK-COMMON-STDERR -input-file=%t_asan_no_recover.stderr %s
+// RUN: %FileCheck --check-prefixes=CHECK-COMMON-STDOUT,CHECK-NO-RECOVER-STDOUT -input-file=%t_asan_no_recover.stdout %s
+
+// We need to test reads via instrumentation not via runtime so try to check
+// for calls to unwanted interceptor/runtime functions.
+// CHECK-IR-NOT: call {{.+}} @__asan_memcpy
+// CHECK-IR-NOT: call {{.+}} @memcpy
+
+// FIXME: We need this so we can flush stdout but this won't
+// work on other Platforms (e.g. Windows).
+#if os(Linux)
+    import Glibc
+#else
+    import Darwin.C
+#endif
+
+// CHECK-COMMON-STDOUT: START
+print("START")
+fflush(stdout)
+
+let size:Int = 128;
+
+// In this test we need multiple issues to occur that ASan can detect.
+// Allocating a buffer and artificially poisoning it seems like the best way to
+// test this because there's no undefined behavior happening. Hopefully this
+// means that the program behaviour after ASan catches an issue should be
+// consistent. If we did something different like two use-after-free issues the
+// behaviour could be very unpredicatable resulting in a flakey test.
+var x = UnsafeMutablePointer<UInt8>.allocate(capacity: size)
+x.initialize(repeating: 0, count: size)
+__asan_poison_memory_region(UnsafeMutableRawPointer(x), size)
+
+// Perform accesses that ASan will catch. Note it is important here that
+// the reads are performed **in** the instrumented code so that the
+// instrumentation catches the access to poisoned memory. I tried doing:
+//
+// ```
+// var x = x.advanced(by: 0).pointee
+// print(x)
+// ```
+//
+// However, this generated code that called into memcpy rather than performing
+// a direct read which meant that ASan caught an issue via its interceptors
+// rather than from instrumentation, which does not test the right thing here.
+//
+// Doing:
+//
+// ```
+// print("Read first element:\(x.advanced(by: 0).pointee)")
+// ```
+//
+// seems to do the right thing right now but this seems really fragile.
+
+// First error
+// NOTE: Testing for stackframe `#0` should ensure that the poison read
+// happened in instrumentation and not in an interceptor.
+// CHECK-COMMON-STDERR: AddressSanitizer: use-after-poison
+// CHECK-COMMON-STDERR: #0 0x{{.+}} in main {{.*}}asan_recover.swift:[[@LINE+1]]
+print("Read first element:\(x.advanced(by: 0).pointee)")
+fflush(stdout)
+// CHECK-RECOVER-STDOUT: Read first element:0
+
+// Second error
+// CHECK-RECOVER-STDERR: AddressSanitizer: use-after-poison
+// CHECK-RECOVER-STDERR: #0 0x{{.+}} in main {{.*}}asan_recover.swift:[[@LINE+1]]
+print("Read second element:\(x.advanced(by: 1).pointee)")
+fflush(stdout)
+// CHECK-RECOVER-STDOUT: Read second element:0
+
+__asan_unpoison_memory_region(UnsafeMutableRawPointer(x), size)
+
+x.deallocate();
+// CHECK-NO-RECOVER-STDOUT-NOT: DONE
+// CHECK-RECOVER-STDOUT: DONE
+print("DONE")
+fflush(stdout)