[RFC v2 16/26] mm/asi: Option to map current task into ASI

[Alexandre Chartre <alexandre.chartre@xxxxxxxxxx>]

Add an option to map the current task into an ASI page-table.
The task is mapped when entering isolation and unmapped on
abort/exit.

Signed-off-by: Alexandre Chartre <alexandre.chartre@xxxxxxxxxx>
---
 arch/x86/include/asm/asi.h  |    2 ++
 arch/x86/mm/asi.c           |   25 +++++++++++++++++++++----
 arch/x86/mm/asi_pagetable.c |    4 ++--
 3 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/arch/x86/include/asm/asi.h b/arch/x86/include/asm/asi.h
index 1ac8fd3..a277e43 100644
--- a/arch/x86/include/asm/asi.h
+++ b/arch/x86/include/asm/asi.h
@@ -17,6 +17,7 @@
  */
 #define ASI_MAP_STACK_CANARY	0x01	/* map stack canary */
 #define ASI_MAP_CPU_PTR		0x02	/* for get_cpu_var()/this_cpu_ptr() */
+#define ASI_MAP_CURRENT_TASK	0x04	/* map the current task */
 
 enum page_table_level {
 	PGT_LEVEL_PTE,
@@ -31,6 +32,7 @@ enum page_table_level {
 struct asi {
 	spinlock_t		lock;		/* protect all attributes */
 	pgd_t			*pgd;		/* ASI page-table */
+	int			mapping_flags;	/* map flags */
 	struct list_head	mapping_list;	/* list of VA range mapping */
 
 	/*
diff --git a/arch/x86/mm/asi.c b/arch/x86/mm/asi.c
index f049438..acd1135 100644
--- a/arch/x86/mm/asi.c
+++ b/arch/x86/mm/asi.c
@@ -28,6 +28,7 @@ struct asi_map_option {
 struct asi_map_option asi_map_percpu_options[] = {
 	{ ASI_MAP_STACK_CANARY, &fixed_percpu_data, sizeof(fixed_percpu_data) },
 	{ ASI_MAP_CPU_PTR, &this_cpu_off, sizeof(this_cpu_off) },
+	{ ASI_MAP_CURRENT_TASK, &current_task, sizeof(current_task) },
 };
 
 static void asi_log_fault(struct asi *asi, struct pt_regs *regs,
@@ -96,8 +97,9 @@ bool asi_fault(struct pt_regs *regs, unsigned long error_code,
 	return true;
 }
 
-static int asi_init_mapping(struct asi *asi, int flags)
+static int asi_init_mapping(struct asi *asi)
 {
+	int flags = asi->mapping_flags;
 	struct asi_map_option *option;
 	int i, err;
 
@@ -164,8 +166,9 @@ struct asi *asi_create(int map_flags)
 	spin_lock_init(&asi->lock);
 	spin_lock_init(&asi->fault_lock);
 	asi_init_backend(asi);
+	asi->mapping_flags = map_flags;
 
-	err = asi_init_mapping(asi, map_flags);
+	err = asi_init_mapping(asi);
 	if (err)
 		goto error;
 
@@ -248,6 +251,15 @@ int asi_enter(struct asi *asi)
 		goto err_clear_asi;
 
 	/*
+	 * Optionally, also map the current task.
+	 */
+	if (asi->mapping_flags & ASI_MAP_CURRENT_TASK) {
+		err = asi_map(asi, current, sizeof(struct task_struct));
+		if (err)
+			goto err_unmap_stack;
+	}
+
+	/*
 	 * Instructions ordering is important here because we should be
 	 * able to deal with any interrupt/exception which will abort
 	 * the isolation and restore CR3 to its original value:
@@ -269,7 +281,7 @@ int asi_enter(struct asi *asi)
 	if (!original_cr3) {
 		WARN_ON(1);
 		err = -EINVAL;
-		goto err_unmap_stack;
+		goto err_unmap_task;
 	}
 	asi_session->original_cr3 = original_cr3;
 
@@ -286,6 +298,9 @@ int asi_enter(struct asi *asi)
 
 	return 0;
 
+err_unmap_task:
+	if (asi->mapping_flags & ASI_MAP_CURRENT_TASK)
+		asi_unmap(asi, current);
 err_unmap_stack:
 	asi_unmap(asi, current->stack);
 err_clear_asi:
@@ -345,8 +360,10 @@ void asi_exit(struct asi *asi)
 	 */
 	asi_session->abort_depth = 0;
 
-	/* unmap stack */
+	/* unmap stack and task */
 	asi_unmap(asi, current->stack);
+	if (asi->mapping_flags & ASI_MAP_CURRENT_TASK)
+		asi_unmap(asi, current);
 }
 EXPORT_SYMBOL(asi_exit);
 
diff --git a/arch/x86/mm/asi_pagetable.c b/arch/x86/mm/asi_pagetable.c
index bcc95f2..8076626 100644
--- a/arch/x86/mm/asi_pagetable.c
+++ b/arch/x86/mm/asi_pagetable.c
@@ -714,7 +714,7 @@ int asi_map_range(struct asi *asi, void *ptr, size_t size,
 	 * Don't log info the current stack because it is mapped/unmapped
 	 * everytime we enter/exit isolation.
 	 */
-	if (ptr != current->stack) {
+	if (ptr != current->stack && ptr != current) {
 		pr_debug("ASI %p: MAP %px/%lx/%d -> %lx-%lx\n",
 			 asi, ptr, size, level, map_addr, map_end);
 		if (map_addr < addr)
@@ -1001,7 +1001,7 @@ void asi_unmap(struct asi *asi, void *ptr)
 	 * Don't log info the current stack because it is mapped/unmapped
 	 * everytime we enter/exit isolation.
 	 */
-	if (ptr != current->stack) {
+	if (ptr != current->stack && ptr != current) {
 		pr_debug("ASI %p: UNMAP %px/%lx/%d\n", asi, ptr,
 			 range_mapping->size, range_mapping->level);
 	}
-- 
1.7.1