compel: shstk: save CET state when CPU supports it

Signed-off-by: Mike Rapoport (IBM) <rppt@kernel.org>

compel/arch/x86/src/lib/include/uapi/asm/cpu.h

@@ -244,6 +244,7 @@ enum cpuid_leafs {
 #define X86_FEATURE_PKU		     (11 * 32 + 3)  /* Protection Keys for Userspace */
 #define X86_FEATURE_OSPKE	     (11 * 32 + 4)  /* OS Protection Keys Enable */
 #define X86_FEATURE_AVX512_VBMI2     (11 * 32 + 6)  /* Additional AVX512 Vector Bit Manipulation Instructions */
+#define X86_FEATURE_SHSTK            (11 * 32 + 7)  /* Shadow Stack */
 #define X86_FEATURE_GFNI	     (11 * 32 + 8)  /* Galois Field New Instructions */
 #define X86_FEATURE_VAES	     (11 * 32 + 9)  /* Vector AES */
 #define X86_FEATURE_VPCLMULQDQ	     (11 * 32 + 10) /* Carry-Less Multiplication Double Quadword */

compel/arch/x86/src/lib/include/uapi/asm/fpu.h

@@ -245,6 +245,14 @@ struct pkru_state {
 	uint32_t pad;
 } __packed;
 
+/*
+ * State component 11 is Control-flow Enforcement user states
+ */
+struct cet_user_state {
+	uint64_t cet;			/* user control-flow settings */
+	uint64_t ssp;			/* user shadow stack pointer */
+};
+
 /*
  * This is our most modern FPU state format, as saved by the XSAVE
  * and restored by the XRSTOR instructions.
@@ -260,7 +268,7 @@ struct pkru_state {
  * Of course it was not ;-) Now using four pages...
  *
  */
-#define EXTENDED_STATE_AREA_SIZE (XSAVE_SIZE - sizeof(struct i387_fxsave_struct) - sizeof(struct xsave_hdr_struct))
+#define EXTENDED_STATE_AREA_SIZE (XSAVE_SIZE - sizeof(struct i387_fxsave_struct) - sizeof(struct xsave_hdr_struct) - sizeof(struct cet_user_state))
 
 /*
  * cpu requires it to be 64 byte aligned
@@ -276,6 +284,7 @@ struct xsave_struct {
 		struct ymmh_struct ymmh;
 		uint8_t extended_state_area[EXTENDED_STATE_AREA_SIZE];
 	};
+	struct cet_user_state cet;
 } __aligned(FP_MIN_ALIGN_BYTES) __packed;
 
 struct xsave_struct_ia32 {

compel/arch/x86/src/lib/include/uapi/asm/infect-types.h

@@ -143,4 +143,7 @@ typedef struct xsave_struct user_fpregs_struct_t;
  */
 #define __NR32_mmap __NR32_mmap2
 
+extern bool __compel_shstk_enabled(user_fpregs_struct_t *ext_regs);
+#define compel_shstk_enabled __compel_shstk_enabled
+
 #endif /* UAPI_COMPEL_ASM_TYPES_H__ */

compel/arch/x86/src/lib/infect.c

@@ -26,6 +26,16 @@
 #ifndef NT_X86_XSTATE
 #define NT_X86_XSTATE 0x202 /* x86 extended state using xsave */
 #endif
+
+#ifndef NT_X86_SHSTK
+#define NT_X86_SHSTK 0x204	/* x86 shstk state */
+#endif
+
+#ifndef ARCH_SHSTK_STATUS
+#define ARCH_SHSTK_STATUS	0x5005
+#define ARCH_SHSTK_SHSTK	(1ULL << 0)
+#endif
+
 #ifndef NT_PRSTATUS
 #define NT_PRSTATUS 1 /* Contains copy of prstatus struct */
 #endif
@@ -250,7 +260,49 @@ static int get_task_xsave(pid_t pid, user_fpregs_struct_t *xsave)
 		// [1] Intel® 64 and IA-32 Architectures Software Developer's
 		//     Manual Volume 1: Basic Architecture
 		//     Section 13.6: Processor tracking of XSAVE-managed state
-		return get_task_fpregs(pid, xsave);
+		if (get_task_fpregs(pid, xsave))
+			return -1;
+	}
+
+	/*
+	 * xsave may be on stack, if we don't clear it explicitly we get
+	 * funky shadow stack state
+	 */
+	memset(&xsave->cet, 0, sizeof(xsave->cet));
+	if (compel_cpu_has_feature(X86_FEATURE_SHSTK)) {
+		unsigned long ssp = 0;
+		unsigned long features = 0;
+
+		if (ptrace(PTRACE_ARCH_PRCTL, pid, (unsigned long)&features, ARCH_SHSTK_STATUS)) {
+			/*
+			 * kernels that don't support shadow stack return
+			 * -EINVAL
+			 */
+			if (errno == EINVAL)
+				return 0;
+
+			pr_perror("shstk: can't get shadow stack status for %d", pid);
+			return -1;
+		}
+
+		if (!(features & ARCH_SHSTK_SHSTK))
+			return 0;
+
+		iov.iov_base = &ssp;
+		iov.iov_len = sizeof(ssp);
+
+		if (ptrace(PTRACE_GETREGSET, pid, (unsigned int)NT_X86_SHSTK, &iov) < 0) {
+			/* ENODEV means CET is not supported by the CPU  */
+			if (errno != ENODEV) {
+				pr_perror("shstk: can't get SSP for %d", pid);
+				return -1;
+			}
+		}
+
+		xsave->cet.cet = features;
+		xsave->cet.ssp = ssp;
+
+		pr_debug("%d: shstk: cet: %lx ssp: %lx\n", pid, xsave->cet.cet, xsave->cet.ssp);
 	}
 
 	return 0;
