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

Junaid Shahid
2022-02-23T05:25:18Z

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