refactor code

Signed-off-by: Aryan-sharma11 <aryan1126.sharma@gmail.com>

KubeArmor/BPF/enforcer.bpf.c

@@ -11,19 +11,11 @@ int BPF_PROG(enforce_proc, struct linux_binprm *bprm, int ret) {
   struct task_struct *t = (struct task_struct *)bpf_get_current_task();
   event *task_info;
   int retval = ret;
-  // variables required for argument matching---------| 
-  struct argVal *argval ;
-  unsigned int num = BPF_CORE_READ(bprm , argc);
-  unsigned int argKey; 
-  unsigned int *x ;
-  u32 arg_k = 0;
-  arg_bufs_k *a_key = bpf_map_lookup_elem(&args_bufk, &arg_k);
-   if (a_key == NULL)
-    return 0;
-  bool argmatch = false;
-  bool frmsource = false;
-  // -------------------------------------------------|
 
+  // no of arguments
+  unsigned int num_of_args = BPF_CORE_READ(bprm , argc);
+  bool argmatch = false;
+
   bool match = false; 
   struct outer_key okey;
   get_outer_key(&okey, t);
@@ -231,43 +223,10 @@ int BPF_PROG(enforce_proc, struct linux_binprm *bprm, int ret) {
 
 decision:  
   if (match) {
-    if (val && (val->processmask & RULE_ARGSET)) {
-      argKey =  bpf_get_current_pid_tgid();
-      argval = bpf_map_lookup_elem(&args_store, &argKey);
-
-     // clearing to avoid processing garbage values 
-      __builtin_memset(&a_key->okey, 0, sizeof(a_key->okey));
-      __builtin_memset(&a_key->store, 0, sizeof(a_key->store));
-
-      bpf_probe_read(&a_key->okey.mnt_ns, sizeof(okey.mnt_ns) , &okey.mnt_ns);
-      bpf_probe_read(&a_key->okey.pid_ns, sizeof(okey.pid_ns) , &okey.pid_ns);
-      bpf_probe_read_str(&a_key->store.path, sizeof(store->path) , store->path);
-
-      if (pk->path[0] == '\0') {
-          bpf_probe_read_str(&a_key->store.source, sizeof(store->source) , store->source);
-      } 
-      if (argval) {
-        for( int i = 0 ; i< num && i < 100; i++ ){
-            __builtin_memset(a_key->arg, 0, sizeof(a_key->arg));
-            bpf_probe_read_str(&a_key->arg, sizeof(a_key->arg), argval->argsArray[i]);
-            x  = bpf_map_lookup_elem(&a_map ,a_key);
-            bpf_printk("a_key->path %s , a_key->source - %s ", a_key->store.path , a_key->store.source);
-            if (x) {
-                bpf_printk("argument matched");
-                argmatch = true;
-                if (i != 0) {  
-                    continue;  
-                }
-            } else {
-                if (i != 0) {
-                    argmatch = false;
-                    break;  
-                }
-            }
-        }
-        }
+    if (val && (val->processmask & RULE_ARGSET)){ 
+      argmatch = matchArguments( num_of_args , &okey , store , pk);
     }
-
+      
     if (val && (val->processmask & RULE_OWNER)) {
       if (!is_owner(bprm->file)) {
         if((val->processmask & RULE_ARGSET) && argmatch){

KubeArmor/BPF/shared.h

@@ -11,6 +11,7 @@
 #include <bpf/bpf_helpers.h>
 #include <bpf/bpf_tracing.h>
 #include "throttling.h"
+#include "common.h"
 
 char LICENSE[] SEC("license") = "Dual BSD/GPL";
 #define EPERM 13
@@ -744,6 +745,53 @@ static inline int match_and_enforce_path_hooks(struct path *f_path, u32 id,
   bpf_ringbuf_submit(task_info, 0);
   return retval;
 }
+static inline bool matchArguments( unsigned int num_of_args , struct outer_key *okey , bufs_k *store , bufs_k *pk ) {
+
+    bool argmatch = false;
+    unsigned int *x ;
+
+    unsigned int argKey;
+    struct argVal *argval ;
+    argKey =  bpf_get_current_pid_tgid();
+    argval = bpf_map_lookup_elem(&args_store, &argKey);
+
+    u32 arg_k = 0;
+    arg_bufs_k *a_key = bpf_map_lookup_elem(&args_bufk, &arg_k);
+    if (a_key == NULL)
+      return 0;
+
+    // clearing to avoid processing garbage values 
+    __builtin_memset(&a_key->okey, 0, sizeof(a_key->okey));
+    __builtin_memset(&a_key->store, 0, sizeof(a_key->store));
+
+    bpf_probe_read(&a_key->okey.mnt_ns, sizeof(okey->mnt_ns) , &okey->mnt_ns);
+    bpf_probe_read(&a_key->okey.pid_ns, sizeof(okey->pid_ns) , &okey->pid_ns);
+    bpf_probe_read_str(&a_key->store.path, sizeof(store->path) , &store->path);
+
+    if (pk->path[0] == '\0') {
+      // pk->path[0] will be null for fromSource rules
+        bpf_probe_read_str(&a_key->store.source, sizeof(store->source) , store->source);
+    }
+    if (argval) {
+      for( int i = 0 ; i< num_of_args && i < 100; i++ ){
+          __builtin_memset(a_key->arg, 0, sizeof(a_key->arg));
+          bpf_probe_read_str(&a_key->arg, sizeof(a_key->arg), argval->argsArray[i]);
+          x = bpf_map_lookup_elem(&a_map ,a_key);
+          if (x) {
+              argmatch = true;
+              if (i != 0) {  
+                    continue;  
+              }
+            } else {
+                if (i != 0) {
+                  argmatch = false;
+                  break;  
+                }
+            }
+      }
+    }
+  return argmatch; 
+}
 
 /*
   How do we check what to deny or not?