@@ -697,3 +749,14 @@ unsigned long compel_task_size(void)
 {
 	return TASK_SIZE;
 }
+
+bool __compel_shstk_enabled(user_fpregs_struct_t *ext_regs)
+{
+	if (!compel_cpu_has_feature(X86_FEATURE_SHSTK))
+		return false;
+
+	if (ext_regs->cet.cet & ARCH_SHSTK_SHSTK)
+		return true;
+
+	return false;
+}

compel/include/uapi/infect.h

@@ -182,4 +182,12 @@ void compel_set_thread_ip(struct parasite_thread_ctl *tctl, uint64_t v);
 
 extern void compel_get_stack(struct parasite_ctl *ctl, void **rstack, void **r_thread_stack);
 
+#ifndef compel_shstk_enabled
+static inline bool compel_shstk_enabled(user_fpregs_struct_t *ext_regs)
+{
+	return false;
+}
+#define compel_shstk_enabled
+#endif
+
 #endif

criu/arch/x86/crtools.c

@@ -133,6 +133,14 @@ int save_task_regs(void *x, user_regs_struct_t *regs, user_fpregs_struct_t *fpre
 #undef assign_array
 #undef assign_xsave
 
+	if (compel_cpu_has_feature(X86_FEATURE_SHSTK)) {
+		UserX86CetEntry *cet = core->thread_info->fpregs->xsave->cet;
+		struct cet_user_state *regs = &fpregs->cet;
+
+		cet->cet = regs->cet;
+		cet->ssp = regs->ssp;
+	}
+
 	return 0;
 }
 
@@ -199,6 +207,13 @@ static int alloc_xsave_extends(UserX86XsaveEntry *xsave)
 			goto err;
 	}
 
+	if (compel_cpu_has_feature(X86_FEATURE_SHSTK)) {
+		xsave->cet = xzalloc(sizeof(UserX86CetEntry));
+		if (!xsave->cet)
+			goto err;
+		user_x86_cet_entry__init(xsave->cet);
+	}
+
 	return 0;
 err:
 	return -1;
@@ -220,6 +235,8 @@ int arch_alloc_thread_info(CoreEntry *core)
 		with_xsave = compel_cpu_has_feature(X86_FEATURE_OSXSAVE);
 		if (with_xsave)
 			sz += sizeof(UserX86XsaveEntry);
+		if (compel_cpu_has_feature(X86_FEATURE_SHSTK))
+			sz += sizeof(UserX86CetEntry);
 	}
 
 	m = xmalloc(sz);

images/core-x86.proto

@@ -41,6 +41,11 @@ message user_x86_regs_entry {
 	optional user_x86_regs_mode	mode		= 28 [default = NATIVE];
 }
 
+message user_x86_cet_entry {
+	required uint64			cet		=  1[(criu).hex = true];
+	required uint64			ssp		=  2[(criu).hex = true];
+}
+
 message user_x86_xsave_entry {
 	/* standard xsave features */
 	required uint64			xstate_bv	=  1;
@@ -60,6 +65,9 @@ message user_x86_xsave_entry {
 	/* Protected keys */
 	repeated uint32			pkru		=  8;
 
+	/* CET */
+	optional user_x86_cet_entry	cet		=  9;
+
 	/*
 	 * Processor trace (PT) and hardware duty cycling (HDC)
 	 * are supervisor state components and only managed by