Fix bugs in probe code to change sp on x86, x86_64 and 32-bit ARM.
authormark.lam@apple.com <mark.lam@apple.com@268f45cc-cd09-0410-ab3c-d52691b4dbfc>
Tue, 25 Jul 2017 20:56:04 +0000 (20:56 +0000)
committermark.lam@apple.com <mark.lam@apple.com@268f45cc-cd09-0410-ab3c-d52691b4dbfc>
Tue, 25 Jul 2017 20:56:04 +0000 (20:56 +0000)
https://bugs.webkit.org/show_bug.cgi?id=174809
<rdar://problem/33504759>

Reviewed by Filip Pizlo.

1. When the probe handler function changes the sp register to point to the
   region of stack in the middle of the ProbeContext on the stack, there is a
   bug where the ProbeContext's register values to be restored can be over-written
   before they can be restored.  This is now fixed.

2. Added more robust probe tests for changing the sp register.

3. Made existing probe tests to ensure that probe handlers were actually called.

4. Added some verification to testProbePreservesGPRS().

5. Change all the probe tests to fail early on discovering an error instead of
   batching till the end of the test.  This helps point a finger to the failing
   issue earlier.

This patch was tested on x86, x86_64, and ARMv7.  ARM64 probe code will be fixed
next in https://bugs.webkit.org/show_bug.cgi?id=174697.

* assembler/MacroAssemblerARM.cpp:
* assembler/MacroAssemblerARMv7.cpp:
* assembler/MacroAssemblerX86Common.cpp:
* assembler/testmasm.cpp:
(JSC::testProbeReadsArgumentRegisters):
(JSC::testProbeWritesArgumentRegisters):
(JSC::testProbePreservesGPRS):
(JSC::testProbeModifiesStackPointer):
(JSC::testProbeModifiesStackPointerToInsideProbeContextOnStack):
(JSC::testProbeModifiesStackPointerToNBytesBelowSP):
(JSC::testProbeModifiesProgramCounter):
(JSC::run):

git-svn-id: https://svn.webkit.org/repository/webkit/trunk@219885 268f45cc-cd09-0410-ab3c-d52691b4dbfc

Source/JavaScriptCore/ChangeLog
Source/JavaScriptCore/assembler/MacroAssemblerARM.cpp
Source/JavaScriptCore/assembler/MacroAssemblerARMv7.cpp
Source/JavaScriptCore/assembler/MacroAssemblerX86Common.cpp
Source/JavaScriptCore/assembler/testmasm.cpp

index f280857..cbb5a6b 100644 (file)
@@ -1,3 +1,42 @@
+2017-07-25  Mark Lam  <mark.lam@apple.com>
+
+        Fix bugs in probe code to change sp on x86, x86_64 and 32-bit ARM.
+        https://bugs.webkit.org/show_bug.cgi?id=174809
+        <rdar://problem/33504759>
+
+        Reviewed by Filip Pizlo.
+
+        1. When the probe handler function changes the sp register to point to the
+           region of stack in the middle of the ProbeContext on the stack, there is a
+           bug where the ProbeContext's register values to be restored can be over-written
+           before they can be restored.  This is now fixed.
+
+        2. Added more robust probe tests for changing the sp register.
+
+        3. Made existing probe tests to ensure that probe handlers were actually called.
+
+        4. Added some verification to testProbePreservesGPRS().
+
+        5. Change all the probe tests to fail early on discovering an error instead of
+           batching till the end of the test.  This helps point a finger to the failing
+           issue earlier.
+
+        This patch was tested on x86, x86_64, and ARMv7.  ARM64 probe code will be fixed
+        next in https://bugs.webkit.org/show_bug.cgi?id=174697.
+
+        * assembler/MacroAssemblerARM.cpp:
+        * assembler/MacroAssemblerARMv7.cpp:
+        * assembler/MacroAssemblerX86Common.cpp:
+        * assembler/testmasm.cpp:
+        (JSC::testProbeReadsArgumentRegisters):
+        (JSC::testProbeWritesArgumentRegisters):
+        (JSC::testProbePreservesGPRS):
+        (JSC::testProbeModifiesStackPointer):
+        (JSC::testProbeModifiesStackPointerToInsideProbeContextOnStack):
+        (JSC::testProbeModifiesStackPointerToNBytesBelowSP):
+        (JSC::testProbeModifiesProgramCounter):
+        (JSC::run):
+
 2017-07-25  Brian Burg  <bburg@apple.com>
 
         Web Automation: add support for uploading files
