@@ -564,7 +564,6 @@ static inline void switch_fpu_prepare(st
 static inline void switch_fpu_finish(struct fpu *new_fpu)
 {
 	u32 pkru_val = init_pkru_value;
-	struct pkru_state *pk;
 
 	if (!static_cpu_has(X86_FEATURE_FPU))
 		return;

@@ -578,11 +577,8 @@ static inline void switch_fpu_finish(str
 	 * PKRU state is switched eagerly because it needs to be valid before we
 	 * return to userland e.g. for a copy_to_user() operation.
 	 */
-	if (current->mm) {
-		pk = get_xsave_addr(&new_fpu->state.xsave, XFEATURE_PKRU);
-		if (pk)
-			pkru_val = pk->pkru;
-	}
+	if (current->mm)
+		pkru_val = xstate_read_pkru(&new_fpu->state.xsave);
 	__write_pkru(pkru_val);
 
 	/*

@@ -124,27 +124,124 @@ static inline u32 read_pkru(void)
 	return 0;
 }
 
+static inline void xfeature_mark_non_init(struct xregs_state *xstate,
+					  int xfeature_nr)
+{
+	/*
+	 * Caller will place data in the @xstate buffer.
+	 * Mark the xfeature as non-init:
+	 */
+	xstate->header.xfeatures |= BIT_ULL(xfeature_nr);
+}
+
+
+/* Set the contents of @xfeature_nr to the hardware init state */
+static inline void xfeature_init_space(struct xregs_state *xstate,
+					     int xfeature_nr)
+{
+	void *state = get_xsave_addr(xstate, xfeature_nr);
+
+	switch (xfeature_nr) {
+	case XFEATURE_PKRU:
+		/* zero the whole state, including reserved bits */
+		memset(state, 0, sizeof(struct pkru_state));
+		break;
+	default:
+		BUG();
+	}
+}
+
+/*
+ * Called when it is necessary to write to an XSAVE
+ * component feature.  Guarantees that a future
+ * XRSTOR of the 'xstate' buffer will not consider
+ * @xfeature_nr to be in its init state.
+ *
+ * The returned buffer may contain old state.  The
+ * caller must be prepared to fill the entire buffer.
+ *
+ * Caller must first ensure that @xfeature_nr is
+ * enabled and present in the @xstate buffer.
+ */
+static inline void *xstatebuf_get_write_ptr(struct xregs_state *xstate,
+					    int xfeature_nr)
+{
+	bool feature_was_init = xstate->header.xfeatures & BIT_ULL(xfeature_nr);
+
+	/*
+	 * xcomp_bv represents whether 'xstate' has space for
+	 * features.  If not, something is horribly wrong and
+	 * a write would corrupt memory.  Perhaps xfeature_nr
+	 * was not enabled.
+	 */
+	BUG_ON(!(xstate->header.xcomp_bv & BIT_ULL(xfeature_nr)));
+
+	/*
+	 * Ensure a sane xfeature_nr, including avoiding
+	 * confusion with XCOMP_BV_COMPACTED_FORMAT.
+	 */
+	BUG_ON(xfeature_nr >= XFEATURE_MAX);
+
+	/* Prepare xstate for a write to the xfeature: */
+	xfeature_mark_non_init(xstate, xfeature_nr);
+
+	/*
+	 * If xfeature_nr was in the init state, update the buffer
+	 * to match the state. Ensures that callers can safely
+	 * write only a part of the state, they are not forced to
+	 * write it in its entirety.
+	 */
+	if (feature_was_init)
+		xfeature_init_space(xstate, xfeature_nr);
+
+	return get_xsave_addr(xstate, xfeature_nr);
+}
+
+/* Caller must ensure X86_FEATURE_OSPKE is enabled. */
+static inline void xstate_write_pkru(struct xregs_state *xstate, u32 pkru)
+{
+	struct pkru_state *pk;
+
+	pk = xstatebuf_get_write_ptr(xstate, XFEATURE_PKRU);
+	pk->pkru = pkru;
+}
+
+/*
+ * What PKRU value is represented in the 'xstate'?  Note,
+ * this returns the *architecturally* represented value,
+ * not the literal in-memory value.  They may be different.
+ */
+static inline u32 xstate_read_pkru(struct xregs_state *xstate)
+{
+	struct pkru_state *pk;
+
+	pk = get_xsave_addr(xstate, XFEATURE_PKRU);
+	/*
+	 * If present, pull PKRU out of the XSAVE buffer.
+	 * Otherwise, use the hardware init value.
+	 */
+	if (pk)
+		return pk->pkru;
+	else
+		return PKRU_HW_INIT_VALUE;
+}
+
 /*
  * Update all of the PKRU state for the current task:
  * PKRU register and PKRU xstate.
  */
 static inline void current_write_pkru(u32 pkru)
 {
-	struct pkru_state *pk;
-
 	if (!boot_cpu_has(X86_FEATURE_OSPKE))
 		return;
 
-	pk = get_xsave_addr(&current->thread.fpu.state.xsave, XFEATURE_PKRU);
-
+	fpregs_lock();
 	/*
 	 * The PKRU value in xstate needs to be in sync with the value that is
 	 * written to the CPU. The FPU restore on return to userland would
 	 * otherwise load the previous value again.
 	 */
-	fpregs_lock();
-	if (pk)
-		pk->pkru = pkru;
+	xstate_write_pkru(&current->thread.fpu.state.xsave, pkru);
 	__write_pkru(pkru);
 	fpregs_unlock();
 }

@@ -854,4 +854,11 @@ enum mds_mitigations {
 	MDS_MITIGATION_VMWERV,
 };
 
+/*
+ * The XSAVE architecture defines an "init state" for
+ * PKRU.  PKRU is set to this value by XRSTOR when it
+ * tries to restore PKRU but has on value in the buffer.
+ */
+#define PKRU_HW_INIT_VALUE	0x0
+
 #endif /* _ASM_X86_PROCESSOR_H */

@@ -466,8 +466,6 @@ static bool pku_disabled;
 
 static __always_inline void setup_pku(struct cpuinfo_x86 *c)
 {
-	struct pkru_state *pk;
-
 	/* check the boot processor, plus compile options for PKU: */
 	if (!cpu_feature_enabled(X86_FEATURE_PKU))
 		return;

@@ -478,9 +476,7 @@ static __always_inline void setup_pku(st
 		return;
 
 	cr4_set_bits(X86_CR4_PKE);
-	pk = get_xsave_addr(&init_fpstate.xsave, XFEATURE_PKRU);
-	if (pk)
-		pk->pkru = init_pkru_value;
+	xstate_write_pkru(&current->thread.fpu.state.xsave, init_pkru_value);
 	/*
 	 * Seting X86_CR4_PKE will cause the X86_FEATURE_OSPKE
 	 * cpuid bit to be set.  We need to ensure that we

@@ -155,7 +155,6 @@ static ssize_t init_pkru_read_file(struc
 static ssize_t init_pkru_write_file(struct file *file,
 		 const char __user *user_buf, size_t count, loff_t *ppos)
 {
-	struct pkru_state *pk;
 	char buf[32];
 	ssize_t len;
 	u32 new_init_pkru;

@@ -178,10 +177,7 @@ static ssize_t init_pkru_write_file(stru
 		return -EINVAL;
 
 	WRITE_ONCE(init_pkru_value, new_init_pkru);
-	pk = get_xsave_addr(&init_fpstate.xsave, XFEATURE_PKRU);
-	if (!pk)
-		return -EINVAL;
-	pk->pkru = new_init_pkru;
+	xstate_write_pkru(&init_fpstate.xsave, new_init_pkru);
 	return count;
 }