From 555e7507d8034b50bed9ca0d3ce861257068ce6c Mon Sep 17 00:00:00 2001
From: Stephen Bates <sbates@raithlin.com>
Date: Thu, 22 Jul 2021 12:47:13 -0600
Subject: [PATCH] ubpf_run: Refactor code to assist in adding instruction
 counting.

Refactor code to add a ubpf_run function so it is easier to add
instruction counting for both the VM and JIT cases. The counting
itself will be added in subsequent patches.
---
 vm/inc/ubpf.h |  7 +++++++
 vm/test.c     | 13 +------------
 vm/ubpf_vm.c  | 21 +++++++++++++++++++++
 3 files changed, 29 insertions(+), 12 deletions(-)

diff --git a/vm/inc/ubpf.h b/vm/inc/ubpf.h
index bf57277bd..6bbeb524b 100644
--- a/vm/inc/ubpf.h
+++ b/vm/inc/ubpf.h
@@ -102,6 +102,13 @@ int ubpf_exec(const struct ubpf_vm *vm, void *mem, size_t mem_len, uint64_t* bpf
 
 ubpf_jit_fn ubpf_compile(struct ubpf_vm *vm, char **errmsg);
 
+/*
+ * A wrapper around some of the compile, load and run commands to
+ * enable instruction counting. Returns the return value of the eBPF
+ * program being executed.
+ */
+uint64_t ubpf_run(struct ubpf_vm *vm, void *mem, size_t mem_len, bool jit);
+
 /*
  * Translate the eBPF byte code to x64 machine code, store in buffer, and 
  * write the resulting count of bytes to size.
diff --git a/vm/test.c b/vm/test.c
index 74e99191b..4437dbf7e 100644
--- a/vm/test.c
+++ b/vm/test.c
@@ -130,18 +130,7 @@ int main(int argc, char **argv)
 
     uint64_t ret;
 
-    if (jit) {
-        ubpf_jit_fn fn = ubpf_compile(vm, &errmsg);
-        if (fn == NULL) {
-            fprintf(stderr, "Failed to compile: %s\n", errmsg);
-            free(errmsg);
-            return 1;
-        }
-        ret = fn(mem, mem_len);
-    } else {
-        if (ubpf_exec(vm, mem, mem_len, &ret) < 0)
-            ret = UINT64_MAX;
-    }
+    ret = ubpf_run(vm, mem, mem_len, jit);
 
     printf("0x%"PRIx64"\n", ret);
 
diff --git a/vm/ubpf_vm.c b/vm/ubpf_vm.c
index b860ba857..53064439a 100644
--- a/vm/ubpf_vm.c
+++ b/vm/ubpf_vm.c
@@ -573,6 +573,27 @@ ubpf_exec(const struct ubpf_vm *vm, void *mem, size_t mem_len, uint64_t* bpf_ret
     }
 }
 
+uint64_t
+ubpf_run(struct ubpf_vm *vm, void *mem, size_t mem_len, bool jit)
+{
+    uint64_t ret;
+    char *errmsg;
+
+    if (jit) {
+	ubpf_jit_fn fn = ubpf_compile(vm, &errmsg);
+	if (fn == NULL) {
+	    fprintf(stderr, "Failed to compile: %s\n", errmsg);
+	    free(errmsg);
+	    return 1;
+	}
+	ret = fn(mem, mem_len);
+    } else {
+	if (ubpf_exec(vm, mem, mem_len, &ret) < 0)
+	    ret = UINT64_MAX;
+    }
+    return ret;
+}
+
 static bool
 validate(const struct ubpf_vm *vm, const struct ebpf_inst *insts, uint32_t num_insts, char **errmsg)
 {