From e1fb8024a65e2cefc6113b0c1a7143e807bea121 Mon Sep 17 00:00:00 2001
From: Michal Rostecki <mrostecki@opensuse.org>
Date: Mon, 27 Dec 2021 16:23:33 +0100
Subject: [PATCH] fentry/fexit: Add template for fentry/fexit programs

Signed-off-by: Michal Rostecki <mrostecki@opensuse.org>
---
 .github/workflows/ci.yml          |  8 +++++---
 cargo-generate.toml               | 21 ++++++++++++++++++-
 test.sh                           |  3 +++
 {{project-name}}-ebpf/src/main.rs | 34 +++++++++++++++++++++++++++++++
 {{project-name}}/src/main.rs      | 14 +++++++++++++
 5 files changed, 76 insertions(+), 4 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index d1de7ad..943737c 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -21,6 +21,8 @@ jobs:
         program:
           - kprobe
           - kretprobe
+          - fentry
+          - fexit
           - uprobe
           - uretprobe
           - sock_ops
@@ -38,13 +40,13 @@ jobs:
       - name: Install latest stable
         uses: actions-rs/toolchain@v1
         with:
-            toolchain: stable
+          toolchain: stable
 
       - name: Install latest nightly
         uses: actions-rs/toolchain@v1
         with:
-            toolchain: nightly
-            components: rust-src
+          toolchain: nightly
+          components: rust-src
 
       - name: Install bpf-linker
         run: cargo +nightly install bpf-linker
diff --git a/cargo-generate.toml b/cargo-generate.toml
index ae12a41..e88d8be 100644
--- a/cargo-generate.toml
+++ b/cargo-generate.toml
@@ -5,13 +5,32 @@ ignore = [".github", "test.sh"]
 [placeholders.program_type]
 type = "string"
 prompt = "Which type of eBPF program?"
-choices = ["kprobe", "kretprobe", "uprobe", "uretprobe", "sock_ops", "sk_msg", "xdp", "classifier", "cgroup_skb", "tracepoint", "lsm", "tp_btf"]
+choices = [
+    "kprobe",
+    "kretprobe",
+    "fentry",
+    "fexit",
+    "uprobe",
+    "uretprobe",
+    "sock_ops",
+    "sk_msg",
+    "xdp",
+    "classifier",
+    "cgroup_skb",
+    "tracepoint",
+    "lsm",
+    "tp_btf"
+]
 default = "xdp"
 
 [conditional.'program_type == "kprobe" || program_type == "kretprobe"'.placeholders.kprobe]
 type = "string"
 prompt = "Where to attach the (k|kret)probe? (e.g try_to_wake_up)"
 
+[conditional.'program_type == "fentry" || program_type == "fexit"'.placeholders.fn_name]
+type = "string"
+prompt = "Where to attach the f(entry|exit)? (e.g try_to_wake_up)"
+
 [conditional.'program_type == "uprobe" || program_type == "uretprobe"'.placeholders.uprobe_target]
 type = "string"
 prompt = "Target to attach the (u|uret)probe? (e.g libc)"
diff --git a/test.sh b/test.sh
index dfeafc1..e9d758f 100755
--- a/test.sh
+++ b/test.sh
@@ -18,6 +18,9 @@ case "$PROG_TYPE" in
     "kprobe"|"kretprobe")
         ADDITIONAL_ARGS="-d kprobe=test"
         ;;
+    "fentry"|"fexit")
+        ADDITIONAL_ARGS="-d fn_name=try_to_wake_up"
+        ;;
     "uprobe"|"uretprobe")
         ADDITIONAL_ARGS="-d uprobe_target=testlib -d uprobe_fn_name=testfn"
         ;;
diff --git a/{{project-name}}-ebpf/src/main.rs b/{{project-name}}-ebpf/src/main.rs
index ec64a1b..f676e99 100644
--- a/{{project-name}}-ebpf/src/main.rs
+++ b/{{project-name}}-ebpf/src/main.rs
@@ -35,6 +35,40 @@ pub fn {{crate_name}}(ctx: ProbeContext) -> u32 {
 unsafe fn try_{{crate_name}}(_ctx: ProbeContext) -> Result<u32, u32> {
     Ok(0)
 }
+{%- when "fentry" %}
+use aya_bpf::{
+    macros::fentry,
+    programs::FEntryContext,
+};
+
+#[fentry(name="{{crate_name}}")]
+pub fn {{crate_name}}(ctx: FEntryContext) -> u32 {
+    match unsafe { try_{{crate_name}}(ctx) } {
+        Ok(ret) => ret,
+        Err(ret) => ret,
+    }
+}
+
+unsafe fn try_{{crate_name}}(_ctx: FEntryContext) -> Result<u32, u32> {
+    Ok(0)
+}
+{%- when "fexit" %}
+use aya_bpf::{
+    macros::fexit,
+    programs::FExitContext,
+};
+
+#[fexit(name="{{crate_name}}")]
+pub fn {{crate_name}}(ctx: FExitContext) -> u32 {
+    match unsafe { try_{{crate_name}}(ctx) } {
+        Ok(ret) => ret,
+        Err(ret) => ret,
+    }
+}
+
+unsafe fn try_{{crate_name}}(_ctx: FExitContext) -> Result<u32, u32> {
+    Ok(0)
+}
 {%- when "uprobe" %}
 use aya_bpf::{
     macros::uprobe,
diff --git a/{{project-name}}/src/main.rs b/{{project-name}}/src/main.rs
index b7f35b7..0b775b4 100644
--- a/{{project-name}}/src/main.rs
+++ b/{{project-name}}/src/main.rs
@@ -2,6 +2,10 @@ use aya::{Bpf, include_bytes_aligned};
 {% case program_type -%}
 {%- when "kprobe", "kretprobe" -%}
 use aya::programs::KProbe;
+{%- when "fentry" -%}
+use aya::{programs::FEntry, Btf};
+{%- when "fexit" -%}
+use aya::{programs::FExit, Btf};
 {%- when "uprobe", "uretprobe" -%}
 use aya::programs::UProbe;
 {%- when "sock_ops" -%}
@@ -71,6 +75,16 @@ fn try_main() -> Result<(), anyhow::Error> {
     let program: &mut KProbe = bpf.program_mut("{{crate_name}}").unwrap().try_into()?;
     program.load()?;
     program.attach("{{kprobe}}", 0)?;
+    {%- when "fentry" -%}
+    let btf = Btf::from_sys_fs()?;
+    let program: &mut FEntry = bpf.program_mut("{{crate_name}}").unwrap().try_into()?;
+    program.load("{{fn_name}}", &btf)?;
+    program.attach()?;
+    {%- when "fexit" -%}
+    let btf = Btf::from_sys_fs()?;
+    let program: &mut FExit = bpf.program_mut("{{crate_name}}").unwrap().try_into()?;
+    program.load("{{fn_name}}", &btf)?;
+    program.attach()?;
     {%- when "uprobe", "uretprobe" -%}
     let program: &mut UProbe = bpf.program_mut("{{crate_name}}").unwrap().try_into()?;
     program.load()?;