void *obj, u64 address, u64 length, u8 mmap_flag)
 {
        struct efa_mmap_entry *entry;
+       u32 next_mmap_page;
        int err;
 
        entry = kmalloc(sizeof(*entry), GFP_KERNEL);
        entry->mmap_flag = mmap_flag;
 
        xa_lock(&ucontext->mmap_xa);
+       if (check_add_overflow(ucontext->mmap_xa_page,
+                              (u32)(length >> PAGE_SHIFT),
+                              &next_mmap_page))
+               goto err_unlock;
+
        entry->mmap_page = ucontext->mmap_xa_page;
-       ucontext->mmap_xa_page += DIV_ROUND_UP(length, PAGE_SIZE);
+       ucontext->mmap_xa_page = next_mmap_page;
        err = __xa_insert(&ucontext->mmap_xa, entry->mmap_page, entry,
                          GFP_KERNEL);
+       if (err)
+               goto err_unlock;
+
        xa_unlock(&ucontext->mmap_xa);
-       if (err){
-               kfree(entry);
-               return EFA_MMAP_INVALID;
-       }
 
        ibdev_dbg(
                &dev->ibdev,
                entry->obj, entry->address, entry->length, get_mmap_key(entry));
 
        return get_mmap_key(entry);
+
+err_unlock:
+       xa_unlock(&ucontext->mmap_xa);
+       kfree(entry);
+       return EFA_MMAP_INVALID;
+
 }
 
 int efa_query_device(struct ib_device *ibdev,