[PATCH] arm64: use sregs for syscall probe registers

[Frank van der Linden <fllinden@amazon.com>]

Since syscall wrappers are now active on arm64 (4.19+), arguments need to
be retrieved the right way, by checking if there is a saved set of system
call registers, and using them if there are.
---
 tapset/arm64/registers.stp | 50 ++++++++++++++++++++++++++++------------------
 1 file changed, 31 insertions(+), 19 deletions(-)

diff --git a/tapset/arm64/registers.stp b/tapset/arm64/registers.stp
index b2e56495d..b001b8efe 100644
--- a/tapset/arm64/registers.stp
+++ b/tapset/arm64/registers.stp
@@ -107,6 +107,36 @@ function u_register:long (name:string) {
 	return _stp_register(name, 0)
 }
 
+function _stp_arg_register:long (argnum:long) %{ /* pure */
+	long val;
+	struct pt_regs *regs;
+
+	if (STAP_ARG_argnum < 1 || STAP_ARG_argnum > 8) {
+		snprintf(CONTEXT->error_buffer, sizeof(CONTEXT->error_buffer),
+				"Cannot access arg(%lld)",
+				(long long)STAP_ARG_argnum);
+		CONTEXT->last_error = CONTEXT->error_buffer;
+		return;
+	}
+
+	/* syscall-in-pt_regs mode, 4.19+ */
+	if (CONTEXT->sregs) {
+		regs = CONTEXT->sregs;
+	} else {
+		regs = (CONTEXT->user_mode_p ? CONTEXT->uregs : CONTEXT->kregs);
+	}
+
+	if (!regs) {
+		CONTEXT->last_error = "No registers available in this context";
+		return;
+	}
+
+	memcpy(&val, ((char *)regs) + ((STAP_ARG_argnum - 1) * sizeof (long)),
+	    sizeof (long));
+
+	STAP_RETVALUE = val;
+%}
+
 /*
  * Return the value of function arg #argnum (1=first arg).
  * If truncate=1, mask off the top 32 bits.
@@ -115,25 +145,7 @@ function u_register:long (name:string) {
  * TODO: 32-bit arm code has different calling conventions than arm64
  */
 function _stp_arg:long (argnum:long, sign_extend:long, truncate:long) {
-	val = 0
-  assert(!(argnum < 1 || argnum > 8), sprintf("Cannot access arg(%d)", argnum))
-
-	if (argnum == 1)
-		val = u_register("x0")
-	else if (argnum == 2)
-		val = u_register("x1")
-	else if (argnum == 3)
-		val = u_register("x2")
-	else if (argnum == 4)
-		val = u_register("x3")
-	else if (argnum == 5)
-		val = u_register("x4")
-	else if (argnum == 6)
-		val = u_register("x5")
-	else if (argnum == 7)
-		val = u_register("x6")
-	else if (argnum == 8)
-		val = u_register("x7")
+	val = _stp_arg_register(argnum)
 
 	if (truncate) {
 		if (sign_extend)
-- 
2.16.5