PCI/TSM: coco/sev-guest: Implement SEV-TIO PCIe TDISP (phase2)

drivers/pci/pci-driver.c static void pci_dma_cleanup(struct device *dev) · line 1743

Quoted 1 line · Alexey Kardashevskiy

drivers/pci/pci-driver.c +1 -0

+	if (device_tcb_trusted(dev))

Where is that function and the rest of the gunk needed for me to apply before I review those?

The patches on that branch:

https://github.com/AMDESE/linux-kvm/commits/tsm-next

which come before those here are enough?

Correct, the patches between v7.3-rc2 and 01/17 (can drop samples and tools, obviously) are enough. It is basically recent tagged upstream + rebased Dan's https://lore.kernel.org/r/[email protected]

The PCIe spec r6 and later defines the MMIO range start in 4K units which was not the intention. The upcoming change makes it a byte address. The structure is already fixed to match the new definition but the comment is stale, fix it.

Signed-off-by Alexey Kardashevskiy <[email protected]>

---
 drivers/pci/tsm/core.c | 1 -
 1 file changed, 1 deletion(-)

drivers/pci/tsm/core.c +0 -1

index 9ac216ad896d..2a18be5be56e 100644
@@ -623,7 +623,6 @@ EXPORT_SYMBOL_GPL(pci_tsm_mmio_teardown);
 #define PCI_TSM_DEVIF_REPORT_MMIO_ATTR_IS_UPDATABLE BIT(3)
 #define PCI_TSM_DEVIF_REPORT_MMIO_ATTR_RANGE_ID GENMASK(31, 16)
 
