Bug 1816000 - Keep FP valid at Wasm entry points. r=jandem

Currenty, FP can be invalid (FailFP) at the entries. For future modifications,
the FP has to be kept intact during throw handling/unwinding.

The patch replaces FailFP-hack with FailInstanceReg one. The InstanceReg pinned
register is always points to valid instance, including at the end of the call.
Using that property to signal when trap has occurred.

Differential Revision: https://phabricator.services.mozilla.com/D170105

js/src/jit/arm/MacroAssembler-arm.cpp

@@ -3478,6 +3478,7 @@ void MacroAssemblerARMCompat::handleFailureWithHandlerTail(
            scratch);
     ma_ldr(Address(sp, ResumeFromException::offsetOfStackPointer()), sp,
            scratch);
+    ma_mov(Imm32(int32_t(wasm::FailInstanceReg)), InstanceReg);
   }
   as_dtr(IsLoad, 32, PostIndex, pc, DTRAddr(sp, DtrOffImm(4)));
 

js/src/jit/arm64/MacroAssembler-arm64.cpp

@@ -365,6 +365,7 @@ void MacroAssemblerCompat::handleFailureWithHandlerTail(Label* profilerExitTail,
       MemOperand(PseudoStackPointer64,
                  ResumeFromException::offsetOfStackPointer()));
   syncStackPtr();
+  Mov(x23, int64_t(wasm::FailInstanceReg));
   ret();
 
   // Found a wasm catch handler, restore state and jump to it.

js/src/jit/loong64/MacroAssembler-loong64.cpp

@@ -5336,6 +5336,7 @@ void MacroAssemblerLOONG64Compat::handleFailureWithHandlerTail(
           FramePointer);
   loadPtr(Address(StackPointer, ResumeFromException::offsetOfStackPointer()),
           StackPointer);
+  ma_li(InstanceReg, ImmWord(wasm::FailInstanceReg));
   ret();
 
   // Found a wasm catch handler, restore state and jump to it.

js/src/jit/mips32/MacroAssembler-mips32.cpp

@@ -1893,6 +1893,7 @@ void MacroAssemblerMIPSCompat::handleFailureWithHandlerTail(
           FramePointer);
   loadPtr(Address(StackPointer, ResumeFromException::offsetOfStackPointer()),
           StackPointer);
+  ma_li(InstanceReg, ImmWord(wasm::FailInstanceReg));
   ret();
 
   // Found a wasm catch handler, restore state and jump to it.

js/src/jit/mips64/MacroAssembler-mips64.cpp

@@ -1918,6 +1918,7 @@ void MacroAssemblerMIPS64Compat::handleFailureWithHandlerTail(
           FramePointer);
   loadPtr(Address(StackPointer, ResumeFromException::offsetOfStackPointer()),
           StackPointer);
+  ma_li(InstanceReg, ImmWord(wasm::FailInstanceReg));
   ret();
 
   // Found a wasm catch handler, restore state and jump to it.

js/src/jit/riscv64/MacroAssembler-riscv64.cpp

@@ -2012,6 +2012,7 @@ void MacroAssemblerRiscv64Compat::handleFailureWithHandlerTail(
           FramePointer);
   loadPtr(Address(StackPointer, ResumeFromException::offsetOfStackPointer()),
           StackPointer);
+  ma_li(InstanceReg, ImmWord(wasm::FailInstanceReg));
   ret();
 
   // Found a wasm catch handler, restore state and jump to it.

js/src/jit/x64/MacroAssembler-x64.cpp

