shadow stack: restore x86 shstk VMA earlier during premapped VMAs remap loop

criu/arch/x86/include/asm/shstk.h

@@ -73,6 +73,23 @@ int arch_shstk_trampoline(struct pstree_item *item, CoreEntry *core,
 		      int (*func)(void *arg), void *arg);
 #define arch_shstk_trampoline arch_shstk_trampoline
 
+static always_inline long shstk_restorer_stack_size(void)
+{
+	return PAGE_SIZE;
+}
+#define shstk_restorer_stack_size shstk_restorer_stack_size
+static always_inline void shstk_set_restorer_stack(struct rst_shstk_info *info, void *ptr)
+{
+	info->tmp_shstk = (unsigned long)ptr;
+}
+#define shstk_set_restorer_stack shstk_set_restorer_stack
+
+static always_inline long shstk_min_mmap_addr(struct rst_shstk_info *info, unsigned long __maybe_unused def)
+{
+	return !(info->cet & ARCH_SHSTK_SHSTK) ? def : (4UL << 30);
+}
+#define shstk_min_mmap_addr shstk_min_mmap_addr
+
 #ifdef CR_NOGLIBC
 
 #include <compel/plugins/std/syscall.h>
@@ -147,31 +164,51 @@ static inline int shstk_finalize(void)
 }
 
 /*
- * Restore contents of the shadow stack and set shadow stack pointer
+ * Create shadow stack vma and restore its content from premmapped anonymous (non-shstk) vma
  */
-static always_inline int shstk_restore(struct rst_shstk_info *cet)
+static always_inline int shstk_vma_restore(VmaEntry *vma_entry)
 {
-	unsigned long *shstk_data = (unsigned long *)cet->premmaped_addr;
-	unsigned long ssp = cet->vma_start + cet->vma_size - 8;
-	unsigned long shstk_top = cet->vma_size / 8 - 1;
-	unsigned long val;
+	long shstk, i;
+	unsigned long *shstk_data = (void *)vma_premmaped_start(vma_entry);
+	unsigned long vma_size = vma_entry_len(vma_entry);
 	long ret;
 
-	if (!(cet->cet & ARCH_SHSTK_SHSTK))
-		return 0;
-
-	if (shstk_map(cet->vma_start, cet->vma_size))
+	shstk = sys_map_shadow_stack(0, vma_size, SHADOW_STACK_SET_TOKEN);
+	if (shstk < 0) {
+		pr_err("Failed to map shadow stack: %ld\n", shstk);
 		return -1;
+	}
+
+	/* restore shadow stack contents */
+	for (i = 0; i < vma_size / 8; i++)
+		wrssq(shstk + i * 8, shstk_data[i]);
+
+	ret = sys_munmap(shstk_data, vma_size);
+	if (ret < 0) {
+		pr_err("Failed to unmap premmaped shadow stack\n");
+		return ret;
+	}
 
 	/*
-	 * Switch shadow stack from temporary location to the actual task's
-	 * shadow stack VMA
+	 * From that point premapped vma is (shstk) and we need
+	 * to mremap() it to the final location. Originally premapped
+	 * (shstk_data) has been unmapped already.
 	 */
-	shstk_switch_ssp(ssp);
+	vma_premmaped_start(vma_entry) = shstk;
 
-	/* restore shadow stack contents */
-	for (; ssp >= cet->ssp; ssp -= 8, shstk_top--)
-		wrssq(ssp, shstk_data[shstk_top]);
+	return 0;
+}
+#define shstk_vma_restore shstk_vma_restore
+
+/*
+ * Restore contents of the shadow stack and set shadow stack pointer
+ */
+static always_inline int shstk_restore(struct rst_shstk_info *cet)
+{
+	unsigned long ssp, val;
+
+	if (!(cet->cet & ARCH_SHSTK_SHSTK))
+		return 0;
 
 	/*
 	 * Add tokens for sigreturn frame and for switch of the shadow stack.
@@ -182,6 +219,7 @@ static always_inline int shstk_restore(struct rst_shstk_info *cet)
 	 */
 
 	/* token for sigreturn frame */
+	ssp = cet->ssp - 8;
 	val = ALIGN_DOWN(cet->ssp, 8) | SHSTK_DATA_BIT;
 	wrssq(ssp, val);
 
@@ -193,12 +231,6 @@ static always_inline int shstk_restore(struct rst_shstk_info *cet)
 	/* reset shadow stack pointer to the proper location */
 	shstk_switch_ssp(ssp);
 
-	ret = sys_munmap(shstk_data, cet->vma_size + PAGE_SIZE);
-	if (ret < 0) {
-		pr_err("Failed to unmap premmaped shadow stack\n");
-		return ret;
-	}
-
 	return shstk_finalize();
 }
 #define arch_shstk_restore shstk_restore

