IRGen: Add option to enable stack protector

Adds frontend option -enable-stack-protector to enable emission of a
stack protector.

Disabled by default.

When enabled enables LLVM's strong stack protection mode.

rdar://93677524

Author: aschwaighofer

include/swift/AST/IRGenOptions.h

@@ -94,6 +94,8 @@ enum class ReflectionMetadataMode : unsigned {
   Runtime,      ///< Make reflection metadata fully available.
 };
 
+enum class StackProtectorMode : bool { NoStackProtector, StackProtector };
+
 using clang::PointerAuthSchema;
 
 struct PointerAuthOptions : clang::PointerAuthOptions {
@@ -280,6 +282,8 @@ class IRGenOptions {
   /// Whether we should run swift specific LLVM optimizations after IRGen.
   unsigned DisableSwiftSpecificLLVMOptzns : 1;
 
+  unsigned EnableStackProtector : 1;
+
   /// Special codegen for playgrounds.
   unsigned Playground : 1;
 
@@ -457,10 +461,10 @@ class IRGenOptions {
         DebugInfoFormat(IRGenDebugInfoFormat::None),
         DisableClangModuleSkeletonCUs(false), UseJIT(false),
         DisableLLVMOptzns(false), DisableSwiftSpecificLLVMOptzns(false),
-        Playground(false), EmitStackPromotionChecks(false),
-        UseSingleModuleLLVMEmission(false), FunctionSections(false),
-        PrintInlineTree(false), EmbedMode(IRGenEmbedMode::None),
-        LLVMLTOKind(IRGenLLVMLTOKind::None),
+        EnableStackProtector(false), Playground(false),
+        EmitStackPromotionChecks(false), UseSingleModuleLLVMEmission(false),
+        FunctionSections(false), PrintInlineTree(false),
+        EmbedMode(IRGenEmbedMode::None), LLVMLTOKind(IRGenLLVMLTOKind::None),
         SwiftAsyncFramePointer(SwiftAsyncFramePointerKind::Auto),
         HasValueNamesSetting(false), ValueNames(false),
         ReflectionMetadata(ReflectionMetadataMode::Runtime),
@@ -543,6 +547,9 @@ class IRGenOptions {
   bool hasMultipleIRGenThreads() const { return !UseSingleModuleLLVMEmission && NumThreads > 1; }
   bool shouldPerformIRGenerationInParallel() const { return !UseSingleModuleLLVMEmission && NumThreads != 0; }
   bool hasMultipleIGMs() const { return hasMultipleIRGenThreads(); }
+  StackProtectorMode getStackProtectorMode() const {
+    return StackProtectorMode(EnableStackProtector);
+  }
 };
 
 } // end namespace swift

include/swift/Option/FrontendOptions.td

@@ -1079,6 +1079,13 @@ def concurrency_model_EQ :
     Joined<["-"], "concurrency-model=">,
     Alias<concurrency_model>;
 
+def enable_stack_protector :
+  Flag<["-"], "enable-stack-protector">,
+  HelpText<"Enable the stack protector">;
+def disable_stack_protector :
+  Flag<["-"], "disable-stack-protector">,
+  HelpText<"Disable the stack-protector">;
+
 def enable_new_llvm_pass_manager :
   Flag<["-"], "enable-new-llvm-pass-manager">,
   HelpText<"Enable the new llvm pass manager">;

lib/Frontend/CompilerInvocation.cpp

@@ -2407,6 +2407,10 @@ static bool ParseIRGenArgs(IRGenOptions &Opts, ArgList &Args,
                    OPT_enable_new_llvm_pass_manager,
                    Opts.LegacyPassManager);
 
+  Opts.EnableStackProtector =
+      Args.hasFlag(OPT_enable_stack_protector, OPT_disable_stack_protector,
+                   Opts.EnableStackProtector);
+
   return false;
 }
 

lib/IRGen/GenDecl.cpp

@@ -2246,11 +2246,11 @@ static bool isPointerTo(llvm::Type *ptrTy, llvm::Type *objTy) {
 }
 
 /// Get or create an LLVM function with these linkage rules.
-llvm::Function *irgen::createFunction(IRGenModule &IGM,
-                                      LinkInfo &linkInfo,
+llvm::Function *irgen::createFunction(IRGenModule &IGM, LinkInfo &linkInfo,
                                       const Signature &signature,
                                       llvm::Function *insertBefore,
-                                      OptimizationMode FuncOptMode) {
+                                      OptimizationMode FuncOptMode,
+                                      StackProtectorMode stackProtect) {
   auto name = linkInfo.getName();
 
   llvm::Function *existing = IGM.Module.getFunction(name);
@@ -2282,7 +2282,7 @@ llvm::Function *irgen::createFunction(IRGenModule &IGM,
       .to(fn, linkInfo.isForDefinition());
 
   llvm::AttrBuilder initialAttrs;
-  IGM.constructInitialFnAttributes(initialAttrs, FuncOptMode);
+  IGM.constructInitialFnAttributes(initialAttrs, FuncOptMode, stackProtect);
   // Merge initialAttrs with attrs.
   auto updatedAttrs = signature.getAttributes().addFnAttributes(
       IGM.getLLVMContext(), initialAttrs);
@@ -3059,7 +3059,7 @@ void IRGenModule::emitDynamicReplacementOriginalFunctionThunk(SILFunction *f) {
   LinkInfo implLink = LinkInfo::get(*this, entity, ForDefinition);
   auto implFn =
       createFunction(*this, implLink, signature, nullptr /*insertBefore*/,
-                     f->getOptimizationMode());
+                     f->getOptimizationMode(), shouldEmitStackProtector(f));
   implFn->addFnAttr(llvm::Attribute::NoInline);
 
   IRGenFunction IGF(*this, implFn);
@@ -3168,6 +3168,10 @@ llvm::Constant *swift::irgen::emitCXXConstructorThunkIfNeeded(
   return thunk;
 }
 
+StackProtectorMode IRGenModule::shouldEmitStackProtector(SILFunction *f) {
+  return IRGen.Opts.getStackProtectorMode();
+}
+
 /// Find the entry point for a SIL function.
 llvm::Function *IRGenModule::getAddrOfSILFunction(
     SILFunction *f, ForDefinition_t forDefinition,
@@ -3263,7 +3267,7 @@ llvm::Function *IRGenModule::getAddrOfSILFunction(
   LinkInfo link = LinkInfo::get(*this, entity, forDefinition);
 
   fn = createFunction(*this, link, signature, insertBefore,
-                      f->getOptimizationMode());
+                      f->getOptimizationMode(), shouldEmitStackProtector(f));
 
   // If `hasCReferences` is true, then the function is either marked with
   // @_silgen_name OR @_cdecl.  If it is the latter, it must have a definition

lib/IRGen/GenDecl.h

@@ -41,12 +41,11 @@ namespace irgen {
                                   llvm::GlobalValue *global,
                                   const LinkEntity &entity);
 
-  llvm::Function *createFunction(IRGenModule &IGM,
-                                 LinkInfo &linkInfo,
-                                 const Signature &signature,
-                                 llvm::Function *insertBefore = nullptr,
-                                 OptimizationMode FuncOptMode =
-                                   OptimizationMode::NotSet);
+  llvm::Function *createFunction(
+      IRGenModule &IGM, LinkInfo &linkInfo, const Signature &signature,
+      llvm::Function *insertBefore = nullptr,
+      OptimizationMode FuncOptMode = OptimizationMode::NotSet,
+      StackProtectorMode stackProtect = StackProtectorMode::NoStackProtector);
 
   llvm::GlobalVariable *
   createVariable(IRGenModule &IGM, LinkInfo &linkInfo, llvm::Type *objectType,

lib/IRGen/IRGenModule.cpp

@@ -1246,8 +1246,9 @@ void IRGenModule::setHasNoFramePointer(llvm::Function *F) {
 }
 
 /// Construct initial function attributes from options.
-void IRGenModule::constructInitialFnAttributes(llvm::AttrBuilder &Attrs,
-                                               OptimizationMode FuncOptMode) {
+void IRGenModule::constructInitialFnAttributes(
+    llvm::AttrBuilder &Attrs, OptimizationMode FuncOptMode,
+    StackProtectorMode stackProtector) {
   // Add the default attributes for the Clang configuration.
   clang::CodeGen::addDefaultFunctionDefinitionAttributes(getClangCGM(), Attrs);
 
@@ -1261,6 +1262,10 @@ void IRGenModule::constructInitialFnAttributes(llvm::AttrBuilder &Attrs,
     Attrs.removeAttribute(llvm::Attribute::MinSize);
     Attrs.removeAttribute(llvm::Attribute::OptimizeForSize);
   }
+  if (stackProtector == StackProtectorMode::StackProtector) {
+    Attrs.addAttribute(llvm::Attribute::StackProtectStrong);
+    Attrs.addAttribute("stack-protector-buffer-size", llvm::utostr(8));
+  }
 }
 
 llvm::AttributeList IRGenModule::constructInitialAttributes() {

lib/IRGen/IRGenModule.h

@@ -22,6 +22,7 @@
 #include "SwiftTargetInfo.h"
 #include "TypeLayout.h"
 #include "swift/AST/Decl.h"
+#include "swift/AST/IRGenOptions.h"
 #include "swift/AST/LinkLibrary.h"
 #include "swift/AST/Module.h"
 #include "swift/AST/ReferenceCounting.h"
@@ -32,8 +33,8 @@
 #include "swift/Basic/OptimizationMode.h"
 #include "swift/Basic/SuccessorMap.h"
 #include "swift/IRGen/ValueWitness.h"
-#include "swift/SIL/SILFunction.h"
 #include "swift/SIL/RuntimeEffect.h"
+#include "swift/SIL/SILFunction.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/Hashing.h"
@@ -1441,12 +1442,14 @@ private:                            \
   /// invalid.
   bool finalize();
 
-  void constructInitialFnAttributes(llvm::AttrBuilder &Attrs,
-                                    OptimizationMode FuncOptMode =
-                                      OptimizationMode::NotSet);
+  void constructInitialFnAttributes(
+      llvm::AttrBuilder &Attrs,
+      OptimizationMode FuncOptMode = OptimizationMode::NotSet,
+      StackProtectorMode stackProtect = StackProtectorMode::NoStackProtector);
   void setHasNoFramePointer(llvm::AttrBuilder &Attrs);
   void setHasNoFramePointer(llvm::Function *F);
   llvm::AttributeList constructInitialAttributes();
+  StackProtectorMode shouldEmitStackProtector(SILFunction *f);
 
   void emitProtocolDecl(ProtocolDecl *D);
   void emitEnumDecl(EnumDecl *D);

test/IRGen/stack_protector.swift

@@ -0,0 +1,18 @@
+// RUN: %target-swift-frontend -enable-stack-protector -emit-ir %s -o - | %FileCheck %s
+
+@_silgen_name("escape")
+func f<T>(_ arg: UnsafePointer<T>)
+
+public func escapeGenericValue<T>(_ t: T) {
+    withUnsafePointer(to: t) { ptr in
+       f(ptr)
+    }
+}
+
+// CHECK: define {{.*}}swiftcc void @"$s15stack_protector21requestStackProtectoryyF"() [[SSPATTRS:#[0-9]+]] {
+public func requestStackProtector() {
+    var x: Int = 0
+    escapeGenericValue(x)
+}
+
+// CHECK: [[SSPATTRS]] = { sspstrong {{.*}}"stack-protector-buffer-size"="8"