[PATCH] mm/secretmem: properly account locked pages

@@ -25,7 +25,8 @@ struct user_struct {
 
 #if defined(CONFIG_PERF_EVENTS) || defined(CONFIG_BPF_SYSCALL) || \
 	defined(CONFIG_NET) || defined(CONFIG_IO_URING) || \
-	defined(CONFIG_VFIO_PCI_ZDEV_KVM) || IS_ENABLED(CONFIG_IOMMUFD)
+	defined(CONFIG_VFIO_PCI_ZDEV_KVM) || IS_ENABLED(CONFIG_IOMMUFD) || \
+	defined(CONFIG_SECRETMEM)
 	atomic_long_t locked_vm;
 #endif
 #ifdef CONFIG_WATCH_QUEUE

@@ -18,6 +18,8 @@
 #include <linux/secretmem.h>
 #include <linux/set_memory.h>
 #include <linux/sched/signal.h>
+#include <linux/sched/user.h>
+#include <linux/cred.h>
 
 #include <uapi/linux/magic.h>
 

@@ -47,10 +49,107 @@ bool secretmem_active(void)
 	return !!atomic_read(&secretmem_users);
 }
 
+struct secretmem_inode_state {
+	struct user_struct	*user;
+	atomic_long_t		nr_pages_accounted;
+};
+
+static bool __secretmem_account_pages(struct user_struct *user,
+		unsigned long nr_pages)
+{
+	unsigned long page_limit, cur_pages, new_pages;
+
+	if (!nr_pages)
+		return true;
+
+	page_limit = rlimit(RLIMIT_MEMLOCK) >> PAGE_SHIFT;
+
+	cur_pages = atomic_long_read(&user->locked_vm);
+	do {
+		new_pages = cur_pages + nr_pages;
+		if (new_pages > page_limit)
+			return false;
+	} while (!atomic_long_try_cmpxchg(&user->locked_vm,
+					  &cur_pages, new_pages));
+	return true;
+}
+
+static bool secretmem_account_folio(struct secretmem_inode_state *state,
+		const struct folio *folio)
+{
+	unsigned long nr_pages;
+
+	if (!state)
+		return true;
+
+	nr_pages = folio_nr_pages(folio);
+	if (!__secretmem_account_pages(state->user, nr_pages))
+		return false;
+
+	atomic_long_add(nr_pages, &state->nr_pages_accounted);
+	return true;
+}
+
+static void __secretmem_unaccount_pages(struct secretmem_inode_state *state,
+		unsigned long nr_pages)
+{
+	atomic_long_sub(nr_pages, &state->user->locked_vm);
+	atomic_long_sub(nr_pages, &state->nr_pages_accounted);
+}
+
+static void secretmem_unaccount_folio(struct secretmem_inode_state *state,
+		struct folio *folio)
+{
+	if (!state)
+		return;
+
+	__secretmem_unaccount_pages(state, folio_nr_pages(folio));
+}
+
+static void secretmem_unaccount_all_folios(struct secretmem_inode_state *state)
+{
+	unsigned long nr_pages_accounted;
+
+	if (!state)
+		return;
+
+	nr_pages_accounted = atomic_long_read(&state->nr_pages_accounted);
+	__secretmem_unaccount_pages(state, nr_pages_accounted);
+}
+
+static void secretmem_destroy_inode_priv(struct inode *inode)
+{
+	struct secretmem_inode_state *state = inode->i_private;
+
+	if (!state)
+		return;
+
+	secretmem_unaccount_all_folios(state);
+	free_uid(state->user);
+	kfree(state);
+}
+
+static int secretmem_init_inode_priv(struct inode *inode)
+{
+	struct secretmem_inode_state *state;
+
+	if (ns_capable_noaudit(&init_user_ns, CAP_IPC_LOCK))
+		return 0;
+
+	state = kzalloc_obj(*state, GFP_KERNEL);
+	if (!state)
+		return -ENOMEM;
+
+	state->user = get_uid(current_user());
+	inode->i_private = state;
+	return 0;
+}
+
 static vm_fault_t secretmem_fault(struct vm_fault *vmf)
 {
 	struct address_space *mapping = vmf->vma->vm_file->f_mapping;
 	struct inode *inode = file_inode(vmf->vma->vm_file);
+	struct secretmem_inode_state *state = inode->i_private;
 	pgoff_t offset = vmf->pgoff;
 	gfp_t gfp = vmf->gfp_mask;
 	unsigned long addr;

@@ -72,8 +171,15 @@ static vm_fault_t secretmem_fault(struct vm_fault *vmf)
 			goto out;
 		}
 