criu/arch/x86/shstk.c

@@ -45,7 +45,6 @@ static int shstk_prepare_task(struct vm_area_list *vmas,
 			shstk->vma_start = vma->e->start;
 			shstk->vma_size = size;
 			shstk->premmaped_addr = premmaped_addr;
-			shstk->tmp_shstk = premmaped_addr + size;
 
 			break;
 		}

criu/cr-restore.c

@@ -2431,16 +2431,15 @@ int cr_restore_tasks(void)
 	return ret;
 }
 
-static long restorer_get_vma_hint(struct list_head *tgt_vma_list, struct list_head *self_vma_list, long vma_len)
+static long restorer_get_vma_hint(struct list_head *tgt_vma_list, struct list_head *self_vma_list, long min_addr, long vma_len)
 {
 	struct vma_area *t_vma, *s_vma;
-	long prev_vma_end = 0;
+	long prev_vma_end = min_addr;
 	struct vma_area end_vma;
 	VmaEntry end_e;
 
 	end_vma.e = &end_e;
 	end_e.start = end_e.end = kdat.task_size;
-	prev_vma_end = kdat.mmap_min_addr;
 
 	s_vma = list_first_entry(self_vma_list, struct vma_area, list);
 	t_vma = list_first_entry(tgt_vma_list, struct vma_area, list);
@@ -3196,7 +3195,7 @@ static int sigreturn_restore(pid_t pid, struct task_restore_args *task_args, uns
 
 	rst_mem_size = rst_mem_lock();
 	memzone_size = round_up(sizeof(struct restore_mem_zone) * current->nr_threads, page_size());
-	task_args->bootstrap_len = restorer_len + memzone_size + alen + rst_mem_size;
+	task_args->bootstrap_len = restorer_len + memzone_size + alen + rst_mem_size + shstk_restorer_stack_size();
 	BUG_ON(task_args->bootstrap_len & (PAGE_SIZE - 1));
 	pr_info("%d threads require %ldK of memory\n", current->nr_threads, KBYTES(task_args->bootstrap_len));
 
@@ -3226,7 +3225,9 @@ static int sigreturn_restore(pid_t pid, struct task_restore_args *task_args, uns
 	 * or inited from scratch).
 	 */
 
-	mem = (void *)restorer_get_vma_hint(&vmas->h, &self_vmas.h, task_args->bootstrap_len);
+	mem = (void *)restorer_get_vma_hint(&vmas->h, &self_vmas.h,
+					    shstk_min_mmap_addr(&task_args->shstk, kdat.mmap_min_addr),
+					    task_args->bootstrap_len);
 	if (mem == (void *)-1) {
 		pr_err("No suitable area for task_restore bootstrap (%ldK)\n", task_args->bootstrap_len);
 		goto err;
@@ -3465,6 +3466,10 @@ static int sigreturn_restore(pid_t pid, struct task_restore_args *task_args, uns
 	 * self-vmas are unmaped.
 	 */
 	mem += rst_mem_size;
+
+	shstk_set_restorer_stack(&task_args->shstk, mem);
+	mem += shstk_restorer_stack_size();
+
 	task_args->vdso_rt_parked_at = (unsigned long)mem;
 	task_args->vdso_maps_rt = vdso_maps_rt;
 	task_args->vdso_rt_size = vdso_rt_size;

criu/include/restore.h

@@ -9,6 +9,7 @@ extern int arch_set_thread_regs_nosigrt(struct pid *pid);
 
 struct task_restore_args;
 struct pstree_item;
+struct rst_shstk_info;
 
 #ifndef arch_shstk_prepare
 static inline int arch_shstk_prepare(struct pstree_item *item,
@@ -38,4 +39,25 @@ static inline int arch_shstk_trampoline(struct pstree_item *item, CoreEntry *cor
 #define arch_shstk_trampoline arch_shstk_trampoline
 #endif
 
+#ifndef shstk_restorer_stack_size
+static always_inline long shstk_restorer_stack_size(void)
+{
+	return 0;
+}
+#endif
+
+#ifndef shstk_set_restorer_stack
+static always_inline long shstk_set_restorer_stack(struct rst_shstk_info *info, void *ptr)
+{
+	return 0;
+}
+#endif
+
+#ifndef shstk_min_mmap_addr
+static always_inline long shstk_min_mmap_addr(struct rst_shstk_info *info, unsigned long def)
+{
+	return def;
+}
+#endif
+
 #endif

criu/include/restorer.h

@@ -357,4 +357,11 @@ static inline int arch_shstk_restore(struct rst_shstk_info *shstk)
 #define arch_shstk_restore arch_shstk_restore
 #endif
 
+#ifndef shstk_vma_restore
+static always_inline int shstk_vma_restore(VmaEntry *vma_entry)
+{
+	return -1;
+}
+#endif
+
 #endif /* __CR_RESTORER_H__ */

criu/mem.c

@@ -787,8 +787,6 @@ int prepare_mm_pid(struct pstree_item *i)
 			ri->vmas.rst_priv_size += vma_area_len(vma);
 			if (vma_has_guard_gap_hidden(vma))
 				ri->vmas.rst_priv_size += PAGE_SIZE;
-			if (vma_area_is(vma, VMA_AREA_SHSTK))
-				ri->vmas.rst_priv_size += PAGE_SIZE;
 		}
 
 		pr_info("vma 0x%" PRIx64 " 0x%" PRIx64 "\n", vma->e->start, vma->e->end);
@@ -931,13 +929,6 @@ static int premap_private_vma(struct pstree_item *t, struct vma_area *vma, void
 
 	size = vma_entry_len(vma->e);
 
-	/*
-	 * map an extra page for shadow stack VMAs, it will be used as a
-	 * temporary shadow stack
-	 */
-	if (vma_area_is(vma, VMA_AREA_SHSTK))
-		size += PAGE_SIZE;
-
 	if (!vma_inherited(vma)) {
 		int flag = 0;
 		/*

criu/pie/restorer.c

@@ -1112,6 +1112,23 @@ static int vma_remap(VmaEntry *vma_entry, int uffd)
 
 	pr_info("Remap %lx->%lx len %lx\n", src, dst, len);
 
+	/*
+	 * SHSTK VMAs are a bit special, in fact we create shstk vma right in the
+	 * shstk_vma_restore() and populate it with contents from a premapped VMA
+	 * (which in turns is just a normal anonymous VMA!). Then, we munmap() this
+	 * premapped VMA. After, we need to adjust vma_premmaped_start(vma_entry)
+	 * to point to a created shstk vma and treat it as a premmaped one in vma_remap().
+	 */
+	if (vma_entry_is(vma_entry, VMA_AREA_SHSTK)) {
+		if (shstk_vma_restore(vma_entry)) {
+			pr_err("Unable to prepare shadow stack vma for remap %lx -> %lx\n", src, dst);
+			return -1;
+		}
+
+		/* shstk_vma_restore() modifies vma premapped address */
+		src = vma_premmaped_start(vma_entry);
+	}
+
 	if (src - dst < len)
 		guard = dst;
 	else if (dst - src < len)
@@ -1811,13 +1828,6 @@ __visible long __export_restore_task(struct task_restore_args *args)
 		if (vma_entry->start > vma_entry->shmid)
 			break;
 
-		/*
-		 * shadow stack VMAs cannot be remapped, they must be
-		 * recreated with map_shadow_stack system call
-		 */
-		if (vma_entry_is(vma_entry, VMA_AREA_SHSTK))
-			continue;
-
 		if (vma_remap(vma_entry, args->uffd))
 			goto core_restore_end;
 	}
@@ -1835,13 +1845,6 @@ __visible long __export_restore_task(struct task_restore_args *args)
 		if (vma_entry->start < vma_entry->shmid)
 			break;
 
-		/*
-		 * shadow stack VMAs cannot be remapped, they must be
-		 * recreated with map_shadow_stack system call
-		 */
-		if (vma_entry_is(vma_entry, VMA_AREA_SHSTK))
-			continue;
-
 		if (vma_remap(vma_entry, args->uffd))
 			goto core_restore_end;
 	}

test/zdtm/Makefile.inc

@@ -66,6 +66,11 @@ endif
         export PKG_CONFIG_PATH
 endif
 
+ifeq ($(SHSTK_ENABLE),1)
+        CFLAGS  += -mshstk
+        LDFLAGS += -Wl,-z,shstk
+endif
+
 define pkg-libs
         $(shell PKG_CONFIG_PATH="$(PKG_CONFIG_PATH)" $(PKG_CONFIG) --libs $(1))
 endef