Implement SNP DOWNLOAD_FIRMWARE_EX support

index d1e9e0ac63b6..22bc4ef27a63 100644
@@ -2381,7 +2381,8 @@ static int sev_ioctl_do_pdh_export(struct sev_issue_cmd *argp, bool writable)
 	return ret;
 }
 
-static int sev_ioctl_do_snp_platform_status(struct sev_issue_cmd *argp)
+static int __sev_do_snp_platform_status(struct sev_user_data_snp_status *status,
+					int *error)
 {
 	struct sev_device *sev = psp_master->sev_data;
 	struct sev_data_snp_addr buf;

@@ -2389,9 +2390,6 @@ static int sev_ioctl_do_snp_platform_status(struct sev_issue_cmd *argp)
 	void *data;
 	int ret;
 
-	if (!argp->data)
-		return -EINVAL;
-
 	status_page = alloc_page(GFP_KERNEL_ACCOUNT);
 	if (!status_page)
 		return -ENOMEM;

@@ -2414,7 +2412,7 @@ static int sev_ioctl_do_snp_platform_status(struct sev_issue_cmd *argp)
 	}
 
 	buf.address = __psp_pa(data);
-	ret = __sev_do_cmd_locked(SEV_CMD_SNP_PLATFORM_STATUS, &buf, &argp->error);
+	ret = __sev_do_cmd_locked(SEV_CMD_SNP_PLATFORM_STATUS, &buf, error);
 
 	if (sev->snp_initialized) {
 		/*

@@ -2429,15 +2427,32 @@ static int sev_ioctl_do_snp_platform_status(struct sev_issue_cmd *argp)
 	if (ret)
 		goto cleanup;
 
-	if (copy_to_user((void __user *)argp->data, data,
-			 sizeof(struct sev_user_data_snp_status)))
-		ret = -EFAULT;
+	memcpy(status, data, sizeof(*status));
 
 cleanup:
 	__free_pages(status_page, 0);
 	return ret;
 }
 
+static int sev_ioctl_do_snp_platform_status(struct sev_issue_cmd *argp)
+{
+	struct sev_user_data_snp_status status;
+	int ret;
+
+	if (!argp->data)
+		return -EINVAL;
+
+	ret = __sev_do_snp_platform_status(&status, &argp->error);
+	if (ret < 0)
+		return ret;
+
+	if (copy_to_user((void __user *)argp->data, &status,
+			 sizeof(struct sev_user_data_snp_status)))
+		ret = -EFAULT;
+
+	return ret;
+}
+
 static int sev_ioctl_do_snp_commit(struct sev_issue_cmd *argp)
 {
 	struct sev_device *sev = psp_master->sev_data;

index 22bc4ef27a63..7ca29ccda0e7 100644
@@ -132,6 +132,9 @@ static void __sev_firmware_shutdown(struct sev_device *sev, bool panic);
 static int snp_shutdown_on_panic(struct notifier_block *nb,
 				 unsigned long reason, void *arg);
 
+static int __sev_do_snp_platform_status(struct sev_user_data_snp_status *status,
+					int *error);
+
 static struct notifier_block snp_panic_notifier = {
 	.notifier_call = snp_shutdown_on_panic,
 };

@@ -1264,19 +1267,12 @@ static int snp_get_platform_data(struct sev_device *sev, int *error)
 {
 	struct sev_data_snp_feature_info snp_feat_info;
 	struct snp_feature_info *feat_info;
-	struct sev_data_snp_addr buf;
 	struct page *page;
 	int rc;
 
-	/*
-	 * This function is expected to be called before SNP is
-	 * initialized.
-	 */
-	if (sev->snp_initialized)
-		return -EINVAL;
-
-	buf.address = __psp_pa(&sev->snp_plat_status);
-	rc = sev_do_cmd(SEV_CMD_SNP_PLATFORM_STATUS, &buf, error);
+	mutex_lock(&sev_cmd_mutex);
+	rc = __sev_do_snp_platform_status(&sev->snp_plat_status, error);
+	mutex_unlock(&sev_cmd_mutex);
 	if (rc) {
 		dev_err(sev->dev, "SNP PLATFORM_STATUS command failed, ret = %d, error = %#x\n",
 			rc, *error);

@@ -1305,17 +1301,32 @@ static int snp_get_platform_data(struct sev_device *sev, int *error)
 		return -ENOMEM;
 
 	feat_info = page_address(page);
+
+	if (sev->snp_initialized) {
+		if (rmp_mark_pages_firmware(__pa(feat_info), 1, false)) {
+			rc = -EFAULT;
+			goto free_page;
+		}
+	}
+
 	snp_feat_info.length = sizeof(snp_feat_info);
 	snp_feat_info.ecx_in = 0;
 	snp_feat_info.feature_info_paddr = __psp_pa(feat_info);
 
 	rc = sev_do_cmd(SEV_CMD_SNP_FEATURE_INFO, &snp_feat_info, error);
+
+	if (sev->snp_initialized) {
+		if (snp_reclaim_pages(__pa(feat_info), 1, false))
+			return -EFAULT;
+	}
+
 	if (!rc)
 		sev->snp_feat_info_0 = *feat_info;
 	else
 		dev_err(sev->dev, "SNP FEATURE_INFO command failed, ret = %d, error = %#x\n",
 			rc, *error);
 
+free_page:
 	__free_page(page);
 
 	return rc;

index 7ca29ccda0e7..defdc1bc226e 100644
@@ -252,6 +252,7 @@ static int sev_cmd_buffer_len(int cmd)
 	case SEV_CMD_SNP_PLATFORM_STATUS:	return sizeof(struct sev_data_snp_addr);
 	case SEV_CMD_SNP_GUEST_REQUEST:		return sizeof(struct sev_data_snp_guest_request);
 	case SEV_CMD_SNP_CONFIG:		return sizeof(struct sev_user_data_snp_config);
+	case SEV_CMD_SNP_DOWNLOAD_FIRMWARE_EX:	return sizeof(struct sev_data_download_firmware_ex);
 	case SEV_CMD_SNP_COMMIT:		return sizeof(struct sev_data_snp_commit);
 	case SEV_CMD_SNP_FEATURE_INFO:		return sizeof(struct sev_data_snp_feature_info);
 	case SEV_CMD_SNP_VLEK_LOAD:		return sizeof(struct sev_user_data_snp_vlek_load);

index d5099a2baca5..5227e901bff2 100644
@@ -855,6 +855,26 @@ struct sev_platform_init_args {
 	unsigned int max_snp_asid;
 };
 
+/**
+ * struct sev_data_download_firmware_ex
+ *
+ * @len: length of the command buffer read by the PSP
+ * @rsvd0: reserved
+ * @fw_paddr: physical address of the start of the firmware blob
+ * @fw_len: length of the firmware blob
+ * @commit: whether to immediately commit the firmware update. If set, this
+ *  operation behaves like DOWNLOAD_FIRMWARE.
+ * @rsvd1: reserved
+ */
+struct sev_data_download_firmware_ex {
+	u32 len;		/* In */
+	u32 rsvd0;
+	u64 fw_paddr;		/* In */
+	u32 fw_len;		/* In */
+	u32 commit:1;		/* In */
+	u32 rsvd1:31;
+} __packed;
+
 /**
  * struct sev_data_snp_commit - SNP_COMMIT structure
  *

index defdc1bc226e..2df621b9f6e2 100644
@@ -948,6 +948,38 @@ int __sev_do_cmd_locked(int cmd, void *data, int *psp_ret)
 
 	/* wait for command completion */
 	ret = sev_wait_cmd_ioc(sev, &reg, psp_timeout);
+
+	/*
+	 * Copy potential output from the PSP back to data.  Do this even on
+	 * failure in case the caller wants to glean something from the error,
+	 * unless the operation timed out, in which case there is nothing to
+	 * copy back.
+	 */
+	if (data) {
+		int ret_reclaim;
+		/*
+		 * Restore the page state after the command completes.
+		 */
+		ret_reclaim = snp_reclaim_cmd_buf(cmd, cmd_buf);
+		if (ret_reclaim) {
+			dev_err(sev->dev,
+				"SEV: failed to reclaim buffer for legacy command %#x. Error: %d\n",
+				cmd, ret_reclaim);
+			return ret_reclaim;
+		}
+
+		if (ret != -ETIMEDOUT)
+			memcpy(data, cmd_buf, buf_len);
+
+		if (sev->cmd_buf_backup_active)
+			sev->cmd_buf_backup_active = false;
+		else
+			sev->cmd_buf_active = false;
+
+		if (snp_unmap_cmd_buf_desc_list(desc_list))
+			return -EFAULT;
+	}
+
 	if (ret) {
 		if (psp_ret)
 			*psp_ret = 0;

@@ -984,34 +1016,6 @@ int __sev_do_cmd_locked(int cmd, void *data, int *psp_ret)
 		ret = sev_write_init_ex_file_if_required(cmd);
 	}
 
-	/*
-	 * Copy potential output from the PSP back to data.  Do this even on
-	 * failure in case the caller wants to glean something from the error.
-	 */
-	if (data) {
-		int ret_reclaim;
-		/*
-		 * Restore the page state after the command completes.
-		 */
-		ret_reclaim = snp_reclaim_cmd_buf(cmd, cmd_buf);
-		if (ret_reclaim) {
-			dev_err(sev->dev,
-				"SEV: failed to reclaim buffer for legacy command %#x. Error: %d\n",
-				cmd, ret_reclaim);
-			return ret_reclaim;
-		}
-
-		memcpy(data, cmd_buf, buf_len);
-
-		if (sev->cmd_buf_backup_active)
-			sev->cmd_buf_backup_active = false;
-		else
-			sev->cmd_buf_active = false;
-
-		if (snp_unmap_cmd_buf_desc_list(desc_list))
-			return -EFAULT;
-	}
-
 	print_hex_dump_debug("(out): ", DUMP_PREFIX_OFFSET, 16, 2, data,
 			     buf_len, false);
 

index 2df621b9f6e2..b4711bf823e8 100644
@@ -2040,6 +2040,53 @@ static int sev_update_firmware(struct device *dev)
 	return ret;
 }
 
+static enum fw_upload_err sev_fw_upload_prepare(struct fw_upload *fw_upload,
+						const u8 *data, u32 size)
+{
+	return FW_UPLOAD_ERR_NONE;
+}
+
+static enum fw_upload_err sev_fw_upload_write(struct fw_upload *fw_upload,
+					      const u8 *data, u32 offset,
+					      u32 size, u32 *written)
+{
+	return FW_UPLOAD_ERR_BUSY;
+}
+
+static enum fw_upload_err sev_fw_upload_poll_complete(struct fw_upload *fw_upload)
+{
+	return FW_UPLOAD_ERR_NONE;
+}
+
+static void sev_fw_upload_cancel(struct fw_upload *fw_upload)
+{
+	/* intentional no-op */
+}
+
+static const struct fw_upload_ops sev_fw_upload_ops = {
+	.prepare = sev_fw_upload_prepare,
+	.write = sev_fw_upload_write,
+	.poll_complete = sev_fw_upload_poll_complete,
+	.cancel = sev_fw_upload_cancel,
+};
+
+static void register_sev_fw_uploader(struct sev_device *sev)
+{
+	struct fw_upload *fwl;
+
+	if (!IS_ENABLED(CONFIG_FW_UPLOAD))
+		return;
+
+	fwl = firmware_upload_register(THIS_MODULE, sev->dev, "sev",
+				       &sev_fw_upload_ops, sev);
+	if (IS_ERR(fwl)) {
+		dev_err(sev->dev, "SEV firmware upload registration failure: %ld\n", PTR_ERR(fwl));
+		return;
+	}
+
+	sev->fwl = fwl;
+}
+
 static int __sev_snp_shutdown_locked(int *error, bool panic)
 {
 	struct psp_device *psp = psp_master;

@@ -2953,6 +3000,7 @@ void sev_pci_init(void)
 			 api_major, api_minor, build,
 			 sev->api_major, sev->api_minor, sev->build);
 
+	register_sev_fw_uploader(sev);
 	return;
 
 err:

@@ -2969,4 +3017,7 @@ void sev_pci_exit(void)
 		return;
 
 	sev_firmware_shutdown(sev);
+
+	if (sev->fwl)
+		firmware_upload_unregister(sev->fwl);
 }

