/*
  * Handle a single data segment in a page-fault WQE or RDMA region.
  *
- * Returns number of OS pages retrieved on success. The caller may continue to
- * the next data segment.
+ * Returns zero on success. The caller may continue to the next data segment.
  * Can return the following error codes:
  * -EAGAIN to designate a temporary error. The caller will abort handling the
  *  page fault and resolve it.
                                         u32 *bytes_committed,
                                         u32 *bytes_mapped)
 {
-       int npages = 0, ret, i, outlen, cur_outlen = 0, depth = 0;
+       int ret, i, outlen, cur_outlen = 0, depth = 0, pages_in_range;
        struct pf_frame *head = NULL, *frame;
        struct mlx5_ib_mkey *mmkey;
        struct mlx5_ib_mr *mr;
        case MLX5_MKEY_MR:
                mr = container_of(mmkey, struct mlx5_ib_mr, mmkey);
 
+               pages_in_range = (ALIGN(io_virt + bcnt, PAGE_SIZE) -
+                                 (io_virt & PAGE_MASK)) >>
+                                PAGE_SHIFT;
                ret = pagefault_mr(mr, io_virt, bcnt, bytes_mapped, 0, false);
                if (ret < 0)
                        goto end;
 
                mlx5_update_odp_stats_with_handled(mr, faults, ret);
 
-               npages += ret;
+               if (ret < pages_in_range) {
+                       ret = -EFAULT;
+                       goto end;
+               }
+
                ret = 0;
                break;
 
        kfree(out);
 
        *bytes_committed = 0;
-       return ret ? ret : npages;
+       return ret;
 }
 
 /*
  *                   the committed bytes).
  * @receive_queue: receive WQE end of sg list
  *
- * Returns the number of pages loaded if positive, zero for an empty WQE, or a
- * negative error code.
+ * Returns zero for success or a negative error code.
  */
 static int pagefault_data_segments(struct mlx5_ib_dev *dev,
                                   struct mlx5_pagefault *pfault,
                                   void *wqe_end, u32 *bytes_mapped,
                                   u32 *total_wqe_bytes, bool receive_queue)
 {
-       int ret = 0, npages = 0;
+       int ret = 0;
        u64 io_virt;
        __be32 key;
        u32 byte_count;
                                                    bytes_mapped);
                if (ret < 0)
                        break;
-               npages += ret;
        }
 
-       return ret < 0 ? ret : npages;
+       return ret;
 }
 
 /*
        free_page((unsigned long)wqe_start);
 }
 
-static int pages_in_range(u64 address, u32 length)
-{
-       return (ALIGN(address + length, PAGE_SIZE) -
-               (address & PAGE_MASK)) >> PAGE_SHIFT;
-}
-
 static void mlx5_ib_mr_rdma_pfault_handler(struct mlx5_ib_dev *dev,
                                           struct mlx5_pagefault *pfault)
 {
        if (ret == -EAGAIN) {
                /* We're racing with an invalidation, don't prefetch */
                prefetch_activated = 0;
-       } else if (ret < 0 || pages_in_range(address, length) > ret) {
+       } else if (ret < 0) {
                mlx5_ib_page_fault_resume(dev, pfault, 1);
                if (ret != -ENOENT)
                        mlx5_ib_dbg(dev, "PAGE FAULT error %d. QP 0x%llx, type: 0x%x\n",