index 822a65b..9ff2023 100644 (file)
@@ -215,8 +215,8 @@ asm (
 
     // MacroAssemblerARM::probe() has already generated code to store some values.
     // The top of stack now looks like this:
-    //     esp[0 * ptrSize]: probeFunction
-    //     esp[1 * ptrSize]: arg
+    //     esp[0 * ptrSize]: probe handler function
+    //     esp[1 * ptrSize]: probe arg
     //     esp[2 * ptrSize]: saved r3 / S0
     //     esp[3 * ptrSize]: saved ip
     //     esp[4 * ptrSize]: saved lr
@@ -288,38 +288,46 @@ asm (
     //     pc from there.
     //
     // 2. Issue 1 means we will need to write to the stack location at
-    //    ProbeContext.cpu.sp - 4. But if the user probe function had  modified
-    //    the value of ProbeContext.cpu.sp to point in the range between
-    //    &ProbeContext.cpu.ip thru &ProbeContext.cpu.aspr, then the action for
-    //    Issue 1 may trash the values to be restored before we can restore
-    //    them.
+    //    ProbeContext.cpu.gprs[sp] - PTR_SIZE. But if the user probe function had
+    //    modified the value of ProbeContext.cpu.gprs[sp] to point in the range between
+    //    &ProbeContext.cpu.gprs[ip] thru &ProbeContext.cpu.sprs[aspr], then the action
+    //    for Issue 1 may trash the values to be restored before we can restore them.
     //
-    //    The solution is to check if ProbeContext.cpu.sp contains a value in
+    //    The solution is to check if ProbeContext.cpu.gprs[sp] contains a value in
     //    the undesirable range. If so, we copy the remaining ProbeContext
-    //    register data to a safe range (at memory lower than where
-    //    ProbeContext.cpu.sp points) first, and restore the remaining register
-    //    from this new range.
+    //    register data to a safe area first, and restore the remaining register
+    //    from this new safe area.
 
-    "add       ip, sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_APSR_OFFSET) "\n"
+    // The restore area for the pc will be located at 1 word below the resultant sp.
+    // All restore values are located at offset <= PROBE_CPU_APSR_OFFSET. Hence,
+    // we need to make sure that resultant sp > offset of apsr + 1.
+    "add       ip, sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_APSR_OFFSET + PTR_SIZE) "\n"
     "ldr       lr, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_SP_OFFSET) "]" "\n"
     "cmp       lr, ip" "\n"
     "bgt     " SYMBOL_STRING(ctiMasmProbeTrampolineEnd) "\n"
 
-    // We get here because the new expected stack pointer location is lower
-    // than where it's supposed to be. This means the safe range of stack
-    // memory where we'll be copying the remaining register restore values to
-    // might be in a region of memory below the sp i.e. unallocated stack
-    // memory. This in turn makes it vulnerable to interrupts potentially
-    // trashing the copied values. To prevent that, we must first allocate the
-    // needed stack memory by adjusting the sp before the copying.
+    // Getting here means that the restore area will overlap the ProbeContext data
+    // that we will need to get the restoration values from. So, let's move that
+    // data to a safe place before we start writing into the restore area.
+    // Let's locate the "safe area" at 2x sizeof(ProbeContext) below where the
+    // restore area. This ensures that:
+    // 1. The safe area does not overlap the restore area.
+    // 2. The safe area does not overlap the ProbeContext.
+    //    This makes it so that we can use memcpy (does not require memmove) semantics
+    //    to copy the restore values to the safe area.
+    
+    // lr already contains [sp, #STRINGIZE_VALUE_OF(PROBE_CPU_SP_OFFSET)].
+    "sub       lr, lr, #(2 * " STRINGIZE_VALUE_OF(PROBE_ALIGNED_SIZE) ")" "\n"
 
-    "sub       lr, lr, #(6 * " STRINGIZE_VALUE_OF(PTR_SIZE)
-    " + " STRINGIZE_VALUE_OF(PROBE_CPU_IP_OFFSET) ")" "\n"
+    "mov       ip, sp" "\n" // Save the original ProbeContext*.
 
-    "mov       ip, sp" "\n"
-    "mov       sp, lr" "\n"
-    "mov       lr, ip" "\n"
+    // Make sure the stack pointer points to the safe area. This ensures that the
+    // safe area is protected from interrupt handlers overwriting it.
+    "mov       sp, lr" "\n" // sp now points to the new ProbeContext in the safe area.
+
+    "mov       lr, ip" "\n" // Use lr as the old ProbeContext*.
 
+    // Copy the restore data to the new ProbeContext*.
     "ldr       ip, [lr, #" STRINGIZE_VALUE_OF(PROBE_CPU_IP_OFFSET) "]" "\n"
     "str       ip, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_IP_OFFSET) "]" "\n"
     "ldr       ip, [lr, #" STRINGIZE_VALUE_OF(PROBE_CPU_SP_OFFSET) "]" "\n"
@@ -331,20 +339,31 @@ asm (
     "ldr       ip, [lr, #" STRINGIZE_VALUE_OF(PROBE_CPU_APSR_OFFSET) "]" "\n"
     "str       ip, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_APSR_OFFSET) "]" "\n"
 
+    // ctiMasmProbeTrampolineEnd expects lr to contain the sp value to be restored.
+    // Since we used it as scratch above, let's restore it.
+    "ldr       lr, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_SP_OFFSET) "]" "\n"
+
     SYMBOL_STRING(ctiMasmProbeTrampolineEnd) ":" "\n"
+
+    // Set up the restore area for sp and pc.
+    // lr already contains [sp, #STRINGIZE_VALUE_OF(PROBE_CPU_SP_OFFSET)].
+
+    // Push the pc on to the restore area.
     "ldr       ip, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_PC_OFFSET) "]" "\n"
-    "ldr       lr, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_SP_OFFSET) "]" "\n"
     "sub       lr, lr, #" STRINGIZE_VALUE_OF(PTR_SIZE) "\n"
     "str       ip, [lr]" "\n"
+    // Point sp to the restore area.
     "str       lr, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_SP_OFFSET) "]" "\n"
 
+    // All done with math i.e. won't trash the status register (apsr) and don't need
+    // scratch registers (lr and ip) anymore. Restore apsr, lr, and ip.
     "ldr       ip, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_APSR_OFFSET) "]" "\n"
     "msr       APSR, ip" "\n"
-    "ldr       ip, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_LR_OFFSET) "]" "\n"
-    "mov       lr, ip" "\n"
+    "ldr       lr, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_LR_OFFSET) "]" "\n"
     "ldr       ip, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_IP_OFFSET) "]" "\n"
-    "ldr       sp, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_SP_OFFSET) "]" "\n"
 
+    // Restore the sp and pc.
+    "ldr       sp, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_SP_OFFSET) "]" "\n"
     "pop       { pc }" "\n"
 );
 #endif // COMPILER(GCC_OR_CLANG)
