[RFC PATCH 22/47] mm: asi: Added refcounting when initilizing an asi

Junaid Shahid
2022-02-23T05:26:23Z

From: Ofir Weisse <oweisse@google.com>

Some KVM tests initilize multiple VMs in a single process. For these
cases, we want to suppurt multiple callse to asi_init() before a single
asi_destroy is called. We want the initilization to happen exactly once.
IF asi_destroy() is called, release the resources only if the counter
reached zero. In our current implementation, asi's are tied to
a specific mm. This may change in a future implementation. In which
case, the mutex for the refcounting will need to move to struct asi.

Signed-off-by: Ofir Weisse <oweisse@google.com>

---
arch/x86/include/asm/asi.h | 1 +
arch/x86/mm/asi.c | 52 +++++++++++++++++++++++++++++++++-----
include/linux/mm_types.h | 2 ++
kernel/fork.c | 3 +++
4 files changed, 51 insertions(+), 7 deletions(-)

diff --git a/arch/x86/include/asm/asi.h b/arch/x86/include/asm/asi.h
index e3cbf6d8801e..2dc465f78bcc 100644
--- a/arch/x86/include/asm/asi.h
+++ b/arch/x86/include/asm/asi.h
@@ -40,6 +40,7 @@ struct asi {
pgd_t *pgd;
struct asi_class *class;
struct mm_struct *mm;
+ int64_t asi_ref_count;
};

DECLARE_PER_CPU_ALIGNED(struct asi_state, asi_cpu_state);
diff --git a/arch/x86/mm/asi.c b/arch/x86/mm/asi.c
index 91e5ff1224ff..ac35323193a3 100644
--- a/arch/x86/mm/asi.c
+++ b/arch/x86/mm/asi.c
@@ -282,9 +282,25 @@ static int __init asi_global_init(void)
}
subsys_initcall(asi_global_init)

+/* We're assuming we hold mm->asi_init_lock */
+static void __asi_destroy(struct asi *asi)
+{
+ if (!boot_cpu_has(X86_FEATURE_ASI))
+ return;
+
+ /* If refcount is non-zero, it means asi_init() was called multiple
+ * times. We free the asi pgd only when the last VM is destroyed. */
+ if (--(asi->asi_ref_count) > 0)
+ return;
+
+ asi_free_pgd(asi);
+ memset(asi, 0, sizeof(struct asi));
+}
+
int asi_init(struct mm_struct *mm, int asi_index, struct asi **out_asi)
{
- struct asi *asi = &mm->asi[asi_index];
+ int err = 0;
+ struct asi *asi = &mm->asi[asi_index];

*out_asi = NULL;

@@ -295,6 +311,15 @@ int asi_init(struct mm_struct *mm, int asi_index, struct asi **out_asi)
WARN_ON(asi_index == 0 || asi_index >= ASI_MAX_NUM);
WARN_ON(asi->pgd != NULL);

+ /* Currently, mm and asi structs are conceptually tied together. In
+ * future implementations an asi object might be unrelated to a specicic
+ * mm. In that future implementation - the mutex will have to be inside
+ * asi. */
+ mutex_lock(&mm->asi_init_lock);
+
+ if (asi->asi_ref_count++ > 0)
+ goto exit_unlock; /* err is 0 */
+
/*
* For now, we allocate 2 pages to avoid any potential problems with
* KPTI code. This won't be needed once KPTI is folded into the ASI
@@ -302,8 +327,10 @@ int asi_init(struct mm_struct *mm, int asi_index, struct asi **out_asi)
*/
asi->pgd = (pgd_t *)__get_free_pages(GFP_PGTABLE_USER,
PGD_ALLOCATION_ORDER);
- if (!asi->pgd)
- return -ENOMEM;
+ if (!asi->pgd) {
+ err = -ENOMEM;
+ goto exit_unlock;
+ }

asi->class = &asi_class[asi_index];
asi->mm = mm;
@@ -328,19 +355,30 @@ int asi_init(struct mm_struct *mm, int asi_index, struct asi **out_asi)
set_pgd(asi->pgd + i, asi_global_nonsensitive_pgd[i]);
}

- *out_asi = asi;
+exit_unlock:
+ if (err)
+ __asi_destroy(asi);

- return 0;
+ /* This unlock signals future asi_init() callers that we finished. */
+ mutex_unlock(&mm->asi_init_lock);
+
+ if (!err)
+ *out_asi = asi;
+ return err;
}
EXPORT_SYMBOL_GPL(asi_init);

void asi_destroy(struct asi *asi)
{
+ struct mm_struct *mm;
+
if (!boot_cpu_has(X86_FEATURE_ASI) || !asi)
return;

- asi_free_pgd(asi);
- memset(asi, 0, sizeof(struct asi));
+ mm = asi->mm;
+ mutex_lock(&mm->asi_init_lock);
+ __asi_destroy(asi);
+ mutex_unlock(&mm->asi_init_lock);
}
EXPORT_SYMBOL_GPL(asi_destroy);

diff --git a/include/linux/mm_types.h b/include/linux/mm_types.h
index f9702d070975..e6980ae31323 100644
--- a/include/linux/mm_types.h
+++ b/include/linux/mm_types.h
@@ -16,6 +16,7 @@
#include <linux/page-flags-layout.h>
#include <linux/workqueue.h>
#include <linux/seqlock.h>
+#include <linux/mutex.h>

#include <asm/mmu.h>
#include <asm/asi.h>
@@ -628,6 +629,7 @@ struct mm_struct {
* these resources for every mm in the system, we expect that
* only VM mm's will have this flag set. */
bool asi_enabled;
+ struct mutex asi_init_lock;
#endif
struct user_namespace *user_ns;

diff --git a/kernel/fork.c b/kernel/fork.c
index dd5a86e913ea..68b3aeab55ac 100644
--- a/kernel/fork.c
+++ b/kernel/fork.c
@@ -1084,6 +1084,9 @@ static struct mm_struct *mm_init(struct mm_struct *mm, struct task_struct *p,

mm->user_ns = get_user_ns(user_ns);

+#ifdef CONFIG_ADDRESS_SPACE_ISOLATION
+ mutex_init(&mm->asi_init_lock);
+#endif
return mm;

fail_noasi:
--
2.35.1.473.g83b2b277ed-goog