-/* An interface report 'pfn' is 4K in size */
 struct pci_tsm_devif_mmio {
 	__le64 phys;
 	__le32 nr_pfns;

Signature · 1 line

2.55.0

(quoted diff) +34 -0

+static int alloc_iommu_tlb_flush_ghcb_pages(void)
+{
+       unsigned int cpu;
+       struct page *pg;
+       void *p;
+
+       /*
+        * Allocate per CPU pages while encrypted DMA is not happening yet
+        * and smashing is cheap.
+        */
+       for_each_possible_cpu(cpu) {
+               if (per_cpu(iommu_tlb_flush_ghcb_page, cpu))
+                       continue;
+
+               pg = alloc_pages_node(cpu_to_node(cpu), GFP_KERNEL, 0);
+               if (!pg)
+                       return -ENOMEM;
+
+               p = page_to_virt(pg);
+               /* Trigger psmash in the host os now to avoid psmash race later */
+               snp_set_memory_shared((unsigned long)p, 1);
+               snp_set_memory_private((unsigned long)p, 1);
+               per_cpu(iommu_tlb_flush_ghcb_page, cpu) = p;
+       }
+
+       return 0;
+}
 
 int sev_tio_op(u32 guest_rid, unsigned int op, u64 *fw_err, u64 *tdi_id)
 {

@@ -111,6 +142,24 @@ int sev_tio_op(u32 guest_rid, unsigned int op, u64 *fw_err, u64 *tdi_id)
        struct ghcb *ghcb;
        int ret;
 
+       if (!(sev_hv_features & GHCB_HV_FT_SNP_SEV_TIO))
+               return -EPERM;
+
+       if (op == SVM_VMGEXIT_SEV_TIO_OP_RUN || op == SVM_VMGEXIT_SEV_TIO_OP_STOP) {
+               if (!(sev_hv_features & GHCB_HV_FT_SNP_IOMMU_TLB_FLUSH))
+                       return -EPERM;
+

IOMMU performs RMP checks when SNP is enabled, the results are cached along with the IOMMU translations. When a VM lowers permission of a mapped page (moves to a lower VMPL level or from read+write to read-only or private to shared), the cached RMP check results require invalidation.

At the moment the only way to invalidate IOMMU cache is the RMPUPDATE instruction which flushes all IOMMU TLBs. It is a host privileged instruction so a VM needs a way to ensure the host has done it. Note that the guest's RMPADJUST/PVALIDATE do not flush IOMMU TLBs.

The host implements a new "IOMMU TLB Flush" VMGEXIT code which is advertised via bit#11 in the GHCB Hypervisor capabilities.

Use RMPUPDATE in the following way:
- allocate a page per VCPU (to allow lockless flushing);
- When invalidation is needed, copy two patterns (A and B) to the page;
- invalidate the page so the host can make it shared;
- use new GHCB call to request RMPUPDATE on the host;
- the host makes the page shared;
- the host clears pattern A;
- the host makes the page private again;
- the host returns to the guest;
- check if pattern A has changed and pattern B has not;
- if the above failed, panic().

The patterns are located far enough to not hit the same cache line to work with the cipher text hiding feature.

The host can choose to not execute the request, WARN_ON if this is the case. Further patches will attempt to handle this in other way.

Signed-off-by Alexey Kardashevskiy <[email protected]>

---
 arch/x86/include/asm/sev-common.h |  2 +
 arch/x86/include/uapi/asm/svm.h   |  3 +
 arch/x86/coco/sev/core.c          | 92 ++++++++++++++++++++
 3 files changed, 97 insertions(+)

arch/x86/include/asm/sev-common.h +2 -0

index ff763c3c5d63..51abf8d061fa 100644
@@ -138,6 +138,7 @@ enum psc_op {
 #define GHCB_HV_FT_SNP_AP_CREATION     BIT_ULL(1)
 #define GHCB_HV_FT_SNP_MULTI_VMPL      BIT_ULL(5)
 #define GHCB_HV_FT_SNP_SEV_TIO         BIT_ULL(7)
+#define GHCB_HV_FT_SNP_IOMMU_TLB_FLUSH BIT_ULL(11)
 
 /*
  * SNP Page State Change NAE event

@@ -210,6 +211,7 @@ struct snp_psc_desc {
 #define GHCB_TERM_SECURE_TSC           10      /* Secure TSC initialization failed */
 #define GHCB_TERM_SVSM_CA_REMAP_FAIL   11      /* SVSM is present but CA could not be remapped */
 #define GHCB_TERM_SAVIC_FAIL           12      /* Secure AVIC-specific failure */
+#define GHCB_TERM_IOMMUTLB_FLUSH       13      /* IOMMUTLB flush failed for SEV-TIO device */
 
 #define GHCB_RESP_CODE(v)              ((v) & GHCB_MSR_INFO_MASK)
 

arch/x86/include/uapi/asm/svm.h +3 -0

index 93597ad492bf..269050942c8e 100644
@@ -160,6 +160,8 @@
 #define SVM_VMGEXIT_SEV_TIO_OP_UNBIND  1
 #define SVM_VMGEXIT_SEV_TIO_OP_RUN     2
 #define SVM_VMGEXIT_SEV_TIO_OP_STOP    3
+#define SVM_VMGEXIT_IOMMU_TLB_FLUSH            0x80000022ull
+#define SVM_VMGEXIT_IOMMU_TLB_FLUSH_NO_ACTION  1
 #define SVM_VMGEXIT_HV_FEATURES                        0x8000fffdull
 #define SVM_VMGEXIT_TERM_REQUEST               0x8000fffeull
 #define SVM_VMGEXIT_TERM_REASON(reason_set, reason_code)       \

@@ -285,6 +287,7 @@
        { SVM_VMGEXIT_AP_CREATION,      "vmgexit_ap_creation" }, \
        { SVM_VMGEXIT_SEV_TIO_GR,       "vmgexit_sev_tio_guest_request" }, \
        { SVM_VMGEXIT_SEV_TIO_OP,       "vmgexit_sev_tio_op" }, \
+       { SVM_VMGEXIT_IOMMU_TLB_FLUSH, "vmgexit_sev_tio_iommu_tlb_flush" }, \
        { SVM_VMGEXIT_HV_FEATURES,      "vmgexit_hypervisor_feature" }, \
        { SVM_EXIT_ERR,         "invalid_guest_state" }
 

arch/x86/coco/sev/core.c +92 -1

index ed0e4546d5e5..aa5a3abb4796 100644
@@ -44,6 +44,7 @@
 #include <asm/cpuid/api.h>
 #include <asm/cmdline.h>
 #include <asm/msr.h>
+#include <asm/archrandom.h>
 
 #include "internal.h"
 

@@ -103,6 +104,36 @@ static unsigned long snp_tsc_freq_khz __ro_after_init;
 
 DEFINE_PER_CPU(struct sev_es_runtime_data*, runtime_data);
 DEFINE_PER_CPU(struct sev_es_save_area *, sev_vmsa);
+DEFINE_PER_CPU(u8 *, iommu_tlb_flush_ghcb_page);
+static atomic_t sev_tio_devices_num;
+
+static int alloc_iommu_tlb_flush_ghcb_pages(void)
+{
+       unsigned int cpu;
+       struct page *pg;
+       void *p;
+
+       /*
+        * Allocate per CPU pages while encrypted DMA is not happening yet
+        * and smashing is cheap.
+        */
+       for_each_possible_cpu(cpu) {
+               if (per_cpu(iommu_tlb_flush_ghcb_page, cpu))
+                       continue;
+
+               pg = alloc_pages_node(cpu_to_node(cpu), GFP_KERNEL, 0);
+               if (!pg)
+                       return -ENOMEM;
+
+               p = page_to_virt(pg);
+               /* Trigger psmash in the host os now to avoid psmash race later */
+               snp_set_memory_shared((unsigned long)p, 1);
+               snp_set_memory_private((unsigned long)p, 1);
+               per_cpu(iommu_tlb_flush_ghcb_page, cpu) = p;
+       }
+
+       return 0;
+}
 
 int sev_tio_op(u32 guest_rid, unsigned int op, u64 *fw_err, u64 *tdi_id)
 {

@@ -111,6 +142,24 @@ int sev_tio_op(u32 guest_rid, unsigned int op, u64 *fw_err, u64 *tdi_id)
        struct ghcb *ghcb;
        int ret;
 
+       if (!(sev_hv_features & GHCB_HV_FT_SNP_SEV_TIO))
+               return -EPERM;
+
+       if (op == SVM_VMGEXIT_SEV_TIO_OP_RUN || op == SVM_VMGEXIT_SEV_TIO_OP_STOP) {
+               if (!(sev_hv_features & GHCB_HV_FT_SNP_IOMMU_TLB_FLUSH))
+                       return -EPERM;
+
+               if (op == SVM_VMGEXIT_SEV_TIO_OP_RUN) {
+                       if (atomic_inc_return(&sev_tio_devices_num) == 1) {
+                               ret = alloc_iommu_tlb_flush_ghcb_pages();
+                               if (ret)
+                                       return ret;
+                       }
+               } else if (atomic_dec_return(&sev_tio_devices_num) == 0) {
+                       /* Do cleanup or leave it like this? */
+               }
+       }
+
        /* __sev_get_ghcb() needs IRQs disabled because it uses per-CPU GHCB. */
        guard(irqsave)();
 

@@ -347,6 +396,42 @@ static int vmgexit_psc(struct ghcb *ghcb, struct snp_psc_desc *desc)
        return ret;
 }
 
+static int ghcb_flush_iommu_tlb(struct ghcb *ghcb)
+{
+       /* AES encrypts with 16 byte blocks */
+       unsigned long s1[BITS_TO_LONGS(128)], s2[BITS_TO_LONGS(128)];
+       void *p = this_cpu_read(iommu_tlb_flush_ghcb_page), *p2;
+       struct es_em_ctxt ctxt;
+       int ret;
+
+       if (!p)
+               return -ENOMEM;
+
+       /* Keep patterns apart far enough to not share the same cache line */
+       p2 = (u8 *) p + 2048;
+
+       vc_ghcb_invalidate(ghcb);
+
+       BUILD_BUG_ON(ARRAY_SIZE(s1) != 2);
+       if (!rdrand_long(s1) || !rdrand_long(s1 + 1) ||
+           !rdrand_long(s2) || !rdrand_long(s2 + 1))
+               return -EFAULT;
+
+       memcpy(p, s1, sizeof(s1));
+       memcpy(p2, s2, sizeof(s2));
+
+       pvalidate((unsigned long) p, RMP_PG_SIZE_4K, false);
+       ret = sev_es_ghcb_hv_call(ghcb, &ctxt, SVM_VMGEXIT_IOMMU_TLB_FLUSH, __pa(p), 0);
+       pvalidate((unsigned long) p, RMP_PG_SIZE_4K, true);
+
+       /* Ensure that the host change is visible */
+       smp_mb();
+
+       if (!memcmp(p, s1, sizeof(s1)) || memcmp(p2, s2, sizeof(s2)))
+               return -EFAULT;
+
+       return 0;
+}
 static unsigned long __set_pages_state(struct snp_psc_desc *data, unsigned long vaddr,
                                       unsigned long vaddr_end, int op)
 {

@@ -404,6 +489,13 @@ static unsigned long __set_pages_state(struct snp_psc_desc *data, unsigned long
        if (!ghcb || vmgexit_psc(ghcb, data))
                sev_es_terminate(SEV_TERM_SET_LINUX, GHCB_TERM_PSC);
 
+       if (atomic_read(&sev_tio_devices_num)) {
+               int ret = ghcb_flush_iommu_tlb(ghcb);
+
+               if (ret)
+                       sev_es_terminate(SEV_TERM_SET_LINUX, GHCB_TERM_IOMMUTLB_FLUSH);
+       }
+
        __sev_put_ghcb(&state);
 
        local_irq_restore(flags);
--

2.55.0

Quoted 41 lines · Alexey Kardashevskiy

(quoted diff) +34 -0

+static int alloc_iommu_tlb_flush_ghcb_pages(void)
+{
+       unsigned int cpu;
+       struct page *pg;
+       void *p;
+
+       /*
+        * Allocate per CPU pages while encrypted DMA is not happening yet
+        * and smashing is cheap.
+        */
+       for_each_possible_cpu(cpu) {
+               if (per_cpu(iommu_tlb_flush_ghcb_page, cpu))
+                       continue;
+
+               pg = alloc_pages_node(cpu_to_node(cpu), GFP_KERNEL, 0);
+               if (!pg)
+                       return -ENOMEM;
+
+               p = page_to_virt(pg);
+               /* Trigger psmash in the host os now to avoid psmash race later */
+               snp_set_memory_shared((unsigned long)p, 1);
+               snp_set_memory_private((unsigned long)p, 1);
+               per_cpu(iommu_tlb_flush_ghcb_page, cpu) = p;
+       }
+
+       return 0;
+}
 
  int sev_tio_op(u32 guest_rid, unsigned int op, u64 *fw_err, u64 *tdi_id)
  {

@@ -111,6 +142,24 @@ int sev_tio_op(u32 guest_rid, unsigned int op, u64 *fw_err, u64 *tdi_id)
         struct ghcb *ghcb;
         int ret;
 
+       if (!(sev_hv_features & GHCB_HV_FT_SNP_SEV_TIO))
+               return -EPERM;
+
+       if (op == SVM_VMGEXIT_SEV_TIO_OP_RUN || op == SVM_VMGEXIT_SEV_TIO_OP_STOP) {
+               if (!(sev_hv_features & GHCB_HV_FT_SNP_IOMMU_TLB_FLUSH))
+                       return -EPERM;
+

Review int sev_tio_op(u32 guest_rid, unsigned int op, u64 *fw_err, u64 *tdi_id) · line 161

Quoted 10 lines

(quoted diff) +10 -0

+               if (op == SVM_VMGEXIT_SEV_TIO_OP_RUN) {
+                       if (atomic_inc_return(&sev_tio_devices_num) == 1) {
+                               ret = alloc_iommu_tlb_flush_ghcb_pages();
+                               if (ret)
+                                       return ret;
+                       }
+               } else if (atomic_dec_return(&sev_tio_devices_num) == 0) {
+                       /* Do cleanup or leave it like this? */
+               }
+       }

Hi Alexey,

When testing this series and accepting a locked TDI in the guest (echo 1 > /sys/bus/pci/devices/.../tsm/accept), the guest immediately terminates with 0x1:0xd (SEV_TERM_SET_LINUX, GHCB_TERM_IOMMUTLB_FLUSH).

In sev_tio_op(), sev_tio_devices_num is incremented from 0 to 1 before alloc_iommu_tlb_flush_ghcb_pages() allocates and initializes the per-CPU iommu_tlb_flush_ghcb_page buffers:

1. atomic_inc_return(&sev_tio_devices_num) sets sev_tio_devices_num = 1 while iommu_tlb_flush_ghcb_page is still NULL on all CPUs.
2. alloc_iommu_tlb_flush_ghcb_pages() allocates p for cpu = 0 and calls snp_set_memory_shared((unsigned long)p, 1) to pre-smash the 2M page before per_cpu(iommu_tlb_flush_ghcb_page, cpu) is assigned (and before other CPUs' pages are allocated, in case this task is running on cpu > 0).
3. snp_set_memory_shared() -> __set_pages_state() sees

    atomic_read(&sev_tio_devices_num) != 0 and calls ghcb_flush_iommu_tlb().

4. ghcb_flush_iommu_tlb() reads this_cpu_read(iommu_tlb_flush_ghcb_page), gets NULL, and returns -ENOMEM.
5. __set_pages_state() treats the non-zero return as fatal and calls sev_es_terminate(SEV_TERM_SET_LINUX, GHCB_TERM_IOMMUTLB_FLUSH).

Calling alloc_iommu_tlb_flush_ghcb_pages() before incrementing sev_tio_devices_num avoids triggering ghcb_flush_iommu_tlb() while the per-CPU pages are still being pre-smashed and initialized:

arch/x86/coco/sev/core.c +6 -7

@@ -150,13 +150,12 @@ int sev_tio_op(...)
              return -EPERM;
 
          if (op == SVM_VMGEXIT_SEV_TIO_OP_RUN) {
-            if (atomic_inc_return(&sev_tio_devices_num) == 1) {
-                ret = alloc_iommu_tlb_flush_ghcb_pages();
-                if (ret)
-                    return ret;
-            }
-        } else if (atomic_dec_return(&sev_tio_devices_num) == 0) {
-            /* Do cleanup or leave it like this? */
+            ret = alloc_iommu_tlb_flush_ghcb_pages();
+            if (ret)
+                return ret;
+            atomic_inc(&sev_tio_devices_num);
+        } else {
+            atomic_dec_if_positive(&sev_tio_devices_num);
          }
      }
 

Quoted 216 lines · Alexey Kardashevskiy

IOMMU performs RMP checks when SNP is enabled, the results are cached along with the IOMMU translations. When a VM lowers permission of a mapped page (moves to a lower VMPL level or from read+write to read-only or private to shared), the cached RMP check results require invalidation.

At the moment the only way to invalidate IOMMU cache is the RMPUPDATE instruction which flushes all IOMMU TLBs. It is a host privileged instruction so a VM needs a way to ensure the host has done it. Note that the guest's RMPADJUST/PVALIDATE do not flush IOMMU TLBs.

The host implements a new "IOMMU TLB Flush" VMGEXIT code which is advertised via bit#11 in the GHCB Hypervisor capabilities.

Use RMPUPDATE in the following way:
- allocate a page per VCPU (to allow lockless flushing);
- When invalidation is needed, copy two patterns (A and B) to the page;
- invalidate the page so the host can make it shared;
- use new GHCB call to request RMPUPDATE on the host;
- the host makes the page shared;
- the host clears pattern A;
- the host makes the page private again;
- the host returns to the guest;
- check if pattern A has changed and pattern B has not;
- if the above failed, panic().

The patterns are located far enough to not hit the same cache line to work with the cipher text hiding feature.

The host can choose to not execute the request, WARN_ON if this is the case. Further patches will attempt to handle this in other way.

Signed-off-by Alexey Kardashevskiy <[email protected]>

---
  arch/x86/include/asm/sev-common.h |  2 +
  arch/x86/include/uapi/asm/svm.h   |  3 +
  arch/x86/coco/sev/core.c          | 92 ++++++++++++++++++++
  3 files changed, 97 insertions(+)

arch/x86/include/asm/sev-common.h +2 -0

index ff763c3c5d63..51abf8d061fa 100644
@@ -138,6 +138,7 @@ enum psc_op {
  #define GHCB_HV_FT_SNP_AP_CREATION     BIT_ULL(1)
  #define GHCB_HV_FT_SNP_MULTI_VMPL      BIT_ULL(5)
  #define GHCB_HV_FT_SNP_SEV_TIO         BIT_ULL(7)
+#define GHCB_HV_FT_SNP_IOMMU_TLB_FLUSH BIT_ULL(11)
 
  /*
   * SNP Page State Change NAE event

@@ -210,6 +211,7 @@ struct snp_psc_desc {
  #define GHCB_TERM_SECURE_TSC           10      /* Secure TSC initialization failed */
  #define GHCB_TERM_SVSM_CA_REMAP_FAIL   11      /* SVSM is present but CA could not be remapped */
  #define GHCB_TERM_SAVIC_FAIL           12      /* Secure AVIC-specific failure */
+#define GHCB_TERM_IOMMUTLB_FLUSH       13      /* IOMMUTLB flush failed for SEV-TIO device */
 
  #define GHCB_RESP_CODE(v)              ((v) & GHCB_MSR_INFO_MASK)
 

arch/x86/include/uapi/asm/svm.h +3 -0

index 93597ad492bf..269050942c8e 100644
@@ -160,6 +160,8 @@
  #define SVM_VMGEXIT_SEV_TIO_OP_UNBIND  1
  #define SVM_VMGEXIT_SEV_TIO_OP_RUN     2
  #define SVM_VMGEXIT_SEV_TIO_OP_STOP    3
+#define SVM_VMGEXIT_IOMMU_TLB_FLUSH            0x80000022ull
+#define SVM_VMGEXIT_IOMMU_TLB_FLUSH_NO_ACTION  1
  #define SVM_VMGEXIT_HV_FEATURES                        0x8000fffdull
  #define SVM_VMGEXIT_TERM_REQUEST               0x8000fffeull
  #define SVM_VMGEXIT_TERM_REASON(reason_set, reason_code)       \

@@ -285,6 +287,7 @@
         { SVM_VMGEXIT_AP_CREATION,      "vmgexit_ap_creation" }, \
         { SVM_VMGEXIT_SEV_TIO_GR,       "vmgexit_sev_tio_guest_request" }, \
         { SVM_VMGEXIT_SEV_TIO_OP,       "vmgexit_sev_tio_op" }, \
+       { SVM_VMGEXIT_IOMMU_TLB_FLUSH, "vmgexit_sev_tio_iommu_tlb_flush" }, \
         { SVM_VMGEXIT_HV_FEATURES,      "vmgexit_hypervisor_feature" }, \
         { SVM_EXIT_ERR,         "invalid_guest_state" }
 

arch/x86/coco/sev/core.c +92 -1

index ed0e4546d5e5..aa5a3abb4796 100644
@@ -44,6 +44,7 @@
  #include <asm/cpuid/api.h>
  #include <asm/cmdline.h>
  #include <asm/msr.h>
+#include <asm/archrandom.h>
 
  #include "internal.h"
 

@@ -103,6 +104,36 @@ static unsigned long snp_tsc_freq_khz __ro_after_init;
 
  DEFINE_PER_CPU(struct sev_es_runtime_data*, runtime_data);
  DEFINE_PER_CPU(struct sev_es_save_area *, sev_vmsa);
+DEFINE_PER_CPU(u8 *, iommu_tlb_flush_ghcb_page);
+static atomic_t sev_tio_devices_num;
+
+static int alloc_iommu_tlb_flush_ghcb_pages(void)
+{
+       unsigned int cpu;
+       struct page *pg;
+       void *p;
+
+       /*
+        * Allocate per CPU pages while encrypted DMA is not happening yet
+        * and smashing is cheap.
+        */
+       for_each_possible_cpu(cpu) {
+               if (per_cpu(iommu_tlb_flush_ghcb_page, cpu))
+                       continue;
+
+               pg = alloc_pages_node(cpu_to_node(cpu), GFP_KERNEL, 0);
+               if (!pg)
+                       return -ENOMEM;
+
+               p = page_to_virt(pg);
+               /* Trigger psmash in the host os now to avoid psmash race later */
+               snp_set_memory_shared((unsigned long)p, 1);
+               snp_set_memory_private((unsigned long)p, 1);
+               per_cpu(iommu_tlb_flush_ghcb_page, cpu) = p;
+       }
+
+       return 0;
+}
 
  int sev_tio_op(u32 guest_rid, unsigned int op, u64 *fw_err, u64 *tdi_id)
  {

@@ -111,6 +142,24 @@ int sev_tio_op(u32 guest_rid, unsigned int op, u64 *fw_err, u64 *tdi_id)
         struct ghcb *ghcb;
         int ret;
 
+       if (!(sev_hv_features & GHCB_HV_FT_SNP_SEV_TIO))
+               return -EPERM;
+
+       if (op == SVM_VMGEXIT_SEV_TIO_OP_RUN || op == SVM_VMGEXIT_SEV_TIO_OP_STOP) {
+               if (!(sev_hv_features & GHCB_HV_FT_SNP_IOMMU_TLB_FLUSH))
+                       return -EPERM;
+
+               if (op == SVM_VMGEXIT_SEV_TIO_OP_RUN) {
+                       if (atomic_inc_return(&sev_tio_devices_num) == 1) {
+                               ret = alloc_iommu_tlb_flush_ghcb_pages();
+                               if (ret)
+                                       return ret;
+                       }
+               } else if (atomic_dec_return(&sev_tio_devices_num) == 0) {
+                       /* Do cleanup or leave it like this? */
+               }
+       }
+
         /* __sev_get_ghcb() needs IRQs disabled because it uses per-CPU GHCB. */
         guard(irqsave)();
 

@@ -347,6 +396,42 @@ static int vmgexit_psc(struct ghcb *ghcb, struct snp_psc_desc *desc)
         return ret;
  }
 
+static int ghcb_flush_iommu_tlb(struct ghcb *ghcb)
+{
+       /* AES encrypts with 16 byte blocks */
+       unsigned long s1[BITS_TO_LONGS(128)], s2[BITS_TO_LONGS(128)];
+       void *p = this_cpu_read(iommu_tlb_flush_ghcb_page), *p2;
+       struct es_em_ctxt ctxt;
+       int ret;
+
+       if (!p)
+               return -ENOMEM;
+
+       /* Keep patterns apart far enough to not share the same cache line */
+       p2 = (u8 *) p + 2048;
+
+       vc_ghcb_invalidate(ghcb);
+
+       BUILD_BUG_ON(ARRAY_SIZE(s1) != 2);
+       if (!rdrand_long(s1) || !rdrand_long(s1 + 1) ||
+           !rdrand_long(s2) || !rdrand_long(s2 + 1))
+               return -EFAULT;
+
+       memcpy(p, s1, sizeof(s1));
+       memcpy(p2, s2, sizeof(s2));
+
+       pvalidate((unsigned long) p, RMP_PG_SIZE_4K, false);
+       ret = sev_es_ghcb_hv_call(ghcb, &ctxt, SVM_VMGEXIT_IOMMU_TLB_FLUSH, __pa(p), 0);
+       pvalidate((unsigned long) p, RMP_PG_SIZE_4K, true);
+
+       /* Ensure that the host change is visible */
+       smp_mb();
+
+       if (!memcmp(p, s1, sizeof(s1)) || memcmp(p2, s2, sizeof(s2)))
+               return -EFAULT;
+
+       return 0;
+}
  static unsigned long __set_pages_state(struct snp_psc_desc *data, unsigned long vaddr,
                                        unsigned long vaddr_end, int op)
  {

@@ -404,6 +489,13 @@ static unsigned long __set_pages_state(struct snp_psc_desc *data, unsigned long
         if (!ghcb || vmgexit_psc(ghcb, data))
                 sev_es_terminate(SEV_TERM_SET_LINUX, GHCB_TERM_PSC);
 
+       if (atomic_read(&sev_tio_devices_num)) {
+               int ret = ghcb_flush_iommu_tlb(ghcb);
+
+               if (ret)
+                       sev_es_terminate(SEV_TERM_SET_LINUX, GHCB_TERM_IOMMUTLB_FLUSH);
+       }
+
         __sev_put_ghcb(&state);
 
         local_irq_restore(flags);
--

2.55.0