@@ -627,6 +627,7 @@ void MacroAssemblerX64::handleFailureWithHandlerTail(Label* profilerExitTail,
   bind(&wasm);
   loadPtr(Address(rsp, ResumeFromException::offsetOfFramePointer()), rbp);
   loadPtr(Address(rsp, ResumeFromException::offsetOfStackPointer()), rsp);
+  movePtr(ImmPtr((const void*)wasm::FailInstanceReg), InstanceReg);
   masm.ret();
 
   // Found a wasm catch handler, restore state and jump to it.

js/src/jit/x86/MacroAssembler-x86.cpp

@@ -632,6 +632,7 @@ void MacroAssemblerX86::handleFailureWithHandlerTail(Label* profilerExitTail,
   bind(&wasm);
   loadPtr(Address(esp, ResumeFromException::offsetOfFramePointer()), ebp);
   loadPtr(Address(esp, ResumeFromException::offsetOfStackPointer()), esp);
+  movePtr(ImmPtr((const void*)wasm::FailInstanceReg), InstanceReg);
   masm.ret();
 
   // Found a wasm catch handler, restore state and jump to it.

js/src/vm/FrameIter.cpp

@@ -180,7 +180,7 @@ void JitFrameIter::settle() {
 
   if (isWasm()) {
     const wasm::WasmFrameIter& wasmFrame = asWasm();
-    if (!wasmFrame.unwoundJitCallerFP()) {
+    if (!wasmFrame.hasUnwoundJitFrame()) {
       return;
     }
 
@@ -195,7 +195,7 @@ void JitFrameIter::settle() {
     // The wasm iterator has saved the previous jit frame pointer for us.
 
     MOZ_ASSERT(wasmFrame.done());
-    uint8_t* prevFP = wasmFrame.unwoundJitCallerFP();
+    uint8_t* prevFP = wasmFrame.unwoundCallerFP();
     jit::FrameType prevFrameType = wasmFrame.unwoundJitFrameType();
 
     if (mustUnwindActivation_) {

js/src/wasm/WasmBuiltins.cpp

@@ -652,10 +652,11 @@ bool wasm::HandleThrow(JSContext* cx, WasmFrameIter& iter,
              "unwinding clears the trapping state");
 
   // In case of no handler, exit wasm via ret().
-  // FailFP signals to wasm stub to do a failure return.
+  // FailInstanceReg signals to wasm stub to do a failure return.
   rfe->kind = ExceptionResumeKind::Wasm;
-  rfe->framePointer = (uint8_t*)wasm::FailFP;
+  rfe->framePointer = (uint8_t*)iter.unwoundCallerFP();
   rfe->stackPointer = (uint8_t*)iter.unwoundAddressOfReturnAddress();
+  rfe->instance = (Instance*)FailInstanceReg;
   rfe->target = nullptr;
   return false;
 }

js/src/wasm/WasmCodegenConstants.h

@@ -32,10 +32,10 @@ static const unsigned MaxResultsForJitInlineCall = MaxResultsForJitEntry;
 // returned in registers.
 static const unsigned MaxRegisterResults = 1;
 
-// A magic value of the FramePointer to indicate after a return to the entry
+// A magic value of the InstanceReg to indicate after a return to the entry
 // stub that an exception has been caught and that we should throw.
 
-static const unsigned FailFP = 0xbad;
+static const unsigned FailInstanceReg = 0xbad;
 
 // The following thresholds were derived from a microbenchmark. If we begin to
 // ship this optimization for more platforms, we will need to extend this list.

js/src/wasm/WasmFrameIter.cpp

@@ -60,8 +60,8 @@ WasmFrameIter::WasmFrameIter(JitActivation* activation, wasm::Frame* fp)
       lineOrBytecode_(0),
       fp_(fp ? fp : activation->wasmExitFP()),
       instance_(nullptr),
-      unwoundJitCallerFP_(nullptr),
-      unwoundJitFrameType_(jit::FrameType(-1)),
+      unwoundCallerFP_(nullptr),
+      unwoundJitFrameType_(),
       unwind_(Unwind::False),
       unwoundAddressOfReturnAddress_(nullptr),
       resumePCinCurrentFrame_(nullptr) {
@@ -97,7 +97,7 @@ WasmFrameIter::WasmFrameIter(JitActivation* activation, wasm::Frame* fp)
   // was Ion, we can just skip the wasm frames.
 
   popFrame();
-  MOZ_ASSERT(!done() || unwoundJitCallerFP_);
+  MOZ_ASSERT(!done() || unwoundCallerFP_);
 }
 
 bool WasmFrameIter::done() const {
@@ -160,11 +160,11 @@ void WasmFrameIter::popFrame() {
     // Mark the frame as such.
     AssertDirectJitCall(fp_->jitEntryCaller());
 
-    unwoundJitCallerFP_ = fp_->jitEntryCaller();
-    unwoundJitFrameType_ = FrameType::Exit;
+    unwoundCallerFP_ = fp_->jitEntryCaller();
+    unwoundJitFrameType_.emplace(FrameType::Exit);
 
     if (unwind_ == Unwind::True) {
-      activation_->setJSExitFP(unwoundJitCallerFP());
+      activation_->setJSExitFP(unwoundCallerFP());
       unwoundAddressOfReturnAddress_ = fp_->addressOfReturnAddress();
     }
 
@@ -183,6 +183,9 @@ void WasmFrameIter::popFrame() {
   resumePCinCurrentFrame_ = returnAddress;
 
   if (codeRange_->isInterpEntry()) {
+    // Interpreter entry has a simple frame, record FP from it.
+    unwoundCallerFP_ = reinterpret_cast<uint8_t*>(fp_);
+
     fp_ = nullptr;
     code_ = nullptr;
     codeRange_ = nullptr;
@@ -211,15 +214,15 @@ void WasmFrameIter::popFrame() {
     //
     // The next value of FP is just a regular jit frame used as a marker to
     // know that we should transition to a JSJit frame iterator.
-    unwoundJitCallerFP_ = reinterpret_cast<uint8_t*>(fp_);
-    unwoundJitFrameType_ = FrameType::JSJitToWasm;
+    unwoundCallerFP_ = reinterpret_cast<uint8_t*>(fp_);
+    unwoundJitFrameType_.emplace(FrameType::JSJitToWasm);
 
     fp_ = nullptr;
     code_ = nullptr;
     codeRange_ = nullptr;
 
     if (unwind_ == Unwind::True) {
-      activation_->setJSExitFP(unwoundJitCallerFP());
+      activation_->setJSExitFP(unwoundCallerFP());
       unwoundAddressOfReturnAddress_ = prevFP->addressOfReturnAddress();
     }
 
@@ -332,10 +335,14 @@ DebugFrame* WasmFrameIter::debugFrame() const {
   return DebugFrame::from(fp_);
 }
 
+bool WasmFrameIter::hasUnwoundJitFrame() const {
+  return unwoundCallerFP_ && unwoundJitFrameType_.isSome();
+}
+
 jit::FrameType WasmFrameIter::unwoundJitFrameType() const {
-  MOZ_ASSERT(unwoundJitCallerFP_);
-  MOZ_ASSERT(unwoundJitFrameType_ != jit::FrameType(-1));
-  return unwoundJitFrameType_;
+  MOZ_ASSERT(unwoundCallerFP_);
+  MOZ_ASSERT(unwoundJitFrameType_.isSome());
+  return *unwoundJitFrameType_;
 }
 
 uint8_t* WasmFrameIter::resumePCinCurrentFrame() const {
@@ -1350,11 +1357,6 @@ bool js::wasm::StartUnwinding(const RegisterState& registers,
         // incomplete and we can't unwind.
         return false;
       }
-      // On the error return path, FP might be set to FailFP. Ignore these
-      // transient frames.
-      if (intptr_t(fp) == (FailFP & ~ExitFPTag)) {
-        return false;
-      }
       // Set fixedFP to the address of the JitFrameLayout on the stack.
       if (offsetFromEntry < SetFP) {
         fixedFP = reinterpret_cast<uint8_t*>(sp);

js/src/wasm/WasmFrameIter.h

@@ -66,8 +66,8 @@ class WasmFrameIter {
   unsigned lineOrBytecode_;
   Frame* fp_;
   Instance* instance_;
-  uint8_t* unwoundJitCallerFP_;
-  jit::FrameType unwoundJitFrameType_;
+  uint8_t* unwoundCallerFP_;
+  mozilla::Maybe<jit::FrameType> unwoundJitFrameType_;
   Unwind unwind_;
   void** unwoundAddressOfReturnAddress_;
   uint8_t* resumePCinCurrentFrame_;
@@ -93,7 +93,8 @@ class WasmFrameIter {
   bool debugEnabled() const;
   DebugFrame* debugFrame() const;
   jit::FrameType unwoundJitFrameType() const;
-  uint8_t* unwoundJitCallerFP() const { return unwoundJitCallerFP_; }
+  bool hasUnwoundJitFrame() const;
+  uint8_t* unwoundCallerFP() const { return unwoundCallerFP_; }
   Frame* frame() const { return fp_; }
   Instance* instance() const { return instance_; }
 

js/src/wasm/WasmStubs.cpp

@@ -757,9 +757,51 @@ static bool GenerateInterpEntry(MacroAssembler& masm, const FuncExport& fe,
   Register argv = ABINonArgReturnReg0;
   Register scratch = ABINonArgReturnReg1;
 
+  // scratch := SP
+  masm.moveStackPtrTo(scratch);
+
+  // Dynamically align the stack since ABIStackAlignment is not necessarily
+  // WasmStackAlignment. Preserve SP so it can be restored after the call.
+#ifdef JS_CODEGEN_ARM64
+  static_assert(WasmStackAlignment == 16, "ARM64 SP alignment");
+#else
+  masm.andToStackPtr(Imm32(~(WasmStackAlignment - 1)));
+#endif
+  masm.assertStackAlignment(WasmStackAlignment);
+
+  // Create a fake frame: just previous RA and an FP.
+  const size_t FakeFrameSize = 2 * sizeof(void*);
+#ifdef JS_CODEGEN_ARM64
+  masm.Ldr(ARMRegister(ABINonArgReturnReg0, 64),
+           MemOperand(ARMRegister(scratch, 64), nonVolatileRegsPushSize));
+#else
+  masm.Push(Address(scratch, nonVolatileRegsPushSize));
+#endif
+  // Store fake wasm register state. Ensure the frame pointer passed by the C++
+  // caller doesn't have the ExitFPTag bit set to not confuse frame iterators.
+  // This bit shouldn't be set if C++ code is using frame pointers, so this has
+  // no effect on native stack unwinders.
+  masm.andPtr(Imm32(int32_t(~ExitFPTag)), FramePointer);
+#ifdef JS_CODEGEN_ARM64
+  masm.asVIXL().Push(ARMRegister(ABINonArgReturnReg0, 64),
+                     ARMRegister(FramePointer, 64));
+  masm.moveStackPtrTo(FramePointer);
+#else
+  masm.Push(FramePointer);
+#endif
+
+  masm.moveStackPtrTo(FramePointer);
+  masm.setFramePushed(FakeFrameSize);
+#ifdef JS_CODEGEN_ARM64
+  const size_t FakeFramePushed = 0;
+#else
+  const size_t FakeFramePushed = sizeof(void*);
+  masm.Push(scratch);
+#endif
+
   // Read the arguments of wasm::ExportFuncPtr according to the native ABI.
   // The entry stub's frame is 1 word.
-  const unsigned argBase = sizeof(void*) + masm.framePushed();
+  const unsigned argBase = sizeof(void*) + nonVolatileRegsPushSize;
   ABIArgGenerator abi;
   ABIArg arg;
 
@@ -768,9 +810,7 @@ static bool GenerateInterpEntry(MacroAssembler& masm, const FuncExport& fe,
   if (arg.kind() == ABIArg::GPR) {
     masm.movePtr(arg.gpr(), argv);
   } else {
-    masm.loadPtr(
-        Address(masm.getStackPointer(), argBase + arg.offsetFromArgBase()),
-        argv);
+    masm.loadPtr(Address(scratch, argBase + arg.offsetFromArgBase()), argv);
   }
 
   // Arg 2: Instance*
@@ -778,9 +818,8 @@ static bool GenerateInterpEntry(MacroAssembler& masm, const FuncExport& fe,
   if (arg.kind() == ABIArg::GPR) {
     masm.movePtr(arg.gpr(), InstanceReg);
   } else {
-    masm.loadPtr(
-        Address(masm.getStackPointer(), argBase + arg.offsetFromArgBase()),
-        InstanceReg);
+    masm.loadPtr(Address(scratch, argBase + arg.offsetFromArgBase()),
+                 InstanceReg);
   }
 
   WasmPush(masm, InstanceReg);
@@ -788,23 +827,8 @@ static bool GenerateInterpEntry(MacroAssembler& masm, const FuncExport& fe,
   // Save 'argv' on the stack so that we can recover it after the call.
   WasmPush(masm, argv);
 
-  // Since we're about to dynamically align the stack, reset the frame depth
-  // so we can still assert static stack depth balancing.
-  const unsigned framePushedBeforeAlign =
-      nonVolatileRegsPushSize + NumExtraPushed * WasmPushSize;
-
-  MOZ_ASSERT(masm.framePushed() == framePushedBeforeAlign);
-  masm.setFramePushed(0);
-
-  // Dynamically align the stack since ABIStackAlignment is not necessarily
-  // WasmStackAlignment. Preserve SP so it can be restored after the call.
-#ifdef JS_CODEGEN_ARM64
-  static_assert(WasmStackAlignment == 16, "ARM64 SP alignment");
-#else
-  masm.moveStackPtrTo(scratch);
-  masm.andToStackPtr(Imm32(~(WasmStackAlignment - 1)));
-  masm.Push(scratch);
-#endif
+  MOZ_ASSERT(masm.framePushed() ==
+             NumExtraPushed * WasmPushSize + FakeFrameSize + FakeFramePushed);
 
   // Reserve stack space for the wasm call.
   unsigned argDecrement =
@@ -815,12 +839,6 @@ static bool GenerateInterpEntry(MacroAssembler& masm, const FuncExport& fe,
   // Copy parameters out of argv and into the wasm ABI registers/stack-slots.
   SetupABIArguments(masm, fe, funcType, argv, scratch);
 
-  // Setup wasm register state. Ensure the frame pointer passed by the C++
-  // caller doesn't have the ExitFPTag bit set to not confuse frame iterators.
-  // This bit shouldn't be set if C++ code is using frame pointers, so this has
-  // no effect on native stack unwinders.
-  masm.andPtr(Imm32(int32_t(~ExitFPTag)), FramePointer);
-
   masm.loadWasmPinnedRegsFromInstance();
 
   masm.storePtr(InstanceReg, Address(masm.getStackPointer(),
@@ -832,39 +850,44 @@ static bool GenerateInterpEntry(MacroAssembler& masm, const FuncExport& fe,
   CallFuncExport(masm, fe, funcPtr);
   masm.assertStackAlignment(WasmStackAlignment);
 
+  // Set the return value based on whether InstanceReg is the FailInstanceReg
+  // magic value (set by the throw stub).
+  Label success, join;
+  masm.branchPtr(Assembler::NotEqual, InstanceReg, Imm32(FailInstanceReg),
+                 &success);
+  masm.move32(Imm32(false), scratch);
+  masm.jump(&join);
+  masm.bind(&success);
+  masm.move32(Imm32(true), scratch);
+  masm.bind(&join);
+
   // Pop the arguments pushed after the dynamic alignment.
   masm.freeStack(argDecrement);
 
-  // Pop the stack pointer to its value right before dynamic alignment.
-#ifdef JS_CODEGEN_ARM64
-  static_assert(WasmStackAlignment == 16, "ARM64 SP alignment");
-#else
-  masm.PopStackPtr();
-#endif
-  MOZ_ASSERT(masm.framePushed() == 0);
-  masm.setFramePushed(framePushedBeforeAlign);
+  masm.setFramePushed(NumExtraPushed * WasmPushSize + FakeFrameSize +
+                      FakeFramePushed);
 
   // Recover the 'argv' pointer which was saved before aligning the stack.
   WasmPop(masm, argv);
 
   WasmPop(masm, InstanceReg);
 
+  // Pop the stack pointer to its value right before dynamic alignment.
+#ifdef JS_CODEGEN_ARM64
+  static_assert(WasmStackAlignment == 16, "ARM64 SP alignment");
+  masm.freeStack(FakeFrameSize);
+#else
+  masm.PopStackPtr();
+#endif
+
   // Store the register result, if any, in argv[0].
   // No widening is required, as the value leaves ReturnReg.
   StoreRegisterResult(masm, fe, funcType, argv);
 
-  // After the ReturnReg is stored into argv[0] but before fp is clobbered by
-  // the PopRegsInMask(NonVolatileRegs) below, set the return value based on
-  // whether fp is the FailFP magic value (set by the throw stub).
-  Label success, join;
-  masm.branchPtr(Assembler::NotEqual, FramePointer, Imm32(FailFP), &success);
-  masm.move32(Imm32(false), ReturnReg);
-  masm.jump(&join);
-  masm.bind(&success);
-  masm.move32(Imm32(true), ReturnReg);
-  masm.bind(&join);
+  masm.move32(scratch, ReturnReg);
 
   // Restore clobbered non-volatile registers of the caller.
+  masm.setFramePushed(nonVolatileRegsPushSize);
   masm.PopRegsInMask(NonVolatileRegs);
   MOZ_ASSERT(masm.framePushed() == 0);
 
@@ -920,10 +943,6 @@ static void GenerateJitEntryThrow(MacroAssembler& masm, unsigned frameSize) {
   masm.freeStack(frameSize);
   MoveSPForJitABI(masm);
 
-  // The frame pointer is still set to FailFP. Restore it before using it to
-  // load the instance.
-  masm.moveStackPtrTo(FramePointer);
-
   GenerateJitEntryLoadInstance(masm);
 
   masm.loadPtr(Address(InstanceReg, Instance::offsetOfCx()), ScratchIonEntry);
@@ -1274,17 +1293,18 @@ static bool GenerateJitEntry(MacroAssembler& masm, size_t funcExportIndex,
   masm.storePtr(InstanceReg, Address(masm.getStackPointer(),
                                      WasmCalleeInstanceOffsetBeforeCall));
 
-  // Call into the real function. Note that, due to the throw stub, fp, instance
+  // Call into the real function. Note that, due to the throw stub, instance
   // and pinned registers may be clobbered.
   masm.assertStackAlignment(WasmStackAlignment);
   CallFuncExport(masm, fe, funcPtr);
   masm.assertStackAlignment(WasmStackAlignment);
 
-  // If fp is equal to the FailFP magic value (set by the throw stub), then
-  // report the exception to the JIT caller by jumping into the exception
-  // stub; otherwise the FP value is still set to the parent ion frame value.
+  // If InstanceReg is equal to the FailInstanceReg magic value (set by the
+  // throw stub), then report the exception to the JIT caller by jumping into
+  // the exception stub.
   Label exception;
-  masm.branchPtr(Assembler::Equal, FramePointer, Imm32(FailFP), &exception);
+  masm.branchPtr(Assembler::Equal, InstanceReg, Imm32(FailInstanceReg),
+                 &exception);
 
   // Pop arguments.
   masm.freeStack(frameSizeExclFP);
@@ -1582,7 +1602,7 @@ void wasm::GenerateDirectCallFromJit(MacroAssembler& masm, const FuncExport& fe,
 #endif
   masm.assertStackAlignment(WasmStackAlignment);
 
-  masm.branchPtr(Assembler::Equal, FramePointer, Imm32(wasm::FailFP),
+  masm.branchPtr(Assembler::Equal, InstanceReg, Imm32(wasm::FailInstanceReg),
                  masm.exceptionLabel());
 
   // Store the return value in the appropriate place.
@@ -1628,10 +1648,8 @@ void wasm::GenerateDirectCallFromJit(MacroAssembler& masm, const FuncExport& fe,
 
   GenPrintf(DebugChannel::Function, masm, "\n");
 
-  // Restore the frame pointer by loading it from the ExitFrameLayout.
-  size_t fpOffset = bytesNeeded + ExitFooterFrame::Size() +
-                    ExitFrameLayout::offsetOfCallerFramePtr();
-  masm.loadPtr(Address(masm.getStackPointer(), fpOffset), FramePointer);
+  // Restore the frame pointer.
+  masm.loadPtr(Address(FramePointer, 0), FramePointer);
 
   // Free args + frame descriptor.
   masm.leaveExitFrame(bytesNeeded + ExitFrameLayout::Size());
@@ -2853,6 +2871,8 @@ static bool GenerateThrowStub(MacroAssembler& masm, Label* throwLabel,
   masm.bind(&leaveWasm);
   masm.loadPtr(Address(ReturnReg, ResumeFromException::offsetOfFramePointer()),
                FramePointer);
+  masm.loadPtr(Address(ReturnReg, ResumeFromException::offsetOfInstance()),
+               InstanceReg);
   masm.loadPtr(Address(ReturnReg, ResumeFromException::offsetOfStackPointer()),
                scratch1);
   masm.moveToStackPtr(scratch1);