[RFC PATCH 3/4] bpf: Charge bpf jit limit in bpf_jit_alloc_exec

[Rick Edgecombe <rick.p.edgecombe@intel.com>]

This moves the calls to bpf_jit_charge_modmem and
bpf_jit_uncharge_modmem to bpf_jit_alloc_exec so the behavior can be
overridden for arch specific capabilities, for example if the arch
supports for allocating outside the module space.

Cc: Daniel Borkmann <daniel@iogearbox.net>
Cc: Alexei Starovoitov <ast@fb.com>
Signed-off-by: Rick Edgecombe <rick.p.edgecombe@intel.com>
---
 include/linux/filter.h |  3 +++
 kernel/bpf/core.c      | 20 +++++++++++---------
 2 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/include/linux/filter.h b/include/linux/filter.h
index ad106d845b22..33c0ae5990e1 100644
--- a/include/linux/filter.h
+++ b/include/linux/filter.h
@@ -859,6 +859,9 @@ bpf_jit_binary_alloc(unsigned int proglen, u8 **image_ptr,
 		     unsigned int alignment,
 		     bpf_jit_fill_hole_t bpf_fill_ill_insns);
 void bpf_jit_binary_free(struct bpf_binary_header *hdr);
+int bpf_jit_charge_modmem(u32 pages);
+void bpf_jit_uncharge_modmem(u32 pages);
+
 
 void bpf_jit_free(struct bpf_prog *fp);
 
diff --git a/kernel/bpf/core.c b/kernel/bpf/core.c
index f908b9356025..ce08a09839c2 100644
--- a/kernel/bpf/core.c
+++ b/kernel/bpf/core.c
@@ -722,7 +722,7 @@ static int __init bpf_jit_charge_init(void)
 }
 pure_initcall(bpf_jit_charge_init);
 
-static int bpf_jit_charge_modmem(u32 pages)
+int bpf_jit_charge_modmem(u32 pages)
 {
 	if (atomic_long_add_return(pages, &bpf_jit_current) >
 	    (bpf_jit_limit >> PAGE_SHIFT)) {
@@ -735,14 +735,22 @@ static int bpf_jit_charge_modmem(u32 pages)
 	return 0;
 }
 
-static void bpf_jit_uncharge_modmem(u32 pages)
+void bpf_jit_uncharge_modmem(u32 pages)
 {
 	atomic_long_sub(pages, &bpf_jit_current);
 }
 
 void *__weak bpf_jit_alloc_exec(unsigned long size)
 {
-	return module_alloc(size);
+	void *ret;
+	u32 pages = size / PAGE_SIZE;
+
+	if (bpf_jit_charge_modmem(pages))
+		return NULL;
+	ret = module_alloc(size);
+	if (!ret)
+		bpf_jit_uncharge_modmem(pages);
+	return ret;
 }
 
 void __weak bpf_jit_free_exec(void *addr)
@@ -765,13 +773,7 @@ bpf_jit_binary_alloc(unsigned int proglen, u8 **image_ptr,
 	size = round_up(proglen + sizeof(*hdr) + 128, PAGE_SIZE);
 	pages = size / PAGE_SIZE;
 
-	if (bpf_jit_charge_modmem(pages))
-		return NULL;
 	hdr = bpf_jit_alloc_exec(size);
-	if (!hdr) {
-		bpf_jit_uncharge_modmem(pages);
-		return NULL;
-	}
 
 	/* Fill space with illegal/arch-dep instructions. */
 	bpf_fill_ill_insns(hdr, size);
-- 
2.17.1