+		if (!secretmem_account_folio(state, folio)) {
+			folio_put(folio);
+			ret = VM_FAULT_SIGBUS;
+			goto out;
+		}
+
 		err = set_direct_map_invalid_noflush(folio_page(folio, 0));
 		if (err) {
+			secretmem_unaccount_folio(state, folio);
 			folio_put(folio);
 			ret = vmf_error(err);
 			goto out;

@@ -82,6 +188,7 @@ static vm_fault_t secretmem_fault(struct vm_fault *vmf)
 		__folio_mark_uptodate(folio);
 		err = filemap_add_folio(mapping, folio, offset, gfp);
 		if (unlikely(err)) {
+			secretmem_unaccount_folio(state, folio);
 			/*
 			 * If a split of large page was required, it
 			 * already happened when we marked the page invalid

@@ -115,19 +222,17 @@ static const struct vm_operations_struct secretmem_vm_ops = {
 static int secretmem_release(struct inode *inode, struct file *file)
 {
 	atomic_dec(&secretmem_users);
+	secretmem_destroy_inode_priv(inode);
+
 	return 0;
 }
 
 static int secretmem_mmap_prepare(struct vm_area_desc *desc)
 {
-	const unsigned long len = vma_desc_size(desc);
-
 	if (!vma_desc_test_any(desc, VMA_SHARED_BIT, VMA_MAYSHARE_BIT))
 		return -EINVAL;
 
-	vma_desc_set_flags(desc, VMA_LOCKED_BIT, VMA_DONTDUMP_BIT);
-	if (!mlock_future_ok(desc->mm, /*is_vma_locked=*/ true, len))
-		return -EAGAIN;
+	vma_desc_set_flags(desc, VMA_DONTDUMP_BIT);
 	desc->vm_ops = &secretmem_vm_ops;
 
 	return 0;

@@ -192,15 +297,23 @@ static struct file *secretmem_file_create(unsigned long flags)
 	struct file *file;
 	struct inode *inode;
 	const char *anon_name = "[secretmem]";
+	int err;
 
 	inode = anon_inode_make_secure_inode(secretmem_mnt->mnt_sb, anon_name, NULL);
 	if (IS_ERR(inode))
 		return ERR_CAST(inode);
 
+	err = secretmem_init_inode_priv(inode);
+	if (err)
+		goto err_free_inode;
+
 	file = alloc_file_pseudo(inode, secretmem_mnt, "secretmem",
 				 O_RDWR | O_LARGEFILE, &secretmem_fops);
-	if (IS_ERR(file))
+	if (IS_ERR(file)) {
+		secretmem_destroy_inode_priv(inode);
+		err = PTR_ERR(file);
 		goto err_free_inode;
+	}
 
 	mapping_set_gfp_mask(inode->i_mapping, GFP_USER);
 	mapping_set_unevictable(inode->i_mapping);

@@ -218,7 +331,7 @@ static struct file *secretmem_file_create(unsigned long flags)
 
 err_free_inode:
 	iput(inode);
-	return file;
+	return ERR_PTR(err);
 }
 
 SYSCALL_DEFINE1(memfd_secret, unsigned int, flags)

@@ -15,6 +15,8 @@
 #include <sys/resource.h>
 #include <sys/capability.h>
 
+#include <setjmp.h>
+#include <signal.h>
 #include <stdlib.h>
 #include <string.h>
 #include <unistd.h>

@@ -22,6 +24,8 @@
 #include <stdio.h>
 #include <fcntl.h>
 
+#include <sys/mman.h>
+
 #include "kselftest.h"
 
 #define fail(fmt, ...) ksft_test_result_fail(fmt, ##__VA_ARGS__)

@@ -31,6 +35,7 @@
 #ifdef __NR_memfd_secret
 
 #define PATTERN	0x55
+#define MLOCK_LIMIT_CAP	(8UL << 20)
 
 static const int prot = PROT_READ | PROT_WRITE;
 static const int mode = MAP_SHARED;

@@ -39,6 +44,13 @@ static unsigned long page_size;
 static unsigned long mlock_limit_cur;
 static unsigned long mlock_limit_max;
 
