[RFC PATCH 13/47] asi: Added ASI memory cgroup flag

[Junaid Shahid <junaids@google.com>]

From: Ofir Weisse <oweisse@google.com>

Adds a cgroup flag to control if ASI is enabled for processes in
that cgroup.

Can be set or cleared by writing to the memory.use_asi file in the
memory cgroup. The flag only affects new processes created after
the flag was set.

In addition to the cgroup flag, we may also want to add a per-process
flag, though it will have to be something that can be set at process
creation time.

Signed-off-by: Ofir Weisse <oweisse@google.com>
Co-developed-by: Junaid Shahid <junaids@google.com>
Signed-off-by: Junaid Shahid <junaids@google.com>


---
 arch/x86/mm/asi.c          | 14 ++++++++++++++
 include/linux/memcontrol.h |  3 +++
 include/linux/mm_types.h   | 17 +++++++++++++++++
 mm/memcontrol.c            | 30 ++++++++++++++++++++++++++++++
 4 files changed, 64 insertions(+)

diff --git a/arch/x86/mm/asi.c b/arch/x86/mm/asi.c
index 71348399baf1..ca50a32ecd7e 100644
--- a/arch/x86/mm/asi.c
+++ b/arch/x86/mm/asi.c
@@ -2,6 +2,7 @@
 
 #include <linux/init.h>
 #include <linux/memblock.h>
+#include <linux/memcontrol.h>
 
 #include <asm/asi.h>
 #include <asm/pgalloc.h>
@@ -322,7 +323,20 @@ EXPORT_SYMBOL_GPL(asi_exit);
 
 void asi_init_mm_state(struct mm_struct *mm)
 {
+	struct mem_cgroup *memcg = get_mem_cgroup_from_mm(mm);
+
 	memset(mm->asi, 0, sizeof(mm->asi));
+	mm->asi_enabled = false;
+
+	/*
+	 * TODO: In addition to a cgroup flag, we may also want a per-process
+	 * flag.
+	 */
+        if (memcg) {
+		mm->asi_enabled = boot_cpu_has(X86_FEATURE_ASI) &&
+				  memcg->use_asi;
+		css_put(&memcg->css);
+	}
 }
 
 static bool is_page_within_range(size_t addr, size_t page_size,
diff --git a/include/linux/memcontrol.h b/include/linux/memcontrol.h
index 0c5c403f4be6..a883cb458b06 100644
--- a/include/linux/memcontrol.h
+++ b/include/linux/memcontrol.h
@@ -259,6 +259,9 @@ struct mem_cgroup {
 	 */
 	bool oom_group;
 
+#ifdef CONFIG_ADDRESS_SPACE_ISOLATION
+        bool use_asi;
+#endif
 	/* protected by memcg_oom_lock */
 	bool		oom_lock;
 	int		under_oom;
diff --git a/include/linux/mm_types.h b/include/linux/mm_types.h
index 5b8028fcfe67..8624d2783661 100644
--- a/include/linux/mm_types.h
+++ b/include/linux/mm_types.h
@@ -607,6 +607,14 @@ struct mm_struct {
 		 * new_owner->alloc_lock is held
 		 */
 		struct task_struct __rcu *owner;
+
+#endif
+#ifdef CONFIG_ADDRESS_SPACE_ISOLATION
+                /* Is ASI enabled for this mm? ASI requires allocating extra
+                 * resources, such as ASI page tables. To prevent allocationg
+                 * these resources for every mm in the system, we expect that
+                 * only VM mm's will have this flag set. */
+		bool asi_enabled;
 #endif
 		struct user_namespace *user_ns;
 
@@ -665,6 +673,15 @@ struct mm_struct {
 
 extern struct mm_struct init_mm;
 
+#ifdef CONFIG_ADDRESS_SPACE_ISOLATION
+static inline bool mm_asi_enabled(struct mm_struct *mm)
+{
+        return mm->asi_enabled;
+}
+#else
+static inline bool mm_asi_enabled(struct mm_struct *mm) { return false; }
+#endif
+
 /* Pointer magic because the dynamic array size confuses some compilers. */
 static inline void mm_init_cpumask(struct mm_struct *mm)
 {
diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index 2ed5f2a0879d..a66d6b222ecf 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -3539,6 +3539,29 @@ static int mem_cgroup_hierarchy_write(struct cgroup_subsys_state *css,
 	return -EINVAL;
 }
 
+#ifdef CONFIG_ADDRESS_SPACE_ISOLATION
+
+static u64 mem_cgroup_asi_read(struct cgroup_subsys_state *css,
+				     struct cftype *cft)
+{
+        return mem_cgroup_from_css(css)->use_asi;
+}
+
+static int mem_cgroup_asi_write(struct cgroup_subsys_state *css,
+				      struct cftype *cft, u64 val)
+{
+	struct mem_cgroup *memcg = mem_cgroup_from_css(css);
+
+	if (val == 1 || val == 0)
+		memcg->use_asi = val;
+	else
+		return -EINVAL;
+
+	return 0;
+}
+
+#endif
+
 static unsigned long mem_cgroup_usage(struct mem_cgroup *memcg, bool swap)
 {
 	unsigned long val;
@@ -4888,6 +4911,13 @@ static struct cftype mem_cgroup_legacy_files[] = {
 		.write_u64 = mem_cgroup_hierarchy_write,
 		.read_u64 = mem_cgroup_hierarchy_read,
 	},
+#ifdef CONFIG_ADDRESS_SPACE_ISOLATION
+	{
+		.name = "use_asi",
+		.write_u64 = mem_cgroup_asi_write,
+		.read_u64 = mem_cgroup_asi_read,
+	},
+#endif
 	{
 		.name = "cgroup.event_control",		/* XXX: for compat */
 		.write = memcg_write_event_control,
-- 
2.35.1.473.g83b2b277ed-goog