[XRay][compiler-rt][x86_64] Align the stack before and after calling handlers

[dberris]

Summary

This change fixes the XRay trampolines aside from the __xray_CustomEvent
trampoline to align the stack to 16-byte boundaries before calling the
handler. Before this change we've not been explicitly aligning the stack
to 16-byte boundaries, which makes it dangerous when calling handlers
that leave the stack in a state that isn't strictly 16-byte aligned
after calling the handlers.

We add a test that makes sure we can handle these cases appropriately
after the changes, and prevents us from regressing the state moving
forward.

Fixes http://llvm.org/PR35294.

[pcc]

LGTM

compiler-rt/trunk/lib/xray/xray_trampoline_x86_64.S

.macro RESTORE_REGISTERS
   movq  24(%rsp), %rsi
   movq  16(%rsp), %rcx
   movq  8(%rsp), %r8
   movq  0(%rsp), %r9
   addq  $192, %rsp
   .cfi_def_cfa_offset 8
 .endm
 
+.macro ALIGNED_CALL_RAX
+  // Call the logging handler, after aligning the stack to a 16-byte boundary.
+  // The approach we're taking here uses additional stack space to stash the
+  // stack pointer twice before aligning the pointer to 16-bytes. If the stack
+  // was 8-byte aligned, it will become 16-byte aligned -- when restoring the
+  // pointer, we can always look -8 bytes from the current position to get
+  // either of the values we've stashed in the first place.
+  pushq %rsp
+  pushq (%rsp)
+  andq $-0x10, %rsp
+  callq *%rax
+  movq 8(%rsp), %rsp
+.endm
+
   .text
   .file "xray_trampoline_x86.S"
 
 //===----------------------------------------------------------------------===//
 
   .globl __xray_FunctionEntry
   .align 16, 0x90
   .type __xray_FunctionEntry,@function
 
 __xray_FunctionEntry:
   .cfi_startproc
   SAVE_REGISTERS
 
   // This load has to be atomic, it's concurrent with __xray_patch().
   // On x86/amd64, a simple (type-aligned) MOV instruction is enough.
   movq  _ZN6__xray19XRayPatchedFunctionE(%rip), %rax
   testq %rax, %rax
   je  .Ltmp0
 
   // The patched function prolog puts its xray_instr_map index into %r10d.
   movl  %r10d, %edi
   xor %esi,%esi
-  callq *%rax
+  ALIGNED_CALL_RAX
+
 .Ltmp0:
   RESTORE_REGISTERS
   retq
 .Ltmp1:
   .size __xray_FunctionEntry, .Ltmp1-__xray_FunctionEntry
   .cfi_endproc
 
 //===----------------------------------------------------------------------===//
__xray_FunctionExit:
   movq  %rax, 8(%rsp)
   movq  %rdx, 0(%rsp)
   movq  _ZN6__xray19XRayPatchedFunctionE(%rip), %rax
   testq %rax,%rax
   je  .Ltmp2
 
   movl  %r10d, %edi
   movl  $1, %esi
-  callq *%rax
+  ALIGNED_CALL_RAX
+
 .Ltmp2:
   // Restore the important registers.
   movq  48(%rsp), %rbp
   movupd  32(%rsp), %xmm0
   movupd  16(%rsp), %xmm1
   movq  8(%rsp), %rax
   movq  0(%rsp), %rdx
   addq  $56, %rsp
__xray_FunctionTailExit:
   SAVE_REGISTERS
 
   movq  _ZN6__xray19XRayPatchedFunctionE(%rip), %rax
   testq %rax,%rax
   je  .Ltmp4
 
   movl  %r10d, %edi
   movl  $2, %esi
-  callq *%rax
+
+  ALIGNED_CALL_RAX
 
 .Ltmp4:
   RESTORE_REGISTERS
   retq
 .Ltmp5:
   .size __xray_FunctionTailExit, .Ltmp5-__xray_FunctionTailExit
   .cfi_endproc
 