index 1aa6273..d193391 100644 (file)
@@ -186,8 +186,8 @@ asm (
 
     // MacroAssemblerARMv7::probe() has already generated code to store some values.
     // The top of stack now looks like this:
-    //     esp[0 * ptrSize]: probeFunction
-    //     esp[1 * ptrSize]: arg
+    //     esp[0 * ptrSize]: probe handler function
+    //     esp[1 * ptrSize]: probe arg
     //     esp[2 * ptrSize]: saved r0
     //     esp[3 * ptrSize]: saved ip
     //     esp[4 * ptrSize]: saved lr
@@ -254,47 +254,55 @@ asm (
     //
     // 1. Normal ARM calling convention relies on moving lr to pc to return to
     //    the caller. In our case, the address to return to is specified by
-    //    ProbeContext.cpu.pc. And at that moment, we won't have any available
+    //    ProbeContext.cpu.gprs[pc]. And at that moment, we won't have any available
     //    scratch registers to hold the return address (lr needs to hold
-    //    ProbeContext.cpu.lr, not the return address).
+    //    ProbeContext.cpu.gprs[lr], not the return address).
     //
     //    The solution is to store the return address on the stack and load the
     //    pc from there.
     //
     // 2. Issue 1 means we will need to write to the stack location at
-    //    ProbeContext.cpu.sp - 4. But if the user probe function had modified
-    //    the value of ProbeContext.cpu.sp to point in the range between
-    //    &ProbeContext.cpu.ip thru &ProbeContext.cpu.aspr, then the action for
-    //    Issue 1 may trash the values to be restored before we can restore
-    //    them.
+    //    ProbeContext.cpu.gprs[sp] - PTR_SIZE. But if the user probe function had
+    //    modified the value of ProbeContext.cpu.gprs[sp] to point in the range between
+    //    &ProbeContext.cpu.gprs[ip] thru &ProbeContext.cpu.sprs[aspr], then the action
+    //    for Issue 1 may trash the values to be restored before we can restore them.
     //
-    //    The solution is to check if ProbeContext.cpu.sp contains a value in
+    //    The solution is to check if ProbeContext.cpu.gprs[sp] contains a value in
     //    the undesirable range. If so, we copy the remaining ProbeContext
-    //    register data to a safe range (at memory lower than where
-    //    ProbeContext.cpu.sp points) first, and restore the remaining register
-    //    from this new range.
+    //    register data to a safe area first, and restore the remaining register
+    //    from this new safe area.
 
-    "add       ip, sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_APSR_OFFSET) "\n"
+    // The restore area for the pc will be located at 1 word below the resultant sp.
+    // All restore values are located at offset <= PROBE_CPU_APSR_OFFSET. Hence,
+    // we need to make sure that resultant sp > offset of apsr + 1.
+    "add       ip, sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_APSR_OFFSET + PTR_SIZE) "\n"
     "ldr       lr, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_SP_OFFSET) "]" "\n"
     "cmp       lr, ip" "\n"
     "it        gt" "\n"
     "bgt     " SYMBOL_STRING(ctiMasmProbeTrampolineEnd) "\n"
 
-    // We get here because the new expected stack pointer location is lower
-    // than where it's supposed to be. This means the safe range of stack
-    // memory where we'll be copying the remaining register restore values to
-    // might be in a region of memory below the sp i.e. unallocated stack
-    // memory. This, in turn, makes it vulnerable to interrupts potentially
-    // trashing the copied values. To prevent that, we must first allocate the
-    // needed stack memory by adjusting the sp before the copying.
-
-    "sub       lr, lr, #(6 * " STRINGIZE_VALUE_OF(PTR_SIZE)
-    " + " STRINGIZE_VALUE_OF(PROBE_CPU_IP_OFFSET) ")" "\n"
-
-    "mov       ip, sp" "\n"
-    "mov       sp, lr" "\n"
-    "mov       lr, ip" "\n"
-
+    // Getting here means that the restore area will overlap the ProbeContext data
+    // that we will need to get the restoration values from. So, let's move that
+    // data to a safe place before we start writing into the restore area.
+    // Let's locate the "safe area" at 2x sizeof(ProbeContext) below where the
+    // restore area. This ensures that:
+    // 1. The safe area does not overlap the restore area.
+    // 2. The safe area does not overlap the ProbeContext.
+    //    This makes it so that we can use memcpy (does not require memmove) semantics
+    //    to copy the restore values to the safe area.
+
+    // lr already contains [sp, #STRINGIZE_VALUE_OF(PROBE_CPU_SP_OFFSET)].
+    "sub       lr, lr, #(2 * " STRINGIZE_VALUE_OF(PROBE_ALIGNED_SIZE) ")" "\n"
+    
+    "mov       ip, sp" "\n" // Save the original ProbeContext*.
+    
+    // Make sure the stack pointer points to the safe area. This ensures that the
+    // safe area is protected from interrupt handlers overwriting it.
+    "mov       sp, lr" "\n" // sp now points to the new ProbeContext in the safe area.
+    
+    "mov       lr, ip" "\n" // Use lr as the old ProbeContext*.
+    
+    // Copy the restore data to the new ProbeContext*.
     "ldr       ip, [lr, #" STRINGIZE_VALUE_OF(PROBE_CPU_IP_OFFSET) "]" "\n"
     "str       ip, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_IP_OFFSET) "]" "\n"
     "ldr       ip, [lr, #" STRINGIZE_VALUE_OF(PROBE_CPU_SP_OFFSET) "]" "\n"
@@ -306,21 +314,32 @@ asm (
     "ldr       ip, [lr, #" STRINGIZE_VALUE_OF(PROBE_CPU_APSR_OFFSET) "]" "\n"
     "str       ip, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_APSR_OFFSET) "]" "\n"
 
+    // ctiMasmProbeTrampolineEnd expects lr to contain the sp value to be restored.
+    // Since we used it as scratch above, let's restore it.
+    "ldr       lr, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_SP_OFFSET) "]" "\n"
+
     ".thumb_func " THUMB_FUNC_PARAM(ctiMasmProbeTrampolineEnd) "\n"
     SYMBOL_STRING(ctiMasmProbeTrampolineEnd) ":" "\n"
+
+    // Set up the restore area for sp and pc.
+    // lr already contains [sp, #STRINGIZE_VALUE_OF(PROBE_CPU_SP_OFFSET)].
+
+    // Push the pc on to the restore area.
     "ldr       ip, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_PC_OFFSET) "]" "\n"
-    "ldr       lr, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_SP_OFFSET) "]" "\n"
     "sub       lr, lr, #" STRINGIZE_VALUE_OF(PTR_SIZE) "\n"
     "str       ip, [lr]" "\n"
+    // Point sp to the restore area.
     "str       lr, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_SP_OFFSET) "]" "\n"
 
+    // All done with math i.e. won't trash the status register (apsr) and don't need
+    // scratch registers (lr and ip) anymore. Restore apsr, lr, and ip.
     "ldr       ip, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_APSR_OFFSET) "]" "\n"
     "msr       APSR, ip" "\n"
-    "ldr       ip, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_LR_OFFSET) "]" "\n"
-    "mov       lr, ip" "\n"
+    "ldr       lr, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_LR_OFFSET) "]" "\n"
     "ldr       ip, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_IP_OFFSET) "]" "\n"
-    "ldr       sp, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_SP_OFFSET) "]" "\n"
 
+    // Restore the sp and pc.
+    "ldr       sp, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_SP_OFFSET) "]" "\n"
     "pop       { pc }" "\n"
 );
 #endif // COMPILER(GCC_OR_CLANG)
