Merge pull request #62057 from eeckstein/enable-stack-protection

SILOptimizer: enable stack protection by default

Author: eeckstein

SwiftCompilerSources/Sources/Optimizer/ModulePasses/StackProtection.swift

@@ -28,24 +28,24 @@ private func log(_ message: @autoclosure () -> String) {
 /// Within safe swift code there shouldn't be any buffer overflows. But if the address
 /// of a stack variable is converted to an unsafe pointer, it's not in the control of
 /// the compiler anymore.
-/// This means, if there is any `address_to_pointer` instruction for an `alloc_stack`,
-/// such a function is marked for stack protection.
-/// Another case is `index_addr` for non-tail allocated memory. This pattern appears if
-/// pointer arithmetic is done with unsafe pointers in swift code.
+/// This means, if an `alloc_stack` ends up at an `address_to_pointer [stack_protection]`,
+/// the `alloc_stack`'s function is marked for stack protection.
+/// Another case is `index_addr [stack_protection]` for non-tail allocated memory. This
+/// pattern appears if pointer arithmetic is done with unsafe pointers in swift code.
 ///
 /// If the origin of an unsafe pointer can only be tracked to a function argument, the
 /// pass tries to find the root stack allocation for such an argument by doing an
-/// inter-procedural analysis. If this is not possible, the fallback is to move the
-/// argument into a temporary `alloc_stack` and do the unsafe pointer operations on
-/// the temporary.
+/// inter-procedural analysis. If this is not possible and the `enableMoveInoutStackProtection`
+/// option is set, the fallback is to move the argument into a temporary `alloc_stack`
+/// and do the unsafe pointer operations on the temporary.
 let stackProtection = ModulePass(name: "stack-protection", {
     (context: ModulePassContext) in
 
   if !context.options.enableStackProtection {
     return
   }
 
-  var optimization = StackProtectionOptimization()
+  var optimization = StackProtectionOptimization(enableMoveInout: context.options.enableMoveInoutStackProtection)
   optimization.processModule(context)
 })
 
@@ -60,13 +60,15 @@ let functionStackProtection = FunctionPass(name: "function-stack-protection", {
     return
   }
 
-  var optimization = StackProtectionOptimization()
+  var optimization = StackProtectionOptimization(enableMoveInout: context.options.enableMoveInoutStackProtection)
   optimization.process(function: function, context)
 })
 
 /// The optimization algorithm.
 private struct StackProtectionOptimization {
 
+  private let enableMoveInout: Bool
+
   // The following members are nil/not used if this utility is used on function-level.
 
   private var moduleContext: ModulePassContext?
@@ -76,7 +78,11 @@ private struct StackProtectionOptimization {
   // Functions (other than the currently processed one) which need stack protection,
   // are added to this array in `findOriginsInCallers`.
   private var needStackProtection: [Function] = []
-  
+
+  init(enableMoveInout: Bool) {
+    self.enableMoveInout = enableMoveInout
+  }
+
   /// The main entry point if running on module-level.
   mutating func processModule(_ moduleContext: ModulePassContext) {
     self.moduleContext = moduleContext
@@ -147,6 +153,8 @@ private struct StackProtectionOptimization {
     case .yes:
       // For example:
       //     %baseAddr = alloc_stack $T
+      log("local: \(function.name) -- \(instruction)")
+
       function.setNeedsStackProtection(context)
 
     case .decidedInCaller(let arg):
@@ -157,7 +165,7 @@ private struct StackProtectionOptimization {
       defer { worklist.deinitialize() }
       worklist.push(arg)
 
-      if !findOriginsInCallers(&worklist) {
+      if findOriginsInCallers(&worklist) == NeedInsertMoves.yes {
         // We don't know the origin of the function argument. Therefore we need to do the
         // conservative default which is to move the value to a temporary stack location.
         if let beginAccess = scope {
@@ -179,22 +187,19 @@ private struct StackProtectionOptimization {
 
       // If the object is passed as an argument to its function, add those arguments
       // to the worklist.
-      switch worklist.push(rootsOf: obj) {
-      case .failed:
-        // If we cannot find the roots, the object is most likely not stack allocated.
-        return
-      case .succeeded(let foundStackAlloc):
-        if foundStackAlloc {
-          // The object is created by an `alloc_ref [stack]`.
-          function.setNeedsStackProtection(context)
-        }
+      let (_, foundStackAlloc) = worklist.push(rootsOf: obj)
+      if foundStackAlloc {
+        // The object is created by an `alloc_ref [stack]`.
+        log("objectIfStackPromoted: \(function.name) -- \(instruction)")
+
+        function.setNeedsStackProtection(context)
       }
       // In case the (potentially) stack allocated object is passed via an argument,
       // process the worklist as we do for indirect arguments (see above).
       // For example:
       //   bb0(%0: $Class):
       //     %baseAddr = ref_element_addr %0 : $Class, #Class.field
-      if !findOriginsInCallers(&worklist),
+      if findOriginsInCallers(&worklist) == NeedInsertMoves.yes,
          let beginAccess = scope {
         // We don't know the origin of the object. Therefore we need to do the
         // conservative default which is to move the value to a temporary stack location.
@@ -206,15 +211,26 @@ private struct StackProtectionOptimization {
       break
     }
   }
-  
+
+  /// Return value of `findOriginsInCallers()`.
+  enum NeedInsertMoves {
+    // Not all call sites could be identified, and if moves are enabled (`enableMoveInout`)
+    // the original argument should be moved to a temporary.
+    case yes
+
+    // Either all call sites could be identified, which means that stack protection is done
+    // in the callers, or moves are not enabled (`enableMoveInout` is false).
+    case no
+  }
+
   /// Find all origins of function arguments in `worklist`.
   /// All functions, which allocate such an origin are added to `self.needStackProtection`.
   /// Returns true if all origins could be found and false, if there are unknown origins.
-  private mutating func findOriginsInCallers(_ worklist: inout ArgumentWorklist) -> Bool {
+  private mutating func findOriginsInCallers(_ worklist: inout ArgumentWorklist) -> NeedInsertMoves {
 
     guard let moduleContext = moduleContext else {
       // Don't do any inter-procedural analysis when used on function-level.
-      return false
+      return enableMoveInout ? .yes : .no
     }
 
     // Put the resulting functions into a temporary array, because we only add them to
@@ -230,28 +246,33 @@ private struct StackProtectionOptimization {
     while let arg = worklist.pop() {
       let f = arg.function
       let uses = functionUses.getUses(of: f)
-      if uses.hasUnknownUses {
-        return false
+      if uses.hasUnknownUses && enableMoveInout {
+        return NeedInsertMoves.yes
       }
 
       for useInst in uses {
         guard let fri = useInst as? FunctionRefInst else {
-          return false
+          if enableMoveInout {
+            return NeedInsertMoves.yes
+          }
+          continue
         }
 
         for functionRefUse in fri.uses {
-          guard let apply = functionRefUse.instruction as? ApplySite else {
-            return false
-          }
-          guard let callerArgIdx = apply.callerArgIndex(calleeArgIndex: arg.index) else {
-            return false
+          guard let apply = functionRefUse.instruction as? ApplySite,
+                let callerArgIdx = apply.callerArgIndex(calleeArgIndex: arg.index) else {
+            if enableMoveInout {
+              return NeedInsertMoves.yes
+            }
+            continue
           }
           let callerArg = apply.arguments[callerArgIdx]
           if callerArg.type.isAddress {
             // It's an indirect argument.
             switch callerArg.accessBase.isStackAllocated {
             case .yes:
               if !callerArg.function.needsStackProtection {
+                log("alloc_stack in caller: \(callerArg.function.name) -- \(callerArg)")
                 newFunctions.push(callerArg.function)
               }
             case .no:
@@ -266,36 +287,38 @@ private struct StackProtectionOptimization {
             case .objectIfStackPromoted(let obj):
               // If the object is passed as an argument to its function,
               // add those arguments to the worklist.
-              switch worklist.push(rootsOf: obj) {
-              case .failed:
-                return false
-              case .succeeded(let foundStackAlloc):
-                if foundStackAlloc && !obj.function.needsStackProtection {
-                  // The object is created by an `alloc_ref [stack]`.
-                  newFunctions.push(obj.function)
-                }
+              let (foundUnknownRoots, foundStackAlloc) = worklist.push(rootsOf: obj)
+              if foundUnknownRoots && enableMoveInout {
+                return NeedInsertMoves.yes
+              }
+              if foundStackAlloc && !obj.function.needsStackProtection {
+                // The object is created by an `alloc_ref [stack]`.
+                log("object in caller: \(obj.function.name) -- \(obj)")
+                newFunctions.push(obj.function)
               }
             case .unknown:
-              return false
+              if enableMoveInout {
+                return NeedInsertMoves.yes
+              }
             }
           } else {
             // The argument is an object. If the object is itself passed as an argument
             // to its function, add those arguments to the worklist.
-            switch worklist.push(rootsOf: callerArg) {
-            case .failed:
-              return false
-            case .succeeded(let foundStackAlloc):
-              if foundStackAlloc && !callerArg.function.needsStackProtection {
-                // The object is created by an `alloc_ref [stack]`.
-                newFunctions.push(callerArg.function)
-              }
+            let (foundUnknownRoots, foundStackAlloc) = worklist.push(rootsOf: callerArg)
+            if foundUnknownRoots && enableMoveInout {
+              return NeedInsertMoves.yes
+            }
+            if foundStackAlloc && !callerArg.function.needsStackProtection {
+              // The object is created by an `alloc_ref [stack]`.
+              log("object arg in caller: \(callerArg.function.name) -- \(callerArg)")
+              newFunctions.push(callerArg.function)
             }
           }
         }
       }
     }
     needStackProtection.append(contentsOf: newFunctions)
-    return true
+    return NeedInsertMoves.no
   }
 
   /// Moves the value of an indirect argument to a temporary stack location, if possible.
@@ -366,9 +389,16 @@ private struct StackProtectionOptimization {
 /// Worklist for inter-procedural analysis of function arguments.
 private struct ArgumentWorklist : ValueUseDefWalker {
   var walkUpCache = WalkerCache<SmallProjectionPath>()
+
+  // Used in `push(rootsOf:)`
   private var foundStackAlloc = false
+  private var foundUnknownRoots = false
 
+  // Contains arguments which are already handled and don't need to be put into the worklist again.
+  // Note that this cannot be a `ValueSet`, because argument can be from different functions.
   private var handled = Set<FunctionArgument>()
+
+  // The actual worklist.
   private var list: Stack<FunctionArgument>
 
   init(_ context: PassContext) {
@@ -385,21 +415,15 @@ private struct ArgumentWorklist : ValueUseDefWalker {
     }
   }
 
-  enum PushResult {
-    case failed
-    case succeeded(foundStackAlloc: Bool)
-  }
-
   /// Pushes all roots of `object`, which are function arguments, to the worklist.
-  /// Returns `.succeeded(true)` if some of the roots are `alloc_ref [stack]` instructions.
-  mutating func push(rootsOf object: Value) -> PushResult {
+  /// If the returned `foundUnknownRoots` is true, it means that not all roots of `object` could
+  /// be tracked to a function argument.
+  /// If the returned `foundStackAlloc` than at least one found root is an `alloc_ref [stack]`.
+  mutating func push(rootsOf object: Value) -> (foundUnknownRoots: Bool, foundStackAlloc: Bool) {
     foundStackAlloc = false
-    switch walkUp(value: object, path: SmallProjectionPath(.anything)) {
-      case .continueWalk:
-        return .succeeded(foundStackAlloc: foundStackAlloc)
-      case .abortWalk:
-        return .failed
-    }
+    foundUnknownRoots = false
+    _ =  walkUp(value: object, path: SmallProjectionPath(.anything))
+    return (foundUnknownRoots, foundStackAlloc)
   }
 
   mutating func pop() -> FunctionArgument? {
@@ -413,15 +437,12 @@ private struct ArgumentWorklist : ValueUseDefWalker {
         if ar.canAllocOnStack {
           foundStackAlloc = true
         }
-        return .continueWalk
       case let arg as FunctionArgument:
-        if handled.insert(arg).0 {
-          list.push(arg)
-        }
-        return .continueWalk
-      default:
-        return .abortWalk
+        push(arg)
+    default:
+        foundUnknownRoots = true
     }
+    return .continueWalk
   }
 }
 
@@ -491,7 +512,6 @@ private extension Instruction {
 private extension Function {
   func setNeedsStackProtection(_ context: PassContext) {
     if !needsStackProtection {
-      log("needs protection: \(name)")
       set(needStackProtection: true, context)
     }
   }

SwiftCompilerSources/Sources/Optimizer/PassManager/Options.swift

@@ -18,4 +18,8 @@ struct Options {
   var enableStackProtection: Bool {
     SILOptions_enableStackProtection(_bridged) != 0
   }
+
+  var enableMoveInoutStackProtection: Bool {
+    SILOptions_enableMoveInoutStackProtection(_bridged) != 0
+  }
 }

include/swift/AST/SILOptions.h

@@ -129,7 +129,11 @@ class SILOptions {
   bool EnablePerformanceAnnotations = false;
 
   /// Enables the emission of stack protectors in functions.
-  bool EnableStackProtection = false;
+  bool EnableStackProtection = true;
+
+  /// Like `EnableStackProtection` and also enables moving of values to
+  /// temporaries for stack protection.
+  bool EnableMoveInoutStackProtection = false;
 
   /// Controls whether or not paranoid verification checks are run.
   bool VerifyAll = false;

include/swift/Option/FrontendOptions.td

@@ -1104,6 +1104,9 @@ def enable_stack_protector :
 def disable_stack_protector :
   Flag<["-"], "disable-stack-protector">,
   HelpText<"Disable the stack-protector">;
+def enable_move_inout_stack_protector :
+  Flag<["-"], "enable-move-inout-stack-protector">,
+  HelpText<"Enable the stack protector by moving values to temporaries">;
 
 def enable_new_llvm_pass_manager :
   Flag<["-"], "enable-new-llvm-pass-manager">,

include/swift/SILOptimizer/OptimizerBridging.h

@@ -184,6 +184,7 @@ OptionalBridgedFunction
 PassContext_loadFunction(BridgedPassContext context, llvm::StringRef name);
 
 SwiftInt SILOptions_enableStackProtection(BridgedPassContext context);
+SwiftInt SILOptions_enableMoveInoutStackProtection(BridgedPassContext context);
 
 SWIFT_END_NULLABILITY_ANNOTATIONS
 

lib/Frontend/CompilerInvocation.cpp

@@ -1825,6 +1825,9 @@ static bool ParseSILArgs(SILOptions &Opts, ArgList &Args,
   Opts.EnableStackProtection =
       Args.hasFlag(OPT_enable_stack_protector, OPT_disable_stack_protector,
                    Opts.EnableStackProtection);
+  Opts.EnableMoveInoutStackProtection =
+      Args.hasFlag(OPT_enable_move_inout_stack_protector, OPT_disable_stack_protector,
+                   Opts.EnableMoveInoutStackProtection);
   Opts.VerifyAll |= Args.hasArg(OPT_sil_verify_all);
   Opts.VerifyNone |= Args.hasArg(OPT_sil_verify_none);
   Opts.DebugSerialization |= Args.hasArg(OPT_sil_debug_serialization);

lib/SILOptimizer/PassManager/PassManager.cpp

@@ -1638,3 +1638,8 @@ SwiftInt SILOptions_enableStackProtection(BridgedPassContext context) {
   SILModule *mod = castToPassInvocation(context)->getPassManager()->getModule();
   return mod->getOptions().EnableStackProtection;
 }
+
+SwiftInt SILOptions_enableMoveInoutStackProtection(BridgedPassContext context) {
+  SILModule *mod = castToPassInvocation(context)->getPassManager()->getModule();
+  return mod->getOptions().EnableMoveInoutStackProtection;
+}