+static sigjmp_buf fault_env;
+
+static void sigbus_handler(int sig)
+{
+	siglongjmp(fault_env, 1);
+}
+
 static int memfd_secret(unsigned int flags)
 {
 	return syscall(__NR_memfd_secret, flags);

@@ -57,10 +69,32 @@ static void test_file_apis(int fd)
 		pass("file IO is blocked as expected\n");
 }
 
-static void test_mlock_limit(int fd)
+/* GUP disallows automatic fault-in of secretmem, so do it manually. */
+static bool fault_in_secretmem(char *mem, size_t len)
+{
+	if (sigsetjmp(fault_env, 1))
+		return false;
+	memset(mem, PATTERN, len);
+	return true;
+}
+
+static void test_mlock_limit(void)
 {
 	size_t len;
 	char *mem;
+	int fd;
+
+	/* Locked pages have an inode lifetime, so need a new fd. */
+	fd = memfd_secret(0);
+	if (fd < 0) {
+		fail("memfd_secret failed: %s\n", strerror(errno));
+		return;
+	}
+
+	if (ftruncate(fd, mlock_limit_max * 2)) {
+		fail("ftruncate failed: %s\n", strerror(errno));
+		goto out_close;
+	}
 
 	len = mlock_limit_cur;
 	if (len % page_size != 0)

@@ -69,19 +103,48 @@ static void test_mlock_limit(int fd)
 	mem = mmap(NULL, len, prot, mode, fd, 0);
 	if (mem == MAP_FAILED) {
 		fail("unable to mmap secret memory\n");
-		return;
+		goto out_close;
+	}
+
+	if (!fault_in_secretmem(mem, len)) {
+		munmap(mem, len);
+		fail("unable to fault in secret memory\n");
+		goto out_close;
 	}
 	munmap(mem, len);
 
 	len = mlock_limit_max * 2;
 	mem = mmap(NULL, len, prot, mode, fd, 0);
-	if (mem != MAP_FAILED) {
-		fail("unexpected mlock limit violation\n");
+	if (mem == MAP_FAILED) {
+		fail("unable to mmap secret memory\n");
+		goto out_close;
+	}
+
+	if (fault_in_secretmem(mem, len)) {
 		munmap(mem, len);
-		return;
+		fail("mlock limit is not respected\n");
+		goto out_close;
 	}
+	munmap(mem, len);
+
+	/* map a page past the limit to assert inode scope. */
+	mem = mmap(NULL, page_size, prot, mode, fd,
+		   mlock_limit_max & ~(page_size - 1));
+	if (mem == MAP_FAILED) {
+		fail("unable to mmap secret memory\n");
+		goto out_close;
+	}
+
+	if (fault_in_secretmem(mem, page_size)) {
+		munmap(mem, page_size);
+		fail("mlock limit is not respected\n");
+		goto out_close;
+	}
+	munmap(mem, page_size);
 
 	pass("mlock limit is respected\n");
+out_close:
+	close(fd);
 }
 
 static void test_vmsplice(int fd, const char *desc)

@@ -292,6 +355,12 @@ static void prepare(void)
 	if (page_size > mlock_limit_max)
 		mlock_limit_max = page_size;
 
+	/* Clamp huge or unlimited. */
+	if (mlock_limit_max > MLOCK_LIMIT_CAP)
+		mlock_limit_max = MLOCK_LIMIT_CAP;
+	if (mlock_limit_cur > mlock_limit_max)
+		mlock_limit_cur = mlock_limit_max;
+
 	if (set_cap_limits(mlock_limit_max))
 		ksft_exit_fail_msg("Unable to set mlock limit: %s\n",
 				   strerror(errno));

@@ -301,6 +370,7 @@ static void prepare(void)
 
 int main(int argc, char *argv[])
 {
+	struct sigaction sa = { .sa_handler = sigbus_handler };
 	int fd;
 
 	prepare();

@@ -316,10 +386,15 @@ int main(int argc, char *argv[])
 			ksft_exit_fail_msg("memfd_secret failed: %s\n",
 					   strerror(errno));
 	}
+
+	sigemptyset(&sa.sa_mask);
+	if (sigaction(SIGBUS, &sa, NULL))
+		ksft_exit_fail_msg("Cannot set up SIGBUS handler");
+
 	if (ftruncate(fd, page_size))
 		ksft_exit_fail_msg("ftruncate failed: %s\n", strerror(errno));
 
-	test_mlock_limit(fd);
+	test_mlock_limit();
 	test_file_apis(fd);
 	/*
 	 * We have to run the first vmsplice test before any secretmem page was