index 34e9218..09655c0 100644 (file)
@@ -171,8 +171,8 @@ asm (
     // this:
     //     esp[0 * ptrSize]: eflags
     //     esp[1 * ptrSize]: return address / saved eip
-    //     esp[2 * ptrSize]: probeFunction
-    //     esp[3 * ptrSize]: arg1
+    //     esp[2 * ptrSize]: probe handler function
+    //     esp[3 * ptrSize]: probe arg
     //     esp[4 * ptrSize]: saved eax
     //     esp[5 * ptrSize]: saved esp
 
@@ -240,35 +240,46 @@ asm (
 
     // There are 6 more registers left to restore:
     //     eax, ecx, ebp, esp, eip, and eflags.
-    // We need to handle these last few restores carefully because:
-    //
-    // 1. We need to push the return address on the stack for ret to use.
-    //    That means we need to write to the stack.
-    // 2. The user probe function may have altered the restore value of esp to
-    //    point to the vicinity of one of the restore values for the remaining
-    //    registers left to be restored.
-    //    That means, for requirement 1, we may end up writing over some of the
-    //    restore values. We can check for this, and first copy the restore
-    //    values to a "safe area" on the stack before commencing with the action
-    //    for requirement 1.
-    // 3. For requirement 2, we need to ensure that the "safe area" is
-    //    protected from interrupt handlers overwriting it. Hence, the esp needs
-    //    to be adjusted to include the "safe area" before we start copying the
-    //    the restore values.
 
+    // The restoration process at ctiMasmProbeTrampolineEnd below works by popping
+    // 5 words off the stack into eflags, eax, ecx, ebp, and eip. These 5 words need
+    // to be pushed on top of the final esp value so that just by popping the 5 words,
+    // we'll get the esp that the probe wants to set. Let's call this area (for storing
+    // these 5 words) the restore area.
+    "movl " STRINGIZE_VALUE_OF(PROBE_CPU_ESP_OFFSET) "(%ebp), %ecx" "\n"
+    "subl $5 * " STRINGIZE_VALUE_OF(PTR_SIZE) ", %ecx" "\n"
+
+    // ecx now points to the restore area.
+
+    // Before we copy values from the ProbeContext to the restore area, we need to
+    // make sure that the restore area does not overlap any of the values that we'll
+    // be copying from in the ProbeContext. All the restore values to be copied from
+    // comes from offset <= PROBE_CPU_EFLAGS_OFFSET in the ProbeContext.
     "movl %ebp, %eax" "\n"
     "addl $" STRINGIZE_VALUE_OF(PROBE_CPU_EFLAGS_OFFSET) ", %eax" "\n"
-    "cmpl %eax, " STRINGIZE_VALUE_OF(PROBE_CPU_ESP_OFFSET) "(%ebp)" "\n"
+    "cmpl %eax, %ecx" "\n"
     "jg " SYMBOL_STRING(ctiMasmProbeTrampolineEnd) "\n"
 
-    // Locate the "safe area" at 2x sizeof(ProbeContext) below where the new
-    // rsp will be. This time we don't have to 32-byte align it because we're
-    // not using this area to store any xmm regs.
-    "movl " STRINGIZE_VALUE_OF(PROBE_CPU_ESP_OFFSET) "(%ebp), %eax" "\n"
+    // Getting here means that the restore area will overlap the ProbeContext data
+    // that we will need to get the restoration values from. So, let's move that
+    // data to a safe place before we start writing into the restore area.
+    // Let's locate the "safe area" at 2x sizeof(ProbeContext) below where the
+    // restore area. This ensures that:
+    // 1. The safe area does not overlap the restore area.
+    // 2. The safe area does not overlap the ProbeContext.
+    //    This makes it so that we can use memcpy (does not require memmove) semantics
+    //    to copy the restore values to the safe area.
+    // Note: the safe area does not have to 32-byte align it because we're not using
+    // it to store any xmm regs.
+    "movl %ecx, %eax" "\n"
     "subl $2 * " STRINGIZE_VALUE_OF(PROBE_ALIGNED_SIZE) ", %eax" "\n"
+
+    // eax now points to the safe area.
+
+    // Make sure the stack pointer points to the safe area. This ensures that the
+    // safe area is protected from interrupt handlers overwriting it.
     "movl %eax, %esp" "\n"
 
-    "subl $" STRINGIZE_VALUE_OF(PROBE_CPU_EAX_OFFSET) ", %eax" "\n"
     "movl " STRINGIZE_VALUE_OF(PROBE_CPU_EAX_OFFSET) "(%ebp), %ecx" "\n"
     "movl %ecx, " STRINGIZE_VALUE_OF(PROBE_CPU_EAX_OFFSET) "(%eax)" "\n"
     "movl " STRINGIZE_VALUE_OF(PROBE_CPU_ECX_OFFSET) "(%ebp), %ecx" "\n"
@@ -283,23 +294,29 @@ asm (
     "movl %ecx, " STRINGIZE_VALUE_OF(PROBE_CPU_EFLAGS_OFFSET) "(%eax)" "\n"
     "movl %eax, %ebp" "\n"
 
-    SYMBOL_STRING(ctiMasmProbeTrampolineEnd) ":" "\n"
-    "movl " STRINGIZE_VALUE_OF(PROBE_CPU_ESP_OFFSET) "(%ebp), %eax" "\n"
-    "subl $5 * " STRINGIZE_VALUE_OF(PTR_SIZE) ", %eax" "\n"
-    // At this point, %esp should be < %eax.
+    // We used ecx above as scratch register. Let's restore it to points to the
+    // restore area.
+    "movl " STRINGIZE_VALUE_OF(PROBE_CPU_ESP_OFFSET) "(%ebp), %ecx" "\n"
+    "subl $5 * " STRINGIZE_VALUE_OF(PTR_SIZE) ", %ecx" "\n"
 
-    "movl " STRINGIZE_VALUE_OF(PROBE_CPU_EFLAGS_OFFSET) "(%ebp), %ecx" "\n"
-    "movl %ecx, 0 * " STRINGIZE_VALUE_OF(PTR_SIZE) "(%eax)" "\n"
-    "movl " STRINGIZE_VALUE_OF(PROBE_CPU_EAX_OFFSET) "(%ebp), %ecx" "\n"
-    "movl %ecx, 1 * " STRINGIZE_VALUE_OF(PTR_SIZE) "(%eax)" "\n"
-    "movl " STRINGIZE_VALUE_OF(PROBE_CPU_ECX_OFFSET) "(%ebp), %ecx" "\n"
-    "movl %ecx, 2 * " STRINGIZE_VALUE_OF(PTR_SIZE) "(%eax)" "\n"
-    "movl " STRINGIZE_VALUE_OF(PROBE_CPU_EBP_OFFSET) "(%ebp), %ecx" "\n"
-    "movl %ecx, 3 * " STRINGIZE_VALUE_OF(PTR_SIZE) "(%eax)" "\n"
-    "movl " STRINGIZE_VALUE_OF(PROBE_CPU_EIP_OFFSET) "(%ebp), %ecx" "\n"
-    "movl %ecx, 4 * " STRINGIZE_VALUE_OF(PTR_SIZE) "(%eax)" "\n"
-    "movl %eax, %esp" "\n"
+    // ecx now points to the restore area.
+
+    SYMBOL_STRING(ctiMasmProbeTrampolineEnd) ":" "\n"
 
+    // Copy remaining restore values from the ProbeContext to the restore area.
+    "movl " STRINGIZE_VALUE_OF(PROBE_CPU_EFLAGS_OFFSET) "(%ebp), %eax" "\n"
+    "movl %eax, 0 * " STRINGIZE_VALUE_OF(PTR_SIZE) "(%ecx)" "\n"
+    "movl " STRINGIZE_VALUE_OF(PROBE_CPU_EAX_OFFSET) "(%ebp), %eax" "\n"
+    "movl %eax, 1 * " STRINGIZE_VALUE_OF(PTR_SIZE) "(%ecx)" "\n"
+    "movl " STRINGIZE_VALUE_OF(PROBE_CPU_ECX_OFFSET) "(%ebp), %eax" "\n"
+    "movl %eax, 2 * " STRINGIZE_VALUE_OF(PTR_SIZE) "(%ecx)" "\n"
+    "movl " STRINGIZE_VALUE_OF(PROBE_CPU_EBP_OFFSET) "(%ebp), %eax" "\n"
+    "movl %eax, 3 * " STRINGIZE_VALUE_OF(PTR_SIZE) "(%ecx)" "\n"
+    "movl " STRINGIZE_VALUE_OF(PROBE_CPU_EIP_OFFSET) "(%ebp), %eax" "\n"
+    "movl %eax, 4 * " STRINGIZE_VALUE_OF(PTR_SIZE) "(%ecx)" "\n"
+    "movl %ecx, %esp" "\n"
+
+    // Do the remaining restoration by popping off the restore area.
     "popfd" "\n"
     "popl %eax" "\n"
     "popl %ecx" "\n"
@@ -321,8 +338,8 @@ asm (
     // this:
     //     esp[0 * ptrSize]: rflags
     //     esp[1 * ptrSize]: return address / saved rip
-    //     esp[2 * ptrSize]: probeFunction
-    //     esp[3 * ptrSize]: arg1
+    //     esp[2 * ptrSize]: probe handler function
+    //     esp[3 * ptrSize]: probe arg
     //     esp[4 * ptrSize]: saved rax
     //     esp[5 * ptrSize]: saved rsp
 
@@ -420,32 +437,44 @@ asm (
 
     // There are 6 more registers left to restore:
     //     rax, rcx, rbp, rsp, rip, and rflags.
-    // We need to handle these last few restores carefully because:
-    //
-    // 1. We need to push the return address on the stack for ret to use
-    //    That means we need to write to the stack.
-    // 2. The user probe function may have altered the restore value of esp to
-    //    point to the vicinity of one of the restore values for the remaining
-    //    registers left to be restored.
-    //    That means, for requirement 1, we may end up writing over some of the
-    //    restore values. We can check for this, and first copy the restore
-    //    values to a "safe area" on the stack before commencing with the action
-    //    for requirement 1.
-    // 3. For both requirement 2, we need to ensure that the "safe area" is
-    //    protected from interrupt handlers overwriting it. Hence, the esp needs
-    //    to be adjusted to include the "safe area" before we start copying the
-    //    the restore values.
 
+    // The restoration process at ctiMasmProbeTrampolineEnd below works by popping
+    // 5 words off the stack into rflags, rax, rcx, rbp, and rip. These 5 words need
+    // to be pushed on top of the final esp value so that just by popping the 5 words,
+    // we'll get the esp that the probe wants to set. Let's call this area (for storing
+    // these 5 words) the restore area.
+    "movq " STRINGIZE_VALUE_OF(PROBE_CPU_ESP_OFFSET) "(%rbp), %rcx" "\n"
+    "subq $5 * " STRINGIZE_VALUE_OF(PTR_SIZE) ", %rcx" "\n"
+
+    // rcx now points to the restore area.
+
+    // Before we copy values from the ProbeContext to the restore area, we need to
+    // make sure that the restore area does not overlap any of the values that we'll
+    // be copying from in the ProbeContext. All the restore values to be copied from
+    // comes from offset <= PROBE_CPU_EFLAGS_OFFSET in the ProbeContext.
     "movq %rbp, %rax" "\n"
     "addq $" STRINGIZE_VALUE_OF(PROBE_CPU_EFLAGS_OFFSET) ", %rax" "\n"
-    "cmpq %rax, " STRINGIZE_VALUE_OF(PROBE_CPU_ESP_OFFSET) "(%rbp)" "\n"
+    "cmpq %rax, %rcx" "\n"
     "jg " SYMBOL_STRING(ctiMasmProbeTrampolineEnd) "\n"
 
-    // Locate the "safe area" at 2x sizeof(ProbeContext) below where the new
-    // rsp will be. This time we don't have to 32-byte align it because we're
-    // not using to store any xmm regs.
-    "movq " STRINGIZE_VALUE_OF(PROBE_CPU_ESP_OFFSET) "(%rbp), %rax" "\n"
+    // Getting here means that the restore area will overlap the ProbeContext data
+    // that we will need to get the restoration values from. So, let's move that
+    // data to a safe place before we start writing into the restore area.
+    // Let's locate the "safe area" at 2x sizeof(ProbeContext) below where the
+    // restore area. This ensures that:
+    // 1. The safe area does not overlap the restore area.
+    // 2. The safe area does not overlap the ProbeContext.
+    //    This makes it so that we can use memcpy (does not require memmove) semantics
+    //    to copy the restore values to the safe area.
+    // Note: the safe area does not have to 32-byte align it because we're not using
+    // it to store any xmm regs.
+    "movq %rcx, %rax" "\n"
     "subq $2 * " STRINGIZE_VALUE_OF(PROBE_ALIGNED_SIZE) ", %rax" "\n"
+
+    // rax now points to the safe area.
+
+    // Make sure the stack pointer points to the safe area. This ensures that the
+    // safe area is protected from interrupt handlers overwriting it.
     "movq %rax, %rsp" "\n"
 
     "movq " STRINGIZE_VALUE_OF(PROBE_CPU_EAX_OFFSET) "(%rbp), %rcx" "\n"
@@ -462,23 +491,29 @@ asm (
     "movq %rcx, " STRINGIZE_VALUE_OF(PROBE_CPU_EFLAGS_OFFSET) "(%rax)" "\n"
     "movq %rax, %rbp" "\n"
 
-    SYMBOL_STRING(ctiMasmProbeTrampolineEnd) ":" "\n"
-    "movq " STRINGIZE_VALUE_OF(PROBE_CPU_ESP_OFFSET) "(%rbp), %rax" "\n"
-    "subq $5 * " STRINGIZE_VALUE_OF(PTR_SIZE) ", %rax" "\n"
-    // At this point, %rsp should be < %rax.
+    // We used rcx above as scratch register. Let's restore it to points to the
+    // restore area.
+    "movq " STRINGIZE_VALUE_OF(PROBE_CPU_ESP_OFFSET) "(%rbp), %rcx" "\n"
+    "subq $5 * " STRINGIZE_VALUE_OF(PTR_SIZE) ", %rcx" "\n"
 
-    "movq " STRINGIZE_VALUE_OF(PROBE_CPU_EFLAGS_OFFSET) "(%rbp), %rcx" "\n"
-    "movq %rcx, 0 * " STRINGIZE_VALUE_OF(PTR_SIZE) "(%rax)" "\n"
-    "movq " STRINGIZE_VALUE_OF(PROBE_CPU_EAX_OFFSET) "(%rbp), %rcx" "\n"
-    "movq %rcx, 1 * " STRINGIZE_VALUE_OF(PTR_SIZE) "(%rax)" "\n"
-    "movq " STRINGIZE_VALUE_OF(PROBE_CPU_ECX_OFFSET) "(%rbp), %rcx" "\n"
-    "movq %rcx, 2 * " STRINGIZE_VALUE_OF(PTR_SIZE) "(%rax)" "\n"
-    "movq " STRINGIZE_VALUE_OF(PROBE_CPU_EBP_OFFSET) "(%rbp), %rcx" "\n"
-    "movq %rcx, 3 * " STRINGIZE_VALUE_OF(PTR_SIZE) "(%rax)" "\n"
-    "movq " STRINGIZE_VALUE_OF(PROBE_CPU_EIP_OFFSET) "(%rbp), %rcx" "\n"
-    "movq %rcx, 4 * " STRINGIZE_VALUE_OF(PTR_SIZE) "(%rax)" "\n"
-    "movq %rax, %rsp" "\n"
+    // rcx now points to the restore area.
+
+    SYMBOL_STRING(ctiMasmProbeTrampolineEnd) ":" "\n"
 
+    // Copy remaining restore values from the ProbeContext to the restore area.
+    "movq " STRINGIZE_VALUE_OF(PROBE_CPU_EFLAGS_OFFSET) "(%rbp), %rax" "\n"
+    "movq %rax, 0 * " STRINGIZE_VALUE_OF(PTR_SIZE) "(%rcx)" "\n"
+    "movq " STRINGIZE_VALUE_OF(PROBE_CPU_EAX_OFFSET) "(%rbp), %rax" "\n"
+    "movq %rax, 1 * " STRINGIZE_VALUE_OF(PTR_SIZE) "(%rcx)" "\n"
+    "movq " STRINGIZE_VALUE_OF(PROBE_CPU_ECX_OFFSET) "(%rbp), %rax" "\n"
+    "movq %rax, 2 * " STRINGIZE_VALUE_OF(PTR_SIZE) "(%rcx)" "\n"
+    "movq " STRINGIZE_VALUE_OF(PROBE_CPU_EBP_OFFSET) "(%rbp), %rax" "\n"
+    "movq %rax, 3 * " STRINGIZE_VALUE_OF(PTR_SIZE) "(%rcx)" "\n"
+    "movq " STRINGIZE_VALUE_OF(PROBE_CPU_EIP_OFFSET) "(%rbp), %rax" "\n"
+    "movq %rax, 4 * " STRINGIZE_VALUE_OF(PTR_SIZE) "(%rcx)" "\n"
+    "movq %rcx, %rsp" "\n"
+
+    // Do the remaining restoration by popping off the restore area.
     "popfq" "\n"
     "popq %rax" "\n"
     "popq %rcx" "\n"
index a7a9736..d0e6d4f 100644 (file)
@@ -33,6 +33,7 @@
 #include "LinkBuffer.h"
 #include <wtf/Compiler.h>
 #include <wtf/DataLog.h>
+#include <wtf/Function.h>
 #include <wtf/Lock.h>
 #include <wtf/NumberOfCores.h>
 #include <wtf/Threading.h>
@@ -55,7 +56,7 @@ namespace {
 
 StaticLock crashLock;
 
-typedef std::function<void(CCallHelpers&)> Generator;
+typedef WTF::Function<void(CCallHelpers&)> Generator;
 
 template<typename T> T nextID(T id) { return static_cast<T>(id + 1); }
 
@@ -102,7 +103,7 @@ bool isFP(MacroAssembler::RegisterID id)
 }
 #endif // ENABLE(MASM_PROBE)
 
-MacroAssemblerCodeRef compile(Generator generate)
+MacroAssemblerCodeRef compile(Generator&& generate)
 {
     CCallHelpers jit;
     generate(jit);
@@ -118,9 +119,9 @@ T invoke(MacroAssemblerCodeRef code, Arguments... arguments)
 }
 
 template<typename T, typename... Arguments>
-T compileAndRun(Generator generator, Arguments... arguments)
+T compileAndRun(Generator&& generator, Arguments... arguments)
 {
-    return invoke<T>(compile(generator), arguments...);
+    return invoke<T>(compile(WTFMove(generator)), arguments...);
 }
 
 void testSimple()
@@ -136,7 +137,7 @@ void testSimple()
 #if ENABLE(MASM_PROBE)
 void testProbeReadsArgumentRegisters()
 {
-    bool success = true;
+    bool probeWasCalled = false;
     compileAndRun<void>([&] (CCallHelpers& jit) {
         jit.emitFunctionPrologue();
 
@@ -157,18 +158,19 @@ void testProbeReadsArgumentRegisters()
 #endif
 
         jit.probe([&] (ProbeContext* context) {
-            success = success && context->gpr(GPRInfo::argumentGPR0) == testWord(0);
-            success = success && context->gpr(GPRInfo::argumentGPR1) == testWord(1);
-            success = success && context->gpr(GPRInfo::argumentGPR2) == testWord(2);
-            success = success && context->gpr(GPRInfo::argumentGPR3) == testWord(3);
-
-            success = success && context->fpr(FPRInfo::fpRegT0) == testWord32(0);
-            success = success && context->fpr(FPRInfo::fpRegT1) == testWord32(1);
+            probeWasCalled = true;
+            CHECK(context->gpr(GPRInfo::argumentGPR0) == testWord(0));
+            CHECK(context->gpr(GPRInfo::argumentGPR1) == testWord(1));
+            CHECK(context->gpr(GPRInfo::argumentGPR2) == testWord(2));
+            CHECK(context->gpr(GPRInfo::argumentGPR3) == testWord(3));
+
+            CHECK(context->fpr(FPRInfo::fpRegT0) == testWord32(0));
+            CHECK(context->fpr(FPRInfo::fpRegT1) == testWord32(1));
         });
         jit.emitFunctionEpilogue();
         jit.ret();
     });
-    CHECK(success);
+    CHECK(probeWasCalled);
 }
 
 void testProbeWritesArgumentRegisters()
@@ -176,7 +178,7 @@ void testProbeWritesArgumentRegisters()
     // This test relies on testProbeReadsArgumentRegisters() having already validated
     // that we can read from argument registers. We'll use that ability to validate
     // that our writes did take effect.
-    bool success = true;
+    unsigned probeCallCount = 0;
     compileAndRun<void>([&] (CCallHelpers& jit) {
         jit.emitFunctionPrologue();
 
@@ -196,7 +198,8 @@ void testProbeWritesArgumentRegisters()
         jit.convertInt32ToDouble(GPRInfo::argumentGPR0, FPRInfo::fpRegT1);
 
         // Write expected values.
-        jit.probe([] (ProbeContext* context) {
+        jit.probe([&] (ProbeContext* context) {
+            probeCallCount++;
             context->gpr(GPRInfo::argumentGPR0) = testWord(0);
             context->gpr(GPRInfo::argumentGPR1) = testWord(1);
             context->gpr(GPRInfo::argumentGPR2) = testWord(2);
@@ -208,19 +211,20 @@ void testProbeWritesArgumentRegisters()
 
         // Validate that expected values were written.
         jit.probe([&] (ProbeContext* context) {
-            success = success && context->gpr(GPRInfo::argumentGPR0) == testWord(0);
-            success = success && context->gpr(GPRInfo::argumentGPR1) == testWord(1);
-            success = success && context->gpr(GPRInfo::argumentGPR2) == testWord(2);
-            success = success && context->gpr(GPRInfo::argumentGPR3) == testWord(3);
-
-            success = success && context->fpr(FPRInfo::fpRegT0) == testWord32(0);
-            success = success && context->fpr(FPRInfo::fpRegT1) == testWord32(1);
+            probeCallCount++;
+            CHECK(context->gpr(GPRInfo::argumentGPR0) == testWord(0));
+            CHECK(context->gpr(GPRInfo::argumentGPR1) == testWord(1));
+            CHECK(context->gpr(GPRInfo::argumentGPR2) == testWord(2));
+            CHECK(context->gpr(GPRInfo::argumentGPR3) == testWord(3));
+
+            CHECK(context->fpr(FPRInfo::fpRegT0) == testWord32(0));
+            CHECK(context->fpr(FPRInfo::fpRegT1) == testWord32(1));
         });
 
         jit.emitFunctionEpilogue();
         jit.ret();
     });
-    CHECK(success);
+    CHECK(probeCallCount == 2);
 }
 
 static NEVER_INLINE NOT_TAIL_CALLED int testFunctionToTrashGPRs(int a, int b, int c, int d, int e, int f, int g, int h, int i, int j)
@@ -241,7 +245,7 @@ void testProbePreservesGPRS()
     // This test relies on testProbeReadsArgumentRegisters() and testProbeWritesArgumentRegisters()
     // having already validated that we can read and write from registers. We'll use these abilities
     // to validate that the probe preserves register values.
-    bool success = true;
+    unsigned probeCallCount = 0;
     MacroAssembler::CPUState originalState;
 
     compileAndRun<void>([&] (CCallHelpers& jit) {
@@ -249,6 +253,7 @@ void testProbePreservesGPRS()
 
         // Write expected values into the registers (except for sp, fp, and pc).
         jit.probe([&] (ProbeContext* context) {
+            probeCallCount++;
             for (auto id = CCallHelpers::firstRegister(); id <= CCallHelpers::lastRegister(); id = nextID(id)) {
                 originalState.gpr(id) = context->gpr(id);
                 if (isPC(id) || isSP(id) || isFP(id))
@@ -263,27 +268,30 @@ void testProbePreservesGPRS()
 
         // Invoke the probe to call a lot of functions and trash register values.
         jit.probe([&] (ProbeContext*) {
-            success = success && (testFunctionToTrashGPRs(0, 1, 2, 3, 4, 5, 6, 7, 8, 9) == 10);
-            success = success && (testFunctionToTrashFPRs(0, 1, 2, 3, 4, 5, 6, 7, 8, 9) == 10);
+            probeCallCount++;
+            CHECK(testFunctionToTrashGPRs(0, 1, 2, 3, 4, 5, 6, 7, 8, 9) == 10);
+            CHECK(testFunctionToTrashFPRs(0, 1, 2, 3, 4, 5, 6, 7, 8, 9) == 10);
         });
 
         // Validate that the registers have the expected values.
         jit.probe([&] (ProbeContext* context) {
+            probeCallCount++;
             for (auto id = CCallHelpers::firstRegister(); id <= CCallHelpers::lastRegister(); id = nextID(id)) {
                 if (isPC(id))
                     continue;
                 if (isSP(id) || isFP(id)) {
-                    success = success && context->gpr(id) == originalState.gpr(id);
+                    CHECK(context->gpr(id) == originalState.gpr(id));
                     continue;
                 }
-                success = success && context->gpr(id) == testWord(id);
+                CHECK(context->gpr(id) == testWord(id));
             }
             for (auto id = CCallHelpers::firstFPRegister(); id <= CCallHelpers::lastFPRegister(); id = nextID(id))
-                success = success && context->fpr(id) == testWord(id);
+                CHECK(context->fpr(id) == testWord(id));
         });
 
         // Restore the original state.
         jit.probe([&] (ProbeContext* context) {
+            probeCallCount++;
             for (auto id = CCallHelpers::firstRegister(); id <= CCallHelpers::lastRegister(); id = nextID(id)) {
                 if (isPC(id) || isSP(id) || isFP(id))
                     continue;
@@ -293,41 +301,145 @@ void testProbePreservesGPRS()
                 context->fpr(id) = originalState.fpr(id);
         });
 
+        // Validate that the original state was restored.
+        jit.probe([&] (ProbeContext* context) {
+            probeCallCount++;
+            for (auto id = CCallHelpers::firstRegister(); id <= CCallHelpers::lastRegister(); id = nextID(id)) {
+                if (isPC(id) || isSP(id) || isFP(id))
+                    continue;
+                CHECK(context->gpr(id) == originalState.gpr(id));
+            }
+            for (auto id = CCallHelpers::firstFPRegister(); id <= CCallHelpers::lastFPRegister(); id = nextID(id))
+                CHECK(context->fpr(id) == originalState.fpr(id));
+        });
+
         jit.emitFunctionEpilogue();
         jit.ret();
     });
-    CHECK(success);
+    CHECK(probeCallCount == 5);
 }
 
-void testProbeModifiesStackPointer()
+void testProbeModifiesStackPointer(WTF::Function<void*(ProbeContext*)> computeModifiedStack)
 {
-    bool success = true;
-    uint8_t* originalSP;
+    unsigned probeCallCount = 0;
+    MacroAssembler::CPUState originalState;
+    uint8_t* originalSP { nullptr };
+    void* modifiedSP { nullptr };
+#if CPU(X86) || CPU(X86_64) || CPU(ARM_THUMB2) || CPU(ARM_TRADITIONAL)
+    uintptr_t modifiedFlags { 0 };
+#endif
 
     compileAndRun<void>([&] (CCallHelpers& jit) {
         jit.emitFunctionPrologue();
 
-        // Preserve original stack pointer and modify the sp.
+        // Preserve original stack pointer and modify the sp, and
+        // write expected values into other registers (except for fp, and pc).
         jit.probe([&] (ProbeContext* context) {
+            probeCallCount++;
+            for (auto id = CCallHelpers::firstRegister(); id <= CCallHelpers::lastRegister(); id = nextID(id)) {
+                originalState.gpr(id) = context->gpr(id);
+                if (isPC(id) || isSP(id) || isFP(id))
+                    continue;
+                context->gpr(id) = testWord(static_cast<int>(id));
+            }
+            for (auto id = CCallHelpers::firstFPRegister(); id <= CCallHelpers::lastFPRegister(); id = nextID(id)) {
+                originalState.fpr(id) = context->fpr(id);
+                context->fpr(id) = testWord(id);
+            }
+#if CPU(X86) || CPU(X86_64)
+            originalState.spr(X86Registers::eflags) = context->spr(X86Registers::eflags);
+            modifiedFlags = originalState.spr(X86Registers::eflags) ^ 0xc5;
+            context->spr(X86Registers::eflags) = modifiedFlags;
+#elif CPU(ARM_THUMB2) || CPU(ARM_TRADITIONAL)
+            originalState.spr(ARMRegisters::apsr) = context->spr(ARMRegisters::apsr);
+            modifiedFlags = originalState.spr(ARMRegisters::apsr) ^ 0xf0000000;
+            context->spr(ARMRegisters::apsr) = modifiedFlags;
+#endif
             originalSP = reinterpret_cast<uint8_t*>(context->sp());
-            context->sp() = originalSP - 1 * KB;
+            modifiedSP = computeModifiedStack(context);
+            context->sp() = modifiedSP;
+        });
+
+        // Validate that the registers have the expected values.
+        jit.probe([&] (ProbeContext* context) {
+            probeCallCount++;
+            for (auto id = CCallHelpers::firstRegister(); id <= CCallHelpers::lastRegister(); id = nextID(id)) {
+                if (isPC(id) || isSP(id))
+                    continue;
+                if (isFP(id)) {
+                    CHECK(context->gpr(id) == originalState.gpr(id));
+                    continue;
+                }
+                CHECK(context->gpr(id) == testWord(id));
+            }
+            for (auto id = CCallHelpers::firstFPRegister(); id <= CCallHelpers::lastFPRegister(); id = nextID(id))
+                CHECK(context->fpr(id) == testWord(id));
+#if CPU(X86) || CPU(X86_64)
+            CHECK(context->spr(X86Registers::eflags) == modifiedFlags);
+#elif CPU(ARM_THUMB2) || CPU(ARM_TRADITIONAL)
+            CHECK(context->spr(ARMRegisters::apsr) == modifiedFlags);
+#endif
+            CHECK(context->sp() == modifiedSP);
         });
 
-        // Validate that the stack pointer has the expected value, and restore the original.
+        // Restore the original state.
         jit.probe([&] (ProbeContext* context) {
-            success = (reinterpret_cast<uint8_t*>(context->sp()) == (originalSP - 1 * KB));
+            probeCallCount++;
+            for (auto id = CCallHelpers::firstRegister(); id <= CCallHelpers::lastRegister(); id = nextID(id)) {
+                if (isPC(id) || isSP(id) || isFP(id))
+                    continue;
+                context->gpr(id) = originalState.gpr(id);
+            }
+            for (auto id = CCallHelpers::firstFPRegister(); id <= CCallHelpers::lastFPRegister(); id = nextID(id))
+                context->fpr(id) = originalState.fpr(id);
+#if CPU(X86) || CPU(X86_64)
+            context->spr(X86Registers::eflags) = originalState.spr(X86Registers::eflags);
+#elif CPU(ARM_THUMB2) || CPU(ARM_TRADITIONAL)
+            context->spr(ARMRegisters::apsr) = originalState.spr(ARMRegisters::apsr);
+#endif
             context->sp() = originalSP;
         });
 
-        // Validate that the original stack pointer was restored.
+        // Validate that the original state was restored.
         jit.probe([&] (ProbeContext* context) {
-            success = (context->sp() == originalSP);
+            probeCallCount++;
+            for (auto id = CCallHelpers::firstRegister(); id <= CCallHelpers::lastRegister(); id = nextID(id)) {
+                if (isPC(id) || isSP(id) || isFP(id))
+                    continue;
+                CHECK(context->gpr(id) == originalState.gpr(id));
+            }
+            for (auto id = CCallHelpers::firstFPRegister(); id <= CCallHelpers::lastFPRegister(); id = nextID(id))
+                CHECK(context->fpr(id) == originalState.fpr(id));
+#if CPU(X86) || CPU(X86_64)
+            CHECK(context->spr(X86Registers::eflags) == originalState.spr(X86Registers::eflags));
+#elif CPU(ARM_THUMB2) || CPU(ARM_TRADITIONAL)
+            CHECK(context->spr(ARMRegisters::apsr) == originalState.spr(ARMRegisters::apsr));
+#endif
+            CHECK(context->sp() == originalSP);
         });
 
         jit.emitFunctionEpilogue();
         jit.ret();
     });
-    CHECK(success);
+    CHECK(probeCallCount == 4);
+}
+
+void testProbeModifiesStackPointerToInsideProbeContextOnStack()
+{
+    for (size_t offset = 0; offset < sizeof(ProbeContext); offset += sizeof(uintptr_t)) {
+        testProbeModifiesStackPointer([=] (ProbeContext* context) -> void* {
+            return reinterpret_cast<uint8_t*>(context) + offset;
+        });
+    }
+}
+
+void testProbeModifiesStackPointerToNBytesBelowSP()
+{
+    for (size_t offset = 0; offset < 1 * KB; offset += sizeof(uintptr_t)) {
+        testProbeModifiesStackPointer([=] (ProbeContext* context) -> void* {
+            return reinterpret_cast<uint8_t*>(context->cpu.sp()) - offset;
+        });
+    }
 }
 
 void testProbeModifiesProgramCounter()
@@ -335,12 +447,14 @@ void testProbeModifiesProgramCounter()
     // This test relies on testProbeReadsArgumentRegisters() and testProbeWritesArgumentRegisters()
     // having already validated that we can read and write from registers. We'll use these abilities
     // to validate that the probe preserves register values.
-    bool success = false;
+    unsigned probeCallCount = 0;
+    bool continuationWasReached = false;
 
     MacroAssemblerCodeRef continuation = compile([&] (CCallHelpers& jit) {
         // Validate that we reached the continuation.
         jit.probe([&] (ProbeContext*) {
-            success = true;
+            probeCallCount++;
+            continuationWasReached = true;
         });
 
         jit.emitFunctionEpilogue();
@@ -352,12 +466,14 @@ void testProbeModifiesProgramCounter()
 
         // Write expected values into the registers.
         jit.probe([&] (ProbeContext* context) {
+            probeCallCount++;
             context->pc() = continuation.code().executableAddress();
         });
 
         jit.breakpoint(); // We should never get here.
     });
-    CHECK(success);
+    CHECK(probeCallCount == 2);
+    CHECK(continuationWasReached);
 }
 #endif // ENABLE(MASM_PROBE)
 
@@ -389,7 +505,8 @@ void run(const char* filter)
     RUN(testProbeReadsArgumentRegisters());
     RUN(testProbeWritesArgumentRegisters());
     RUN(testProbePreservesGPRS());
-    RUN(testProbeModifiesStackPointer());
+    RUN(testProbeModifiesStackPointerToInsideProbeContextOnStack());
+    RUN(testProbeModifiesStackPointerToNBytesBelowSP());
     RUN(testProbeModifiesProgramCounter());
 #endif