From 089f6279752adfb01386600d119913403ed326ee Mon Sep 17 00:00:00 2001
From: Oleg Pigal <olegpi@google.com>
Date: Tue, 21 Apr 2020 19:46:52 -0400
Subject: [PATCH] Introduce toggle to enable/disable bounds_check.

bounds_check routine assumes that memory to read / write may only be on the stack or within ubpf code itself.
The toggle allows disabling the check when the assumption is not correct.
---
 vm/inc/ubpf.h |  9 +++++++++
 vm/ubpf_int.h |  1 +
 vm/ubpf_vm.c  | 18 ++++++++++++++----
 3 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/vm/inc/ubpf.h b/vm/inc/ubpf.h
index 6d23dbf7..c72194d5 100644
--- a/vm/inc/ubpf.h
+++ b/vm/inc/ubpf.h
@@ -26,6 +26,15 @@ typedef uint64_t (*ubpf_jit_fn)(void *mem, size_t mem_len);
 struct ubpf_vm *ubpf_create(void);
 void ubpf_destroy(struct ubpf_vm *vm);
 
+/*
+ * Enable / disable bounds_check
+ *
+ * Bounds check is enabled by default, but it may be too restrictive
+ * Pass true to enable, false to disable
+ * Returns previous state
+ */
+bool toggle_bounds_check(struct ubpf_vm *vm, bool enable);
+
 /*
  * Register an external function
  *
diff --git a/vm/ubpf_int.h b/vm/ubpf_int.h
index 62727f3a..680c699a 100644
--- a/vm/ubpf_int.h
+++ b/vm/ubpf_int.h
@@ -33,6 +33,7 @@ struct ubpf_vm {
     size_t jitted_size;
     ext_func *ext_funcs;
     const char **ext_func_names;
+    bool bounds_check_enabled;
 };
 
 char *ubpf_error(const char *fmt, ...);
diff --git a/vm/ubpf_vm.c b/vm/ubpf_vm.c
index 7116c9e7..9e35a85d 100644
--- a/vm/ubpf_vm.c
+++ b/vm/ubpf_vm.c
@@ -27,7 +27,14 @@
 #define MAX_EXT_FUNCS 64
 
 static bool validate(const struct ubpf_vm *vm, const struct ebpf_inst *insts, uint32_t num_insts, char **errmsg);
-static bool bounds_check(void *addr, int size, const char *type, uint16_t cur_pc, void *mem, size_t mem_len, void *stack);
+static bool bounds_check(const struct ubpf_vm *vm, void *addr, int size, const char *type, uint16_t cur_pc, void *mem, size_t mem_len, void *stack);
+
+bool toggle_bounds_check(struct ubpf_vm *vm, bool enable)
+{
+  bool old = vm->bounds_check_enabled;
+  vm->bounds_check_enabled = enable;
+  return old;
+}
 
 struct ubpf_vm *
 ubpf_create(void)
@@ -49,6 +56,7 @@ ubpf_create(void)
         return NULL;
     }
 
+    vm->bounds_check_enabled = true;
     return vm;
 }
 
@@ -366,13 +374,13 @@ ubpf_exec(const struct ubpf_vm *vm, void *mem, size_t mem_len)
          */
 #define BOUNDS_CHECK_LOAD(size) \
     do { \
-        if (!bounds_check((void *)reg[inst.src] + inst.offset, size, "load", cur_pc, mem, mem_len, stack)) { \
+        if (!bounds_check(vm, (void *)reg[inst.src] + inst.offset, size, "load", cur_pc, mem, mem_len, stack)) { \
             return UINT64_MAX; \
         } \
     } while (0)
 #define BOUNDS_CHECK_STORE(size) \
     do { \
-        if (!bounds_check((void *)reg[inst.dst] + inst.offset, size, "store", cur_pc, mem, mem_len, stack)) { \
+        if (!bounds_check(vm, (void *)reg[inst.dst] + inst.offset, size, "store", cur_pc, mem, mem_len, stack)) { \
             return UINT64_MAX; \
         } \
     } while (0)
@@ -732,8 +740,10 @@ validate(const struct ubpf_vm *vm, const struct ebpf_inst *insts, uint32_t num_i
 }
 
 static bool
-bounds_check(void *addr, int size, const char *type, uint16_t cur_pc, void *mem, size_t mem_len, void *stack)
+bounds_check(const struct ubpf_vm *vm, void *addr, int size, const char *type, uint16_t cur_pc, void *mem, size_t mem_len, void *stack)
 {
+    if (!vm->bounds_check_enabled)
+        return true;
     if (mem && (addr >= mem && (addr + size) <= (mem + mem_len))) {
         /* Context access */
         return true;