compel: Provide compel_set_task_ext_regs()

Arch-dependend way to restore extended registers set.
Use it straight-away to restore per-thread registers.

Signed-off-by: Dmitry Safonov <dima@arista.com>

8 files changed, 165 insertions, 8 deletions

diff --git a/compel/arch/aarch64/src/lib/infect.c b/compel/arch/aarch64/src/lib/infect.c
index 39428cd8d..586eedac0 100644
--- a/compel/arch/aarch64/src/lib/infect.c
+++ b/compel/arch/aarch64/src/lib/infect.c
@@ -89,6 +89,21 @@ err:
 	return ret;
 }
 
+int compel_set_task_ext_regs(pid_t pid, user_fpregs_struct_t *ext_regs)
+{
+	struct iovec iov;
+
+	pr_info("Restoring GP/FPU registers for %d\n", pid);
+
+	iov.iov_base = ext_regs;
+	iov.iov_len = sizeof(*ext_regs);
+	if (ptrace(PTRACE_SETREGSET, pid, NT_PRFPREG, &iov)) {
+		pr_perror("Failed to set FPU registers for %d", pid);
+		return -1;
+	}
+	return 0;
+}
+
 int compel_syscall(struct parasite_ctl *ctl, int nr, long *ret,
 					unsigned long arg1,
 					unsigned long arg2,
diff --git a/compel/arch/arm/src/lib/infect.c b/compel/arch/arm/src/lib/infect.c
index dd5aa9784..1ddb48654 100644
--- a/compel/arch/arm/src/lib/infect.c
+++ b/compel/arch/arm/src/lib/infect.c
@@ -105,6 +105,17 @@ err:
 	return ret;
 }
 
