From b8c43a05dc74a8571d1410ed7df8074baa5e82be Mon Sep 17 00:00:00 2001
From: pdliyan <pdliyan@didiglobal.com>
Date: Tue, 10 Oct 2023 17:31:00 +0800
Subject: [PATCH] Use argument results to decide whether or not to read
 stack_argument.

---
 test/integration-ebpf/src/stack_argument.rs | 24 +++++++++++++++++----
 1 file changed, 20 insertions(+), 4 deletions(-)

diff --git a/test/integration-ebpf/src/stack_argument.rs b/test/integration-ebpf/src/stack_argument.rs
index c642072f..616b5c0e 100644
--- a/test/integration-ebpf/src/stack_argument.rs
+++ b/test/integration-ebpf/src/stack_argument.rs
@@ -12,14 +12,30 @@ static ARGS: HashMap<u32, u64> = HashMap::with_max_entries(24, 0);
 
 #[uprobe]
 pub fn test_stack_argument(ctx: ProbeContext) -> i32 {
-    match try_stack_argument(ctx) {
-        Ok(ret) => ret,
-        Err(_) => 0,
-    }
+    try_stack_argument(ctx).unwrap_or(0)
 }
 
 // read function arguments, and set to map.
 fn try_stack_argument(ctx: ProbeContext) -> Result<i32, i64> {
+    let mut stack = false;
+    let mut arg = 0;
+    loop {
+        if arg > 7 {
+            break;
+        }
+        if stack {
+            let _ = ARGS.insert(&arg, &ctx.arg(arg as usize).ok_or(255)?, 0);
+        } else {
+            let arg_v: Option<u64> = ctx.arg(arg as usize);
+            if let None = arg_v {
+                // assume that we shall read from stack now.
+                stack = true;
+                continue;
+            }
+            let _ = ARGS.insert(&arg, &arg_v.unwrap(), 0);
+        }
+        arg += 1;
+    }
     let _ = ARGS.insert(&0, &ctx.arg(0).ok_or(255)?, 0);
     let _ = ARGS.insert(&1, &ctx.arg(1).ok_or(255)?, 0);
     let _ = ARGS.insert(&2, &ctx.arg(2).ok_or(255)?, 0);