index b1cd556bbbf6..2f5781cb7bb1 100644
@@ -24,6 +24,7 @@
 #include <linux/psp-sev.h>
 #include <linux/miscdevice.h>
 #include <linux/capability.h>
+#include <linux/firmware.h>
 
 #define SEV_CMDRESP_CMD			GENMASK(26, 16)
 #define SEV_CMD_COMPLETE		BIT(1)

@@ -66,6 +67,8 @@ struct sev_device {
 
 	struct tsm_dev *tsmdev;
 	struct sev_tio_status *tio_status;
+
+	struct fw_upload *fwl;
 };
 
 int sev_dev_init(struct psp_device *psp);

index b4711bf823e8..e7fe6dbf69c2 100644
@@ -90,6 +90,8 @@ MODULE_FIRMWARE("amd/amd_sev_fam19h_model1xh.sbin"); /* 4th gen EPYC */
 
 static bool psp_dead;
 static int psp_timeout;
+static bool dlfwex_wants_rollback;
+static bool sev_firmware_needs_reinit;
 
 enum snp_hv_fixed_pages_state {
 	ALLOCATED,

@@ -2046,11 +2048,251 @@ static enum fw_upload_err sev_fw_upload_prepare(struct fw_upload *fw_upload,
 	return FW_UPLOAD_ERR_NONE;
 }
 
+static int sev_download_firmware_ex(struct sev_device *sev, const u8 *data,
+				    u32 size)
+{
+	struct sev_data_download_firmware_ex sev_data = {0};
+	int ret, error = 0, order;
+	struct page *p;
+	void *fw_blob;
+
+	order = get_order(size);
+	p = alloc_pages(GFP_KERNEL, order);
+	if (!p)
+		return -ENOMEM;
+
+	fw_blob = page_address(p);
+	memcpy(fw_blob, data, size);
+
+	sev_data.len = sizeof(sev_data);
+	sev_data.fw_paddr = __psp_pa(fw_blob);
+	sev_data.fw_len = size;
+	sev_data.commit = 0;
+
+	ret = __sev_do_cmd_locked(SEV_CMD_SNP_DOWNLOAD_FIRMWARE_EX, &sev_data, &error);
+
+	/*
+	 * Quirk: firmware upgrades across 1.58.03 give ETIMEDOUT for
+	 * DLFWEX, even though the command actually succeeds. If we're
+	 * in this case, test that we can do SNP_PLATFORM_STATUS, and
+	 * if so, continue as normal.
+	 */
+	if (ret == -ETIMEDOUT) {
+		struct sev_user_data_snp_status status;
+
+		dev_info(sev->dev, "Firmware update timed out, checking status for quirk...\n");
+		psp_dead = false;
+
+		ret = __sev_do_snp_platform_status(&status, &error);
+		if (ret) {
+			dev_err(sev->dev, "SNP STATUS failed after firmware upgrade, ret = %d, error = %#x\n",
+				ret, error);
+			psp_dead = true;
+			goto out;
+		}
+	}
+
+	if (ret < 0 && error != 0)
+		ret = error;
+
+out:
+	__free_pages(p, order);
+	return ret;
+}
+
+static int sev_firmware_shutdown_if_sev_initialized(struct sev_device *sev)
+{
+	int rc, error, sev_plat_state;
+
+	lockdep_assert_held(&sev_cmd_mutex);
+
+	error = 0;
+	rc = sev_get_platform_state(&sev_plat_state, &error);
+	if (rc < 0) {
+		if (error)
+			rc = error;
+		dev_dbg(sev->dev, "SEV get platform state failed %d\n", rc);
+		return rc;
+	}
+
+	switch (sev_plat_state) {
+	case SEV_STATE_UNINIT:
+		return 0;
+	case SEV_STATE_INIT:
+		error = 0;
+		rc = __sev_platform_shutdown_locked(&error);
+		if (rc) {
+			if (error)
+				rc = error;
+			dev_err(sev->dev, "SEV platform shutdown failed %d\n", rc);
+			return rc;
+		}
+		sev_firmware_needs_reinit = true;
+		return 0;
+	case SEV_STATE_WORKING:
+		return -EBUSY;
+	default:
+		dev_err(sev->dev, "Unknown SEV firmware state: %d\n", sev_plat_state);
+		return -EINVAL;
+	}
+}
+
+static void sev_firmware_reinit_if_shutdown(struct sev_device *sev)
+{
+	int rc, error;
+
+	guard(mutex)(&sev_cmd_mutex);
+
+	if (!sev_firmware_needs_reinit)
+		return;
+
+	sev_firmware_needs_reinit = false;
+	error = 0;
+	rc = __sev_platform_init_locked(&error);
+	if (rc) {
+		if (error)
+			rc = error;
+		dev_err(sev->dev, "SEV platform re-init failed %d\n", rc);
+	}
+}
+
 static enum fw_upload_err sev_fw_upload_write(struct fw_upload *fw_upload,
 					      const u8 *data, u32 offset,
 					      u32 size, u32 *written)
 {
-	return FW_UPLOAD_ERR_BUSY;
+	struct sev_device *sev = fw_upload->dd_handle;
+	u8 old_major, old_minor, old_build;
+	int rc, error = 0;
+	enum fw_upload_err ret;
+
+	if (offset != 0)
+		return FW_UPLOAD_ERR_INVALID_SIZE;
+
+	old_major = sev->api_major;
+	old_minor = sev->api_minor;
+	old_build = sev->build;
+
+	mutex_lock(&sev_cmd_mutex);
+
+	/*
+	 * If the last firmware update returned RESTORE_REQUIRED, allow only
+	 * this DLFWEX command so the admin can restore the previous FW
+	 * version. If we are in this state the legacy firmware has previously
+	 * been shut down, so no need to do it again.
+	 */
+	if (dlfwex_wants_rollback && psp_dead) {
+		dlfwex_wants_rollback = false;
+		psp_dead = false;
+	} else {
+		rc = sev_firmware_shutdown_if_sev_initialized(sev);
+		if (rc) {
+			ret = FW_UPLOAD_ERR_BUSY;
+			goto unlock;
+		}
+	}
+
+	rc = sev_download_firmware_ex(sev, data, size);
+	if (rc) {
+		ret = FW_UPLOAD_ERR_FW_INVALID;
+		switch (rc) {
+		case SEV_RET_INVALID_PLATFORM_STATE:
+			fallthrough;
+		case SEV_RET_INVALID_ADDRESS:
+			/* these are probably kernel bugs */
+			WARN_ON_ONCE(true);
+			ret = FW_UPLOAD_ERR_BUSY;
+			goto unlock;
+		case SEV_RET_INVALID_LEN:
+			ret = FW_UPLOAD_ERR_INVALID_SIZE;
+			goto unlock;
+		case SEV_RET_INVALID_PARAM:
+			dev_err(sev->dev, "SEV firmware image is not well formed\n");
+			goto unlock;
+		case SEV_RET_SHUTDOWN_REQUIRED:
+			dev_err(sev->dev, "SEV firmware too far, shutdown required\n");
+			goto unlock;
+		case SEV_RET_INVALID_CONFIG:
+			dev_err(sev->dev, "SEV firmware upgrade would rollback SVN\n");
+			goto unlock;
+		case SEV_RET_BAD_SIGNATURE:
+			dev_err(sev->dev, "SEV firmware upgrade bad signature\n");
+			goto unlock;
+		case SEV_RET_BAD_VERSION:
+			dev_err(sev->dev, "SEV firmware upgrade less than CommittedVersion\n");
+			goto unlock;
+		case SEV_RET_UNSUPPORTED:
+			dev_err(sev->dev, "SEV firmware required feature not supported\n");
+			goto unlock;
+		case SEV_RET_UPDATE_FAILED:
+			/*
+			 * Update failed but fw rolled back on its own,
+			 * operation can continue normally.
+			 */
+			dev_err(sev->dev, "SEV firmware update failed\n");
+			ret = FW_UPLOAD_ERR_HW_ERROR;
+			goto unlock;
+		case SEV_RET_HWSEV_RET_UNSAFE:
+			/*
+			 * "Following a return of HARDWARE_UNSAFE, operation of
+			 * the SEV firmware is indeterminate
+			 * and the recommendation is to reboot the platform."
+			 */
+			dev_err(sev->dev, "SEV firmware no longer safe to operate\n");
+			psp_dead = true;
+			ret = FW_UPLOAD_ERR_HW_ERROR;
+			goto unlock;
+		case SEV_RET_RESTORE_REQUIRED:
+			/*
+			 * FW asked us to roll back; we don't hold onto the
+			 * last FW image, so we can't. We can set a flag to
+			 * allow the admin to rollback if they happen to have
+			 * the old firmware image handy.
+			 */
+			dev_err(sev->dev, "SEV firmware update failed, please roll back\n");
+			psp_dead = true;
+			dlfwex_wants_rollback = true;
+			ret = FW_UPLOAD_ERR_HW_ERROR;
+			goto unlock;
+		default:
+			dev_err(sev->dev, "Unknown SEV firmware err %d\n", rc);
+			ret = FW_UPLOAD_ERR_HW_ERROR;
+			goto unlock;
+		}
+	}
+
+	*written = size;
+	ret = FW_UPLOAD_ERR_NONE;
+
+unlock:
+	mutex_unlock(&sev_cmd_mutex);
+
+	/*
+	 * sev_get_api_version() updates the SEV and SNP statuses, SNP feature
+	 * info if available, build numbers, etc. cached in struct sev_device.
+	 * Update these if they may have changed for new firmware.
+	 */
+	if (ret == FW_UPLOAD_ERR_NONE) {
+		error = 0;
+
+		rc = sev_get_api_version();
+		if (rc) {
+			dev_warn(sev->dev,
+				 "SNP platform data refresh after firmware update failed %d\n",
+				 rc);
+		} else if (sev->api_major != old_major ||
+			   sev->api_minor != old_minor ||
+			   sev->build != old_build) {
+			dev_info(sev->dev, "SEV firmware updated to %d.%d build %d\n",
+				 sev->api_major, sev->api_minor, sev->build);
+		} else {
+			dev_info(sev->dev, "SEV firmware not updated\n");
+		}
+	}
+
+	if (!dlfwex_wants_rollback)
+		sev_firmware_reinit_if_shutdown(sev);
+
+	return ret;
 }
 
 static enum fw_upload_err sev_fw_upload_poll_complete(struct fw_upload *fw_upload)