And specify the memory to be decrypted when running in an SEV guest
so that the VMM can access the memory successfully.

Signed-off-by: Suravee Suthikulpanit <suravee.suthikulpanit@amd.com>
---
 drivers/iommu/amd/amd_iommu_types.h |  2 +-
 drivers/iommu/amd/init.c            | 17 +++++++++++------
 drivers/iommu/amd/iommu.c           |  2 +-
 3 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/drivers/iommu/amd/amd_iommu_types.h b/drivers/iommu/amd/amd_iommu_types.h
index d9159f2e3f0f..653955ab120d 100644
--- a/drivers/iommu/amd/amd_iommu_types.h
+++ b/drivers/iommu/amd/amd_iommu_types.h
@@ -733,7 +733,7 @@ struct amd_iommu {
 	u64 exclusion_length;
 
 	/* command buffer virtual address */
-	u8 *cmd_buf;
+	struct amd_iommu_mem cmd_buf_mem;
 	u32 cmd_buf_head;
 	u32 cmd_buf_tail;
 
diff --git a/drivers/iommu/amd/init.c b/drivers/iommu/amd/init.c
index c68ff602d534..77147dc3b79f 100644
--- a/drivers/iommu/amd/init.c
+++ b/drivers/iommu/amd/init.c
@@ -735,10 +735,15 @@ static void __init free_alias_table(struct amd_iommu_pci_seg *pci_seg)
  */
 static int __init alloc_command_buffer(struct amd_iommu *iommu)
 {
-	iommu->cmd_buf = iommu_alloc_pages(GFP_KERNEL,
-					   get_order(CMD_BUFFER_SIZE));
+	struct amd_iommu_mem *mem = &iommu->cmd_buf_mem;
 
-	return iommu->cmd_buf ? 0 : -ENOMEM;
+	mem->modes = ALLOC_MODE_GUEST_MEM_DECRYPT;
+	mem->order = get_order(CMD_BUFFER_SIZE);
+	mem->buf = amd_iommu_get_zeroed_mem(GFP_KERNEL, mem);
+	if (!mem->buf)
+		return -ENOMEM;
+
+	return 0;
 }
 
 /*
@@ -812,9 +817,9 @@ static void iommu_enable_command_buffer(struct amd_iommu *iommu)
 {
 	u64 entry;
 
-	BUG_ON(iommu->cmd_buf == NULL);
+	BUG_ON(iommu->cmd_buf_mem.buf == NULL);
 
-	entry = iommu_virt_to_phys(iommu->cmd_buf);
+	entry = amd_iommu_mem_to_phys(&iommu->cmd_buf_mem);
 	entry |= MMIO_CMD_SIZE_512;
 
 	memcpy_toio(iommu->mmio_base + MMIO_CMD_BUF_OFFSET,
@@ -833,7 +838,7 @@ static void iommu_disable_command_buffer(struct amd_iommu *iommu)
 
 static void __init free_command_buffer(struct amd_iommu *iommu)
 {
-	iommu_free_pages(iommu->cmd_buf, get_order(CMD_BUFFER_SIZE));
+	amd_iommu_free_mem(&iommu->cmd_buf_mem);
 }
 
 void *__init iommu_alloc_4k_pages(struct amd_iommu *iommu, gfp_t gfp,
diff --git a/drivers/iommu/amd/iommu.c b/drivers/iommu/amd/iommu.c
index 3df07e8ef002..2b18134f1eb5 100644
--- a/drivers/iommu/amd/iommu.c
+++ b/drivers/iommu/amd/iommu.c
@@ -1078,7 +1078,7 @@ static void copy_cmd_to_buffer(struct amd_iommu *iommu,
 
 	/* Copy command to buffer */
 	tail = iommu->cmd_buf_tail;
-	target = iommu->cmd_buf + tail;
+	target = ((u8 *)iommu->cmd_buf_mem.buf) + tail;
 	memcpy(target, cmd, sizeof(*cmd));
 
 	tail = (tail + sizeof(*cmd)) % CMD_BUFFER_SIZE;
-- 
2.34.1