rst: Lock rst memory allocations earlier

After we got the total remapable rst memory size, we no longer
can allocate from it, otherwise the bootstrap area will not
have enough size.
Signed-off-by: Pavel Emelyanov <xemul@parallels.com>

cr-restore.c

@@ -2639,6 +2639,7 @@ static int sigreturn_restore(pid_t pid, CoreEntry *core)
 	long ret;
 
 	long restore_bootstrap_len;
+	long rst_mem_size;
 
 	struct task_restore_args *task_args;
 	struct thread_restore_args *thread_args;
@@ -2734,8 +2735,8 @@ static int sigreturn_restore(pid_t pid, CoreEntry *core)
 	if (ret < 0)
 		goto err;
 
-	restore_bootstrap_len = restorer_len + args_len +
-				rst_mem_remap_size();
+	rst_mem_size = rst_mem_lock();
+	restore_bootstrap_len = restorer_len + args_len + rst_mem_size;
 
 #ifdef CONFIG_VDSO
 	/*
@@ -2855,7 +2856,7 @@ static int sigreturn_restore(pid_t pid, CoreEntry *core)
 	task_args->task_entries = rst_mem_remap_ptr(task_entries_pos, RM_SHREMAP);
 
 	task_args->rst_mem = mem;
-	task_args->rst_mem_size = rst_mem_remap_size();
+	task_args->rst_mem_size = rst_mem_size;
 
 	task_args->bootstrap_start = (void *)exec_mem_hint;
 	task_args->bootstrap_len = restore_bootstrap_len;
@@ -2972,7 +2973,7 @@ static int sigreturn_restore(pid_t pid, CoreEntry *core)
 	 * since we need it being accessible even when own
 	 * self-vmas are unmaped.
 	 */
-	mem += rst_mem_remap_size();
+	mem += rst_mem_size;
 	task_args->vdso_rt_parked_at = (unsigned long)mem + vdso_rt_delta;
 	task_args->vdso_sym_rt = vdso_sym_rt;
 	task_args->vdso_rt_size = vdso_rt_size;

include/rst-malloc.h

@@ -62,7 +62,7 @@ extern void rst_mem_free_last(int type);
 /*
  * Routines to remap SHREMAP and PRIVATE into restorer address space
  */
-extern unsigned long rst_mem_remap_size(void);
+extern unsigned long rst_mem_lock(void);
 extern int rst_mem_remap(void *to);
 
 #endif /* __CR_RST_MALLOC__H__ */

rst-malloc.c

@@ -169,8 +169,15 @@ void rst_mem_free_last(int type)
 	t->last = 0; /* next free_last would be no-op */
 }
 
-unsigned long rst_mem_remap_size(void)
+unsigned long rst_mem_lock(void)
 {
+	/*
+	 * Don't allow further allocations from rst_mem since we're
+	 * going to get the bootstrap area and remap all the stuff
+	 * into it. The SHREMAP and SHARED should be already locked
+	 * in the rst_mem_switch_to_private().
+	 */
+	rst_mems[RM_PRIVATE].enabled = false;
 	return rst_mems[RM_PRIVATE].size + rst_mems[RM_SHREMAP].size;
 }
 
@@ -178,7 +185,7 @@ static int rst_mem_remap_one(struct rst_mem_type_s *t, void *to)
 {
 	void *aux;
 
-	BUG_ON(!t->remapable);
+	BUG_ON(!t->remapable || t->enabled);
 
 	if (!t->buf)
 		/*
@@ -196,7 +203,6 @@ static int rst_mem_remap_one(struct rst_mem_type_s *t, void *to)
 	}
 
 	t->buf = aux;
-	t->enabled = false;
 	return 0;
 }