1 files changed, 2 insertions, 7 deletions

diff --git a/net/core/sock_map.c b/net/core/sock_map.c
index be6092ac69f8..52d4faeee18b 100644
--- a/net/core/sock_map.c
+++ b/net/core/sock_map.c
@@ -44,13 +44,7 @@ static struct bpf_map *sock_map_alloc(union bpf_attr *attr)
 
         /* Make sure page count doesn't overflow. */
         cost = (u64) stab->map.max_entries * sizeof(struct sock *);
-        if (cost >= U32_MAX - PAGE_SIZE) {
-                err = -EINVAL;
-                goto free_stab;
-        }
-
-        stab->map.pages = round_up(cost, PAGE_SIZE) >> PAGE_SHIFT;
-        err = bpf_map_precharge_memlock(stab->map.pages);
+        err = bpf_map_charge_init(&stab->map.memory, cost);
         if (err)
                 goto free_stab;
 
@@ -60,6 +54,7 @@ static struct bpf_map *sock_map_alloc(union bpf_attr *attr)
         if (stab->sks)
                 return &stab->map;
         err = -ENOMEM;
+        bpf_map_charge_finish(&stab->map.memory);
 free_stab:
         kfree(stab);
         return ERR_PTR(err);