diff --git a/fs/userfaultfd.c b/fs/userfaultfd.c
index b5b3cd07da206e70ed2fd2c924ba0bb4d362a353..a89281fc2b683b065fde8a6002f110b4cefb0cb9 100644
--- a/fs/userfaultfd.c
+++ b/fs/userfaultfd.c
@@ -843,7 +843,7 @@ static int userfaultfd_release(struct inode *inode, struct file *file)
 	struct vm_area_struct *vma, *prev;
 	/* len == 0 means wake all */
 	struct userfaultfd_wake_range range = { .len = 0, };
-	unsigned long new_flags;
+	unsigned long new_flags, userfault_flags;
 
 	WRITE_ONCE(ctx->released, true);
 
@@ -861,14 +861,18 @@ static int userfaultfd_release(struct inode *inode, struct file *file)
 	mmap_write_lock(mm);
 	prev = NULL;
 	for (vma = mm->mmap; vma; vma = vma->vm_next) {
+		userfault_flags = VM_UFFD_MISSING | VM_UFFD_WP;
+#ifdef CONFIG_USERSWAP
+		userfault_flags |= VM_USWAP;
+#endif
 		cond_resched();
 		BUG_ON(!!vma->vm_userfaultfd_ctx.ctx ^
-		       !!(vma->vm_flags & (VM_UFFD_MISSING | VM_UFFD_WP)));
+		       !!(vma->vm_flags & userfault_flags));
 		if (vma->vm_userfaultfd_ctx.ctx != ctx) {
 			prev = vma;
 			continue;
 		}
-		new_flags = vma->vm_flags & ~(VM_UFFD_MISSING | VM_UFFD_WP);
+		new_flags = vma->vm_flags & ~userfault_flags;
 		prev = vma_merge(mm, prev, vma->vm_start, vma->vm_end,
 				 new_flags, vma->anon_vma,
 				 vma->vm_file, vma->vm_pgoff,
@@ -1293,6 +1297,8 @@ static int userfaultfd_register(struct userfaultfd_ctx *ctx,
 	 */
 	if (uffdio_register.mode & UFFDIO_REGISTER_MODE_USWAP) {
 		uffdio_register.mode &= ~UFFDIO_REGISTER_MODE_USWAP;
+		if (!uffdio_register.mode)
+			goto out;
 		vm_flags |= VM_USWAP;
 		end = uffdio_register.range.start + uffdio_register.range.len - 1;
 		vma = find_vma(mm, uffdio_register.range.start);
diff --git a/mm/memory.c b/mm/memory.c
index 0be6fd3198a8052596329edd308f428abcf07dbe..d15b9ccfc108e52cde6a9a9a44b1b742d3a00431 100644
--- a/mm/memory.c
+++ b/mm/memory.c
@@ -3317,7 +3317,6 @@ vm_fault_t do_swap_page(struct vm_fault *vmf)
 					vmf->address, vma->vm_flags);
 			goto skip_uswap;
 		}
-		BUG_ON(!(vma->vm_flags & VM_UFFD_MISSING));
 		ret = handle_userfault(vmf, VM_UFFD_MISSING | VM_USWAP);
 		return ret;
 	}