[turbofan] Support poisoning arguments in JavaScript.

This adds support for poisoning the stack pointer and implicit register
arguments like the context register and the function register in the
prologue of generated code with JavaScript linkage. The speculation
poison is computed similarly to the interpreter by matching expected
with actual code start addresses.

R=jarin@chromium.org,rmcilroy@chromium.org
BUG=chromium:798964

Change-Id: I5fa48844745459cf7b3d00c407a7b835f61c857b
Reviewed-on: https://chromium-review.googlesource.com/919167
Commit-Queue: Michael Starzinger <mstarzinger@chromium.org>
Reviewed-by: Jaroslav Sevcik <jarin@chromium.org>
Reviewed-by: Ross McIlroy <rmcilroy@chromium.org>
Cr-Commit-Position: refs/heads/master@{#51553}

src/arm/assembler-arm.cc

@@ -1505,6 +1505,10 @@ void Assembler::and_(Register dst, Register src1, const Operand& src2,
   AddrMode1(cond | AND | s, dst, src1, src2);
 }
 
+void Assembler::and_(Register dst, Register src1, Register src2, SBit s,
+                     Condition cond) {
+  and_(dst, src1, Operand(src2), s, cond);
+}
 
 void Assembler::eor(Register dst, Register src1, const Operand& src2,
                     SBit s, Condition cond) {

src/arm/assembler-arm.h

@@ -741,6 +741,8 @@ class Assembler : public AssemblerBase {
 
   void and_(Register dst, Register src1, const Operand& src2,
             SBit s = LeaveCC, Condition cond = al);
+  void and_(Register dst, Register src1, Register src2, SBit s = LeaveCC,
+            Condition cond = al);
 
   void eor(Register dst, Register src1, const Operand& src2,
            SBit s = LeaveCC, Condition cond = al);

src/compilation-info.cc

@@ -45,6 +45,7 @@ CompilationInfo::CompilationInfo(Zone* zone, Isolate* isolate,
   if (FLAG_function_context_specialization) MarkAsFunctionContextSpecializing();
   if (FLAG_turbo_splitting) MarkAsSplittingEnabled();
   if (!FLAG_turbo_disable_switch_jump_table) SetFlag(kSwitchJumpTableEnabled);
+  if (FLAG_untrusted_code_mitigations) MarkAsPoisoningRegisterArguments();
 
   // Collect source positions for optimized code when profiling or if debugger
   // is active, to be able to get more precise source positions at the price of
@@ -59,7 +60,7 @@ CompilationInfo::CompilationInfo(Vector<const char> debug_name, Zone* zone,
     : CompilationInfo(debug_name, static_cast<AbstractCode::Kind>(code_kind),
                       zone) {
   if (code_kind == Code::BYTECODE_HANDLER && has_untrusted_code_mitigations()) {
-    SetFlag(CompilationInfo::kGenerateSpeculationPoison);
+    SetFlag(CompilationInfo::kGenerateSpeculationPoisonOnEntry);
   }
 }
 

src/compilation-info.h

@@ -53,7 +53,8 @@ class V8_EXPORT_PRIVATE CompilationInfo final {
     kLoopPeelingEnabled = 1 << 11,
     kUntrustedCodeMitigations = 1 << 12,
     kSwitchJumpTableEnabled = 1 << 13,
-    kGenerateSpeculationPoison = 1 << 14,
+    kGenerateSpeculationPoisonOnEntry = 1 << 14,
+    kPoisonRegisterArguments = 1 << 15,
   };
 
   // TODO(mtrofin): investigate if this might be generalized outside wasm, with
@@ -177,8 +178,19 @@ class V8_EXPORT_PRIVATE CompilationInfo final {
     return GetFlag(kSwitchJumpTableEnabled);
   }
 
-  bool is_speculation_poison_enabled() const {
-    bool enabled = GetFlag(kGenerateSpeculationPoison);
+  bool is_generating_speculation_poison_on_entry() const {
+    bool enabled = GetFlag(kGenerateSpeculationPoisonOnEntry);
+    DCHECK_IMPLIES(enabled, has_untrusted_code_mitigations());
+    return enabled;
+  }
+
+  void MarkAsPoisoningRegisterArguments() {
+    DCHECK(has_untrusted_code_mitigations());
+    SetFlag(kGenerateSpeculationPoisonOnEntry);
+    SetFlag(kPoisonRegisterArguments);
+  }
+  bool is_poisoning_register_arguments() const {
+    bool enabled = GetFlag(kPoisonRegisterArguments);
     DCHECK_IMPLIES(enabled, has_untrusted_code_mitigations());
     return enabled;
   }

src/compiler/arm/code-generator-arm.cc

@@ -633,6 +633,12 @@ void CodeGenerator::GenerateSpeculationPoison() {
   __ csdb();
 }
 
+void CodeGenerator::AssembleRegisterArgumentPoisoning() {
+  __ and_(kJSFunctionRegister, kJSFunctionRegister, kSpeculationPoisonRegister);
+  __ and_(kContextRegister, kContextRegister, kSpeculationPoisonRegister);
+  __ and_(sp, sp, kSpeculationPoisonRegister);
+}
+
 // Assembles an instruction after register allocation, producing machine code.
 CodeGenerator::CodeGenResult CodeGenerator::AssembleArchInstruction(
     Instruction* instr) {

src/compiler/arm64/code-generator-arm64.cc

@@ -576,6 +576,17 @@ void CodeGenerator::GenerateSpeculationPoison() {
   __ Csdb();
 }
 
+void CodeGenerator::AssembleRegisterArgumentPoisoning() {
+  UseScratchRegisterScope temps(tasm());
+  Register scratch = temps.AcquireX();
+
+  __ Mov(scratch, sp);
+  __ And(kJSFunctionRegister, kJSFunctionRegister, kSpeculationPoisonRegister);
+  __ And(kContextRegister, kContextRegister, kSpeculationPoisonRegister);
+  __ And(scratch, scratch, kSpeculationPoisonRegister);
+  __ Mov(sp, scratch);
+}
+
 // Assembles an instruction after register allocation, producing machine code.
 CodeGenerator::CodeGenResult CodeGenerator::AssembleArchInstruction(
     Instruction* instr) {

src/compiler/code-generator.cc

@@ -153,24 +153,33 @@ void CodeGenerator::AssembleCode() {
   // Check that {kJavaScriptCallCodeStartRegister} has been set correctly.
   if (FLAG_debug_code & (info->code_kind() == Code::OPTIMIZED_FUNCTION ||
                          info->code_kind() == Code::BYTECODE_HANDLER)) {
+    tasm()->RecordComment("-- Prologue: check code start register --");
     AssembleCodeStartRegisterCheck();
   }
 
-  if (info->is_speculation_poison_enabled()) {
-    GenerateSpeculationPoison();
-  } else {
-    InitializePoisonForLoadsIfNeeded();
-  }
-
-  // TODO(jupvfranco): This should be the first thing in the code after
-  // generating speculation poison, or otherwise MaybeCallEntryHookDelayed may
-  // happen twice (for optimized and deoptimized code). We want to bailout only
-  // from JS functions, which are the only ones that are optimized.
+  // TODO(jupvfranco): This should be the first thing in the code, otherwise
+  // MaybeCallEntryHookDelayed may happen twice (for optimized and deoptimized
+  // code). We want to bailout only from JS functions, which are the only ones
+  // that are optimized.
   if (info->IsOptimizing()) {
     DCHECK(linkage()->GetIncomingDescriptor()->IsJSFunctionCall());
+    tasm()->RecordComment("-- Prologue: check for deoptimization --");
     BailoutIfDeoptimized();
   }
 
+  // Initialize {kSpeculationPoisonRegister} either by comparing the expected
+  // with the actual call target, or by unconditionally using {-1} initially.
+  // Masking register arguments with it only makes sense in the first case.
+  if (info->is_generating_speculation_poison_on_entry()) {
+    tasm()->RecordComment("-- Prologue: generate speculation poison --");
+    GenerateSpeculationPoison();
+    if (info->is_poisoning_register_arguments()) {
+      AssembleRegisterArgumentPoisoning();
+    }
+  } else {
+    InitializePoisonForLoadsIfNeeded();
+  }
+
   // Define deoptimization literals for all inlined functions.
   DCHECK_EQ(0u, deoptimization_literals_.size());
   for (CompilationInfo::InlinedFunctionHolder& inlined :

src/compiler/code-generator.h

@@ -210,6 +210,10 @@ class CodeGenerator final : public GapResolver::Assembler {
   // the code is executing speculatively.
   void GenerateSpeculationPoison();
 
+  // Generates code to poison the stack pointer and implicit register arguments
+  // like the context register and the function register.
+  void AssembleRegisterArgumentPoisoning();
+
   // Generates an architecture-specific, descriptor-specific prologue
   // to set up a stack frame.
   void AssembleConstructFrame();

src/compiler/ia32/code-generator-ia32.cc

@@ -513,8 +513,8 @@ void CodeGenerator::AssembleCodeStartRegisterCheck() {
 //    3. if it is not zero then it jumps to the builtin.
 void CodeGenerator::BailoutIfDeoptimized() {
   int offset = Code::kCodeDataContainerOffset - Code::kHeaderSize;
-  __ mov(ecx, Operand(kJavaScriptCallCodeStartRegister, offset));
-  __ test(FieldOperand(ecx, CodeDataContainer::kKindSpecificFlagsOffset),
+  __ mov(ebx, Operand(kJavaScriptCallCodeStartRegister, offset));
+  __ test(FieldOperand(ebx, CodeDataContainer::kKindSpecificFlagsOffset),
           Immediate(1 << Code::kMarkedForDeoptimizationBit));
   Handle<Code> code = isolate()->builtins()->builtin_handle(
       Builtins::kCompileLazyDeoptimizedCode);
@@ -535,6 +535,12 @@ void CodeGenerator::GenerateSpeculationPoison() {
   __ pop(eax);  // Restore eax.
 }
 
+void CodeGenerator::AssembleRegisterArgumentPoisoning() {
+  __ and_(kJSFunctionRegister, kSpeculationPoisonRegister);
+  __ and_(kContextRegister, kSpeculationPoisonRegister);
+  __ and_(esp, kSpeculationPoisonRegister);
+}
+
 // Assembles an instruction after register allocation, producing machine code.
 CodeGenerator::CodeGenResult CodeGenerator::AssembleArchInstruction(
     Instruction* instr) {

src/compiler/mips/code-generator-mips.cc

@@ -634,12 +634,12 @@ void CodeGenerator::AssembleCodeStartRegisterCheck() {
 //    3. if it is not zero then it jumps to the builtin.
 void CodeGenerator::BailoutIfDeoptimized() {
   int offset = Code::kCodeDataContainerOffset - Code::kHeaderSize;
-  __ lw(a2, MemOperand(kJavaScriptCallCodeStartRegister, offset));
-  __ lw(a2, FieldMemOperand(a2, CodeDataContainer::kKindSpecificFlagsOffset));
-  __ And(a2, a2, Operand(1 << Code::kMarkedForDeoptimizationBit));
+  __ lw(at, MemOperand(kJavaScriptCallCodeStartRegister, offset));
+  __ lw(at, FieldMemOperand(at, CodeDataContainer::kKindSpecificFlagsOffset));
+  __ And(at, at, Operand(1 << Code::kMarkedForDeoptimizationBit));
   Handle<Code> code = isolate()->builtins()->builtin_handle(
       Builtins::kCompileLazyDeoptimizedCode);
-  __ Jump(code, RelocInfo::CODE_TARGET, ne, a2, Operand(zero_reg));
+  __ Jump(code, RelocInfo::CODE_TARGET, ne, at, Operand(zero_reg));
 }
 
 void CodeGenerator::GenerateSpeculationPoison() {
@@ -661,6 +661,12 @@ void CodeGenerator::GenerateSpeculationPoison() {
          kSpeculationPoisonRegister);
 }
 
+void CodeGenerator::AssembleRegisterArgumentPoisoning() {
+  __ And(kJSFunctionRegister, kJSFunctionRegister, kSpeculationPoisonRegister);
+  __ And(kContextRegister, kContextRegister, kSpeculationPoisonRegister);
+  __ And(sp, sp, kSpeculationPoisonRegister);
+}
+
 // Assembles an instruction after register allocation, producing machine code.
 CodeGenerator::CodeGenResult CodeGenerator::AssembleArchInstruction(
     Instruction* instr) {

src/compiler/mips64/code-generator-mips64.cc

@@ -650,12 +650,12 @@ void CodeGenerator::AssembleCodeStartRegisterCheck() {
 //    3. if it is not zero then it jumps to the builtin.
 void CodeGenerator::BailoutIfDeoptimized() {
   int offset = Code::kCodeDataContainerOffset - Code::kHeaderSize;
-  __ Ld(a2, MemOperand(kJavaScriptCallCodeStartRegister, offset));
-  __ Lw(a2, FieldMemOperand(a2, CodeDataContainer::kKindSpecificFlagsOffset));
-  __ And(a2, a2, Operand(1 << Code::kMarkedForDeoptimizationBit));
+  __ Ld(at, MemOperand(kJavaScriptCallCodeStartRegister, offset));
+  __ Lw(at, FieldMemOperand(at, CodeDataContainer::kKindSpecificFlagsOffset));
+  __ And(at, at, Operand(1 << Code::kMarkedForDeoptimizationBit));
   Handle<Code> code = isolate()->builtins()->builtin_handle(
       Builtins::kCompileLazyDeoptimizedCode);
-  __ Jump(code, RelocInfo::CODE_TARGET, ne, a2, Operand(zero_reg));
+  __ Jump(code, RelocInfo::CODE_TARGET, ne, at, Operand(zero_reg));
 }
 
 void CodeGenerator::GenerateSpeculationPoison() {
@@ -677,6 +677,12 @@ void CodeGenerator::GenerateSpeculationPoison() {
          kSpeculationPoisonRegister);
 }
 
+void CodeGenerator::AssembleRegisterArgumentPoisoning() {
+  __ And(kJSFunctionRegister, kJSFunctionRegister, kSpeculationPoisonRegister);
+  __ And(kContextRegister, kContextRegister, kSpeculationPoisonRegister);
+  __ And(sp, sp, kSpeculationPoisonRegister);
+}
+
 // Assembles an instruction after register allocation, producing machine code.
 CodeGenerator::CodeGenResult CodeGenerator::AssembleArchInstruction(
     Instruction* instr) {

src/compiler/pipeline.cc

@@ -1540,7 +1540,7 @@ struct InstructionSelectionPhase {
         data->info()->switch_jump_table_enabled()
             ? InstructionSelector::kEnableSwitchJumpTable
             : InstructionSelector::kDisableSwitchJumpTable,
-        data->info()->is_speculation_poison_enabled()
+        data->info()->is_generating_speculation_poison_on_entry()
             ? InstructionSelector::kEnableSpeculationPoison
             : InstructionSelector::kDisableSpeculationPoison,
         data->info()->is_source_positions_enabled()

src/compiler/x64/code-generator-x64.cc

@@ -616,6 +616,12 @@ void CodeGenerator::GenerateSpeculationPoison() {
   __ cmovq(equal, kSpeculationPoisonRegister, rbx);
 }
 
+void CodeGenerator::AssembleRegisterArgumentPoisoning() {
+  __ andq(kJSFunctionRegister, kSpeculationPoisonRegister);
+  __ andq(kContextRegister, kSpeculationPoisonRegister);
+  __ andq(rsp, kSpeculationPoisonRegister);
+}
+
 // Assembles an instruction after register allocation, producing machine code.
 CodeGenerator::CodeGenResult CodeGenerator::AssembleArchInstruction(
     Instruction* instr) {