.Larg1entryLog:
   // First argument will become the third
   movq  %rdi, %rdx
 
   // XRayEntryType::LOG_ARGS_ENTRY into the second
   mov $0x3, %esi
 
   // 32-bit function ID becomes the first
   movl  %r10d, %edi
-  callq *%rax
+  ALIGNED_CALL_RAX
 
 .Larg1entryFail:
   RESTORE_REGISTERS
   retq
 
 .Larg1entryEnd:
   .size __xray_ArgLoggerEntry, .Larg1entryEnd-__xray_ArgLoggerEntry
   .cfi_endproc
__xray_CustomEvent:
 
   // We take two arguments to this trampoline, which should be in rdi and rsi
   // already. We also make sure that we stash %rax because we use that register
   // to call the logging handler.
   movq _ZN6__xray22XRayPatchedCustomEventE(%rip), %rax
   testq %rax,%rax
   je .LcustomEventCleanup
 
-  // At this point we know that rcx and rdx already has the data, so we just
-  // call the logging handler, after aligning the stack to a 16-byte boundary.
-  // The approach we're taking here uses additional stack space to stash the
-  // stack pointer twice before aligning the pointer to 16-bytes. If the stack
-  // was 8-byte aligned, it will become 16-byte aligned -- when restoring the
-  // pointer, we can always look -8 bytes from the current position to get
-  // either of the values we've stashed in the first place.
-  pushq %rsp
-  pushq (%rsp)
-  andq $-0x10, %rsp
-  callq *%rax
-  movq 8(%rsp), %rsp
+  ALIGNED_CALL_RAX
 
 .LcustomEventCleanup:
   RESTORE_REGISTERS
   retq
 
 .Ltmp8:
   .size __xray_CustomEvent, .Ltmp8-__xray_CustomEvent
   .cfi_endproc
 
 NO_EXEC_STACK_DIRECTIVE

compiler-rt/trunk/test/xray/TestCases/Linux/common-trampoline-alignment.cc

+// Make sure that we're aligning the stack properly to support handlers that
+// expect 16-byte alignment of the stack.
+//
+// RUN: %clangxx_xray -std=c++11 %s -o %t
+// RUN: XRAY_OPTIONS="patch_premain=false verbosity=1 xray_naive_log=false" \
+// RUN:     %run %t 2>&1
+// REQUIRES: x86_64-linux
+// REQUIRES: built-in-llvm-tree
+#include "xray/xray_interface.h"
+#include <stdio.h>
+#include <xmmintrin.h>
+
+[[clang::xray_never_instrument]] __attribute__((weak)) __m128 f(__m128 *i) {
+  return *i;
+}
+
+[[clang::xray_always_instrument]] __attribute__((noinline)) void noarg() {
+  __m128 v = {};
+  f(&v);
+}
+
+[[ clang::xray_always_instrument, clang::xray_log_args(1) ]]
+__attribute__((noinline)) void arg1(int) {
+  __m128 v = {};
+  f(&v);
+}
+
+[[clang::xray_always_instrument]] __attribute__((noinline))
+void no_alignment() {}
+
+[[clang::xray_never_instrument]] void noarg_handler(int32_t,
+                                                        XRayEntryType) {
+  printf("noarg handler called\n");
+  __m128 v = {};
+  f(&v);
+}
+
+[[clang::xray_never_instrument]] void arg1_handler(int32_t, XRayEntryType,
+                                                   uint64_t) {
+  printf("arg1 handler called\n");
+  __m128 v = {};
+  f(&v);
+}
+
+int main(int argc, char *argv[]) {
+  __xray_set_handler(noarg_handler);
+  __xray_set_handler_arg1(arg1_handler);
+  __xray_patch();
+  noarg();    // CHECK: noarg handler called
+  arg1(argc); // CHECK: arg1 handler called
+  no_alignment();
+  __xray_unpatch();
+  __xray_remove_handler();
+  __xray_remove_handler_arg1();
+  noarg();
+  arg1(argc);
+}