1// SPDX-License-Identifier: GPL-2.0-only
2
3#include <linux/pagemap.h>
4#include <linux/xarray.h>
5#include <linux/slab.h>
6#include <linux/swap.h>
7#include <linux/swapops.h>
8#include <asm/mte.h>
9
10static DEFINE_XARRAY(mte_pages);
11
12void *mte_allocate_tag_storage(void)
13{
14	/* tags granule is 16 bytes, 2 tags stored per byte */
15	return kmalloc(MTE_PAGE_TAG_STORAGE, GFP_KERNEL);
16}
17
18void mte_free_tag_storage(char *storage)
19{
20	kfree(storage);
21}
22
23int mte_save_tags(struct page *page)
24{
25	void *tag_storage, *ret;
26
27	if (!page_mte_tagged(page))
28		return 0;
29
30	tag_storage = mte_allocate_tag_storage();
31	if (!tag_storage)
32		return -ENOMEM;
33
34	mte_save_page_tags(page_address(page), tag_storage);
35
36	/* lookup the swap entry.val from the page */
37	ret = xa_store(&mte_pages, page_swap_entry(page).val, tag_storage,
38		       GFP_KERNEL);
39	if (WARN(xa_is_err(ret), "Failed to store MTE tags")) {
40		mte_free_tag_storage(tag_storage);
41		return xa_err(ret);
42	} else if (ret) {
43		/* Entry is being replaced, free the old entry */
44		mte_free_tag_storage(ret);
45	}
46
47	return 0;
48}
49
50void mte_restore_tags(swp_entry_t entry, struct page *page)
51{
52	void *tags = xa_load(&mte_pages, entry.val);
53
54	if (!tags)
55		return;
56
57	if (try_page_mte_tagging(page)) {
58		mte_restore_page_tags(page_address(page), tags);
59		set_page_mte_tagged(page);
60	}
61}
62
63void mte_invalidate_tags(int type, pgoff_t offset)
64{
65	swp_entry_t entry = swp_entry(type, offset);
66	void *tags = xa_erase(&mte_pages, entry.val);
67
68	mte_free_tag_storage(tags);
69}
70
71void mte_invalidate_tags_area(int type)
72{
73	swp_entry_t entry = swp_entry(type, 0);
74	swp_entry_t last_entry = swp_entry(type + 1, 0);
75	void *tags;
76
77	XA_STATE(xa_state, &mte_pages, entry.val);
78
79	xa_lock(&mte_pages);
80	xas_for_each(&xa_state, tags, last_entry.val - 1) {
81		__xa_erase(&mte_pages, xa_state.xa_index);
82		mte_free_tag_storage(tags);
83	}
84	xa_unlock(&mte_pages);
85}
86