1 files changed, 5 insertions, 9 deletions

diff --git a/kernel/bpf/hashtab.c b/kernel/bpf/hashtab.c
index 0f2708fde5f7..d92e05d9979b 100644
--- a/kernel/bpf/hashtab.c
+++ b/kernel/bpf/hashtab.c
@@ -360,14 +360,8 @@ static struct bpf_map *htab_map_alloc(union bpf_attr *attr)
         else
                cost += (u64) htab->elem_size * num_possible_cpus();
 
-        if (cost >= U32_MAX - PAGE_SIZE)
-                /* make sure page count doesn't overflow */
-                goto free_htab;
-
-        htab->map.pages = round_up(cost, PAGE_SIZE) >> PAGE_SHIFT;
-
-        /* if map size is larger than memlock limit, reject it early */
-        err = bpf_map_precharge_memlock(htab->map.pages);
+        /* if map size is larger than memlock limit, reject it */
+        err = bpf_map_charge_init(&htab->map.memory, cost);
         if (err)
                 goto free_htab;
 
@@ -376,7 +370,7 @@ static struct bpf_map *htab_map_alloc(union bpf_attr *attr)
                                            sizeof(struct bucket),
                                            htab->map.numa_node);
         if (!htab->buckets)
-                goto free_htab;
+                goto free_charge;
 
         if (htab->map.map_flags & BPF_F_ZERO_SEED)
                 htab->hashrnd = 0;
@@ -409,6 +403,8 @@ free_prealloc:
         prealloc_destroy(htab);
 free_buckets:
         bpf_map_area_free(htab->buckets);
+free_charge:
+        bpf_map_charge_finish(&htab->map.memory);
 free_htab:
         kfree(htab);
         return ERR_PTR(err);