Skip to content

SILOptimizer: enable stack protection by default #62057

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Nov 12, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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)
})

Expand All @@ -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?
Expand All @@ -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
Expand Down Expand Up @@ -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):
Expand All @@ -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 {
Expand All @@ -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.
Expand All @@ -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
Expand All @@ -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:
Expand All @@ -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.
Expand Down Expand Up @@ -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) {
Expand All @@ -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? {
Expand All @@ -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
}
}

Expand Down Expand Up @@ -491,7 +512,6 @@ private extension Instruction {
private extension Function {
func setNeedsStackProtection(_ context: PassContext) {
if !needsStackProtection {
log("needs protection: \(name)")
set(needStackProtection: true, context)
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,4 +18,8 @@ struct Options {
var enableStackProtection: Bool {
SILOptions_enableStackProtection(_bridged) != 0
}

var enableMoveInoutStackProtection: Bool {
SILOptions_enableMoveInoutStackProtection(_bridged) != 0
}
}
6 changes: 5 additions & 1 deletion include/swift/AST/SILOptions.h
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
3 changes: 3 additions & 0 deletions include/swift/Option/FrontendOptions.td
Original file line number Diff line number Diff line change
Expand Up @@ -1112,6 +1112,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">,
Expand Down
1 change: 1 addition & 0 deletions include/swift/SILOptimizer/OptimizerBridging.h
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
3 changes: 3 additions & 0 deletions lib/Frontend/CompilerInvocation.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -1831,6 +1831,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);
Expand Down
5 changes: 5 additions & 0 deletions lib/SILOptimizer/PassManager/PassManager.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Loading