+int compel_set_task_ext_regs(pid_t pid, user_fpregs_struct_t *ext_regs)
+{
+	pr_info("Restoring GP/FPU registers for %d\n", pid);
+
+	if (ptrace(PTRACE_SETVFPREGS, pid, NULL, ext_regs)) {
+		pr_perror("Can't set FPU registers for %d", pid);
+		return -1;
+	}
+	return 0;
+}
+
 int compel_syscall(struct parasite_ctl *ctl, int nr, long *ret,
 		unsigned long arg1,
 		unsigned long arg2,
diff --git a/compel/arch/mips/src/lib/infect.c b/compel/arch/mips/src/lib/infect.c
index 6a53b3c3b..0c8067f88 100755
--- a/compel/arch/mips/src/lib/infect.c
+++ b/compel/arch/mips/src/lib/infect.c
@@ -129,6 +129,8 @@ int compel_get_task_regs(pid_t pid, user_regs_struct_t *regs,
 	user_fpregs_struct_t xsave = {  }, *xs = ext_regs ? ext_regs : &xsave;
 	int ret = -1;
 
+	pr_info("Dumping GP/FPU registers for %d\n", pid);
+
 	if (ptrace(PTRACE_GETFPREGS, pid, NULL, xs)) {
 	    pr_perror("Can't obtain FPU registers for %d", pid);
 	    return ret;
@@ -156,6 +158,17 @@ int compel_get_task_regs(pid_t pid, user_regs_struct_t *regs,
 	return ret;
 }
 
+int compel_set_task_ext_regs(pid_t pid, user_fpregs_struct_t *ext_regs)
+{
+	pr_info("Restoring GP/FPU registers for %d\n", pid);
+
+	if (ptrace(PTRACE_SETFPREGS, pid, NULL, ext_regs)) {
+		pr_perror("Can't set FPU registers for %d", pid);
+		return -1;
+	}
+	return 0;
+}
+
 int compel_syscall(struct parasite_ctl *ctl, int nr, long *ret,
 		unsigned long arg1,
 		unsigned long arg2,
diff --git a/compel/arch/ppc64/src/lib/infect.c b/compel/arch/ppc64/src/lib/infect.c
index 339e01209..5797fb16d 100644
--- a/compel/arch/ppc64/src/lib/infect.c
+++ b/compel/arch/ppc64/src/lib/infect.c
@@ -386,6 +386,34 @@ int compel_get_task_regs(pid_t pid, user_regs_struct_t *regs,
 	return save(arg, regs, fpregs);
 }
 
+int compel_set_task_ext_regs(pid_t pid, user_fpregs_struct_t *ext_regs)
+{
+	int ret = 0;
+
+	pr_info("Restoring GP/FPU registers for %d\n", pid);
+
+	/* XXX: should restore TM registers somehow? */
+	if (ext_regs->flags & USER_FPREGS_FL_FP) {
+		if (ptrace(PTRACE_SETFPREGS, pid, 0, (void *)&ext_regs->fpregs) < 0) {
+			pr_perror("Couldn't set floating-point registers");
+			ret = -1;
+		}
+	}
+
+	if (ext_regs->flags & USER_FPREGS_FL_ALTIVEC) {
+		if (ptrace(PTRACE_SETVRREGS, pid, 0, (void*)&ext_regs->vrregs) < 0) {
+			pr_perror("Couldn't set Altivec registers");
+			ret = -1;
+		}
+		if (ptrace(PTRACE_SETVSRREGS, pid, 0, (void*)ext_regs->vsxregs) < 0) {
+			pr_perror("Couldn't set VSX registers");
+			ret = -1;
+		}
+	}
+
+	return ret;
+}
+
 int compel_syscall(struct parasite_ctl *ctl, int nr, long *ret,
 		unsigned long arg1,
 		unsigned long arg2,
diff --git a/compel/arch/s390/src/lib/infect.c b/compel/arch/s390/src/lib/infect.c
index 5b3c0d62e..557b21203 100644
--- a/compel/arch/s390/src/lib/infect.c
+++ b/compel/arch/s390/src/lib/infect.c
@@ -370,6 +370,61 @@ int compel_get_task_regs(pid_t pid, user_regs_struct_t *regs,
 	return save(arg, regs, fpregs);
 }
 
+int compel_set_task_ext_regs(pid_t pid, user_fpregs_struct_t *ext_regs)
+{
+	struct iovec iov;
+	int ret = 0;
+
+	iov.iov_base = &ext_regs->prfpreg;
+	iov.iov_len = sizeof(ext_regs->prfpreg);
+	if (ptrace(PTRACE_SETREGSET, pid, NT_PRFPREG, &iov) < 0) {
+		pr_perror("Couldn't set floating-point registers");
+		ret = -1;
+	}
+
+	if (ext_regs->flags & USER_FPREGS_VXRS) {
+		iov.iov_base = &ext_regs->vxrs_low;
+		iov.iov_len = sizeof(ext_regs->vxrs_low);
+		if (ptrace(PTRACE_SETREGSET, pid, NT_S390_VXRS_LOW, &iov) < 0) {
+			pr_perror("Couldn't set VXRS_LOW\n");
+			ret = -1;
+		}
+
+		iov.iov_base = &ext_regs->vxrs_high;
+		iov.iov_len = sizeof(ext_regs->vxrs_high);
+		if (ptrace(PTRACE_SETREGSET, pid, NT_S390_VXRS_HIGH, &iov) < 0) {
+			pr_perror("Couldn't set VXRS_HIGH\n");
+			ret = -1;
+		}
+	}
+
+	if (ext_regs->flags & USER_GS_CB) {
+		iov.iov_base = &ext_regs->gs_cb;
+		iov.iov_len = sizeof(ext_regs->gs_cb);
+		if (ptrace(PTRACE_SETREGSET, pid, NT_S390_GS_CB, &iov) < 0) {
+			pr_perror("Couldn't set GS_CB\n");
+			ret = -1;
+		}
+		iov.iov_base = &ext_regs->gs_bc;
+		iov.iov_len = sizeof(ext_regs->gs_bc);
+		if (ptrace(PTRACE_SETREGSET, pid, NT_S390_GS_BC, &iov) < 0) {
+			pr_perror("Couldn't set GS_BC\n");
+			ret = -1;
+		}
+	}
+
+	if (ext_regs->flags & USER_RI_CB) {
+		iov.iov_base = &ext_regs->ri_cb;
+		iov.iov_len = sizeof(ext_regs->ri_cb);
+		if (ptrace(PTRACE_SETREGSET, pid, NT_S390_RI_CB, &iov) < 0) {
+			pr_perror("Couldn't set RI_CB\n");
+			ret = -1;
+		}
+	}
+
+	return ret;
+}
+
 /*
  * Injected syscall instruction
  */
diff --git a/compel/arch/x86/src/lib/infect.c b/compel/arch/x86/src/lib/infect.c
index 9d81083f7..566238d2d 100644
--- a/compel/arch/x86/src/lib/infect.c
+++ b/compel/arch/x86/src/lib/infect.c
@@ -398,6 +398,31 @@ err:
 	return ret;
 }
 
+int compel_set_task_ext_regs(pid_t pid, user_fpregs_struct_t *ext_regs)
+{
+	struct iovec iov;
+
+	pr_info("Restoring GP/FPU registers for %d\n", pid);
+
+	if (!compel_cpu_has_feature(X86_FEATURE_OSXSAVE)) {
+		if (ptrace(PTRACE_SETFPREGS, pid, NULL, ext_regs)) {
+			pr_perror("Can't set FPU registers for %d", pid);
+			return -1;
+		}
+		return 0;
+	}
+
+	iov.iov_base = ext_regs;
+	iov.iov_len = sizeof(*ext_regs);
+
+	if (ptrace(PTRACE_SETREGSET, pid, (unsigned int)NT_X86_XSTATE, &iov) < 0) {
+		pr_perror("Can't set FPU registers for %d", pid);
+		return -1;
+	}
+
+	return 0;
+}
+
 int compel_syscall(struct parasite_ctl *ctl, int nr, long *ret,
 		unsigned long arg1,
 		unsigned long arg2,
diff --git a/compel/include/infect-priv.h b/compel/include/infect-priv.h
index 16114164c..d4f19307d 100644
--- a/compel/include/infect-priv.h
+++ b/compel/include/infect-priv.h
@@ -73,6 +73,7 @@ extern bool arch_can_dump_task(struct parasite_ctl *ctl);
 extern int compel_get_task_regs(pid_t pid, user_regs_struct_t *regs,
 			user_fpregs_struct_t *ext_regs, save_regs_t save,
 			void *arg, unsigned long flags);
+extern int compel_set_task_ext_regs(pid_t pid, user_fpregs_struct_t *ext_regs);
 extern int arch_fetch_sas(struct parasite_ctl *ctl, struct rt_sigframe *s);
 extern int sigreturn_prep_regs_plain(struct rt_sigframe *sigframe,
 				     user_regs_struct_t *regs,
diff --git a/compel/src/lib/infect.c b/compel/src/lib/infect.c
index d50356f6c..1bb07ace5 100644
--- a/compel/src/lib/infect.c
+++ b/compel/src/lib/infect.c
@@ -475,7 +475,8 @@ err_sig:
 	return -1;
 }
 
-static int restore_thread_ctx(int pid, struct thread_ctx *ctx)
+static int restore_thread_ctx(int pid, struct thread_ctx *ctx,
+			      bool restore_ext_regs)
 {
 	int ret = 0;
 
@@ -483,6 +484,10 @@ static int restore_thread_ctx(int pid, struct thread_ctx *ctx)
 		pr_perror("Can't restore registers (pid: %d)", pid);
 		ret = -1;
 	}
+
+	if (restore_ext_regs && compel_set_task_ext_regs(pid, &ctx->ext_regs))
+		ret = -1;
+
 	if (ptrace(PTRACE_SETSIGMASK, pid, sizeof(k_rtsigset_t), &ctx->sigmask)) {
 		pr_perror("Can't block signals");
 		ret = -1;
@@ -491,11 +496,11 @@ static int restore_thread_ctx(int pid, struct thread_ctx *ctx)
 	return ret;
 }
 
-
 /* we run at @regs->ip */
 static int parasite_trap(struct parasite_ctl *ctl, pid_t pid,
 				user_regs_struct_t *regs,
-				struct thread_ctx *octx)
+				struct thread_ctx *octx,
+				bool may_use_extended_regs)
 {
 	siginfo_t siginfo;
 	int status;
@@ -540,7 +545,7 @@ static int parasite_trap(struct parasite_ctl *ctl, pid_t pid,
 	 */
 	ret = 0;
 err:
-	if (restore_thread_ctx(pid, octx))
+	if (restore_thread_ctx(pid, octx, may_use_extended_regs))
 		ret = -1;
 
 	return ret;
@@ -567,7 +572,7 @@ int compel_execute_syscall(struct parasite_ctl *ctl,
 
 	err = parasite_run(pid, PTRACE_CONT, ctl->ictx.syscall_ip, 0, regs, &ctl->orig);
 	if (!err)
-		err = parasite_trap(ctl, pid, regs, &ctl->orig);
+		err = parasite_trap(ctl, pid, regs, &ctl->orig, false);
 
 	if (ptrace_poke_area(pid, (void *)code_orig,
 			     (void *)ctl->ictx.syscall_ip, sizeof(code_orig))) {
@@ -585,7 +590,7 @@ int compel_run_at(struct parasite_ctl *ctl, unsigned long ip, user_regs_struct_t
 
 	ret = parasite_run(ctl->rpid, PTRACE_CONT, ip, 0, &regs, &ctl->orig);
 	if (!ret)
-		ret = parasite_trap(ctl, ctl->rpid, ret_regs ? ret_regs : &regs, &ctl->orig);
+		ret = parasite_trap(ctl, ctl->rpid, ret_regs ? ret_regs : &regs, &ctl->orig, false);
 	return ret;
 }
 
@@ -1471,7 +1476,7 @@ int compel_run_in_thread(struct parasite_thread_ctl *tctl, unsigned int cmd)
 
 	ret = parasite_run(pid, PTRACE_CONT, ctl->parasite_ip, stack, &regs, octx);
 	if (ret == 0)
-		ret = parasite_trap(ctl, pid, &regs, octx);
+		ret = parasite_trap(ctl, pid, &regs, octx, true);
 	if (ret == 0)
 		ret = (int)REG_RES(regs);
 
@@ -1499,7 +1504,11 @@ int compel_unmap(struct parasite_ctl *ctl, unsigned long addr)
 	ret = compel_stop_on_syscall(1, __NR(munmap, 0),
 			__NR(munmap, 1), TRACE_ENTER);
 
-	if (restore_thread_ctx(pid, &ctl->orig))
+	/*
+	 * Don't touch extended registers here: they were restored
+	 * with rt_sigreturn from sigframe.
+	 */
+	if (restore_thread_ctx(pid, &ctl->orig, false))
 		ret = -1;
 err:
 	return ret;