From 7ffb47cc6b712b3daeca0a16dd7a502f6d0743aa Mon Sep 17 00:00:00 2001
From: William Findlay <william@williamfindlay.com>
Date: Wed, 27 Oct 2021 11:05:48 -0400
Subject: [PATCH] template: add a new template for LSM programs

---
 .github/workflows/ci.yml          |  3 ++-
 cargo-generate.toml               |  9 +++++++--
 test.sh                           |  3 +++
 {{project-name}}-ebpf/src/main.rs | 19 ++++++++++++++++++-
 {{project-name}}/src/main.rs      | 10 ++++++++--
 5 files changed, 38 insertions(+), 6 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 32a5305..2fc7790 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -4,7 +4,7 @@ on:
   push:
     branches:
       - main
-   
+
   pull_request:
     branches:
       - main
@@ -29,6 +29,7 @@ jobs:
           - classifier
           - cgroup_skb
           - tracepoint
+          - lsm
 
     steps:
       - uses: actions/checkout@v2
diff --git a/cargo-generate.toml b/cargo-generate.toml
index a8eca88..83fca28 100644
--- a/cargo-generate.toml
+++ b/cargo-generate.toml
@@ -5,7 +5,7 @@ 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"]
+choices = ["kprobe", "kretprobe", "uprobe", "uretprobe", "sock_ops", "sk_msg", "xdp", "classifier", "cgroup_skb", "tracepoint", "lsm"]
 default = "xdp"
 
 [conditional.'program_type == "kprobe" || program_type == "kretprobe"'.placeholders.kprobe]
@@ -38,4 +38,9 @@ regex = "[a-z]+"
 [conditional.'program_type == "tracepoint"'.placeholders.tracepoint_name]
 type = "string"
 prompt = "Which tracepoint name? (e.g sched_switch, net_dev_queue)"
-regex = "[a-z]+"
\ No newline at end of file
+regex = "[a-z]+"
+
+[conditional.'program_type == "lsm"'.placeholders.lsm_hook]
+type = "string"
+prompt = "Which lsm hook? (e.g file_open, task_alloc) You can find a list of hooks in include/linux/lsm_hooks.h in the kernel source tree."
+regex = "[a-z]+"
diff --git a/test.sh b/test.sh
index 7d642d3..353450f 100755
--- a/test.sh
+++ b/test.sh
@@ -30,6 +30,9 @@ case "$PROG_TYPE" in
     "sk_msg")
         ADDITIONAL_ARGS="-d sock_map=TEST"
         ;;
+    "lsm")
+        ADDITIONAL_ARGS="-d lsm_hook=file_open"
+        ;;
     *)
         ADDITIONAL_ARGS=''
 esac
diff --git a/{{project-name}}-ebpf/src/main.rs b/{{project-name}}-ebpf/src/main.rs
index 69185ce..c373228 100644
--- a/{{project-name}}-ebpf/src/main.rs
+++ b/{{project-name}}-ebpf/src/main.rs
@@ -177,9 +177,26 @@ pub fn {{crate_name}}(ctx: TracePointContext) -> u32 {
 unsafe fn try_{{crate_name}}(_ctx: TracePointContext) -> Result<u32, u32> {
     Ok(0)
 }
+{%- when "lsm" %}
+use aya_bpf::{
+    macros::lsm,
+    programs::LsmContext,
+};
+
+#[lsm(name="{{lsm_hook}}")]
+pub fn {{lsm_hook}}(ctx: LsmContext) -> i32 {
+    match unsafe { try_{{lsm_hook}}(ctx) } {
+        Ok(ret) => ret,
+        Err(ret) => ret,
+    }
+}
+
+unsafe fn try_{{lsm_hook}}(_ctx: LsmContext) -> Result<i32, i32> {
+    Ok(0)
+}
 {%- endcase %}
 
 #[panic_handler]
 fn panic(_info: &core::panic::PanicInfo) -> ! {
     unsafe { core::hint::unreachable_unchecked() }
-}
\ No newline at end of file
+}
diff --git a/{{project-name}}/src/main.rs b/{{project-name}}/src/main.rs
index cc29a49..47a929c 100644
--- a/{{project-name}}/src/main.rs
+++ b/{{project-name}}/src/main.rs
@@ -18,6 +18,8 @@ use aya::programs::{tc, SchedClassifier, TcAttachType};
 use aya::programs::{CgroupSkb, CgroupSkbAttachType};
 {%- when "tracepoint" -%}
 use aya::programs::TracePoint;
+{%- when "lsm" -%}
+use aya::programs::Lsm;
 {%- endcase %}
 use std::{
     convert::{TryFrom,TryInto},
@@ -91,8 +93,12 @@ fn try_main() -> Result<(), anyhow::Error> {
     let program: &mut TracePoint = bpf.program_mut("{{crate_name}}")?.try_into()?;
     program.load()?;
     program.attach("{{tracepoint_category}}", "{{tracepoint_name}}")?;
+    {%- when "lsm" -%}
+    let program: &mut Lsm = bpf.program_mut("{{lsm_hook}}")?.try_into()?;
+    program.load("{{lsm_hook}}")?;
+    program.attach()?;
     {%- endcase %}
-    
+
     let running = Arc::new(AtomicBool::new(true));
     let r = running.clone();
 
@@ -107,4 +113,4 @@ fn try_main() -> Result<(), anyhow::Error> {
     println!("Exiting...");
 
     Ok(())
-}
\ No newline at end of file
+}