[PATCH v2 00/13] KVM: ITS hardening for pKVM

[PATCH v2 00/13] KVM: ITS hardening for pKVM

[Sebastian Ene]

This series introduces the necessary machinery to perform trap & emulate
on device access in pKVM. Furthermore, it hardens the GIC/ITS controller to
prevent an attacker from tampering with the hypervisor protected memory
through this device. 

In pKVM, the host kernel is initially trusted to manage the boot process but
its permissions are revoked once KVM initializes. The GIC/ITS device is
configured before the kernel deprivileges itself. Once the hypervisor
becomes available, sanitize the accesses to the ITS controller by
trapping and emulating certain registers and by shadowing some memory
structures used by the ITS.

This is required because the ITS is in complete host control and can issue
transactions on the memory bus without having any protection in front of it
(IOMMU or SMMU). With pKVM this makes it an interesting target for crossing the
hypervisor privilege boundary because with pKVM the security model is
that the host is untrusted. 

Patches overview
================

The first patch is an adaptation of Mostafa's patch for donating MMIO to
the hypervisor which can be found here [1]

[1] https://lore.kernel.org/all/20260715115906.2664882-3-smostafa@google.com/

The only difference is that this maps MMIO in the linear space of the
hypervisor, but after discussing with him on the list [2]
(https://lore.kernel.org/all/anMwxy9gEEt_Wu6M@google.com/) he seems to be
more inclined to having the IO in the private VA mapping space. For
simplicity, I used the linear map but I am happy to change it this is
required as it doesn't change the code too much.

The next 2 patches add some basic infrastructure to be able to keep
track of the unmapped MMIO regions from the host and invoke a handler
when a stage-2 data abort happens in that region.

The 4th patch looks up the ITS node from the device tree and adds it to
an array of unmapped devices. It install a handler that forwards all the
MMIO request to mediate the host access inside the emulation layer and
to prevent breaking ITS functionality. 

The 5th patch changes the GIC/ITS driver to exposes two new methods
which will be called from the KVM layer to setup the shadow state and
to take the appropriate locks.

Patch 6 adds the entry point into the emulation setup and sets up the
shadow command queue. It adds some helper macros to define the offset
register and the associate action that we want to execute in the
emulation. It also unmaps the state passed from the host kernel
to prevent it from playing nasty games later on. The patch
traps accesses to CWRITER register and copies the commands from the
host command queue to the shadow command queue. 

Patch 7 prevents the host from directly accessing the first layer of the
indirect tables held in GITS_BASER<n>. It also prevents the host from
directly accesssing the last layer of the Device Table (since the entries
in this table hold the address of the ITT table) and of the vPE Table
(since the vPE table entries hold the address of the virtual LPI pending
table.

Patches [8-10] sanitize the commands sent to the ITS and their
arguments.

The next patches restrict the access of the host to certain registers
and prevent undefined behaviour. Prevent the host from re-programming
the tables held in the GITS_BASER register.

The last patch introduces an hvc to setup the ITS emulation and calls
into the ITS driver to setup the shadow state. 

Design
======


1. Command queue shadowing

The ITS hardware supports a command queue which is programmed by the driver
in the GITS_CBASER register. When ITS emulation is enabled, the ITS
driver allocated a copy of the original command queue and uses the copy
and the emulation layer restricts the access to the original queue for
the host. When the driver sends a command to the mirrored queue, the
emulation traps on the write to GITS_CWRITER register and validates the
command before writing it over to the original queue. 


2. Indirect tables first level shadowing

The ITS hardware supports indirection to minimize the space required to
accommodate large tables (eg. deviceId space used to index the Device Table
is quite sparse). This is a 2-level indirection, with entries from the
first table pointing to a second table.

An attacker in control of the host can insert an address that points to
the hypervisor protected memory in the first level table and then use
subsequent ITS commands to write to this memory (MAPD).

To shadow this tables, we rely on the driver to allocate space for it
and we copy the original content from the table into the copy. When
pKVM becomes available we switch the pointers that hold the orginal
tables to point to the copy.
To keep the tables from the hypervisor in sync with what the host
has, we update the tables when commands are sent to the ITS.


3. Hiding the last layer of the Device Table and vPE Table from the host

An attacker in control of the host kernel can alter the content of these
tables directly (the Arm IHI 0069H.b spec says that is undefined behavior
if entries are created by software). Normally these entries are created in
response of commands sent to the ITS.

A Device Table entry that has the following structure:

type DeviceTableEntry is (
	boolean Valid,
	Address ITT_base,
	bits(5) ITT_size
) 

This can be maliciously created by an attacker and the ITT_base can be
pointed to hypervisor protected memory. The MAPTI command can then be
used to write over the ITT_base with an ITE entry.

Similarly a vCPU Table entry has the following structure:

type VCPUTableEntry is (
	boolean Valid,
	bits(32) RDbase,
	Address VPT_base,
	bits(5) VPT_size
)

VPT_base can be pointed to hypervisor protected memory and then a
command can be used to raise interrupts and set the corresponding
bit. This would give a 1-bit write primitive so is not "as generous"
as the others.

Testing
=======

Verified that it boots in Qemu using the following arguments:
"-machine virt,virtualization=true,gic-version=4,its=true,iommu=smmuv3,acpi=off"

and enabled Qemu tracing using the following argument: '--trace "gicv3_its_*"'
Qemu reports a bunch of events due to virtio-block using the ITS to
raise MSIs and I can see the accesses made by the device to the
GITS_TRANSLATER register.

Wrote a PoC [3] which uses a DEVMEM like interface to poke the ITS
directly from userspace to achieve controlled write to hyp protected
memory and verified that this is not possible anymore when the
sanitization is in place.

[3] [https://github.com/sebastianene07/gic-its-poc/blob/main/its.c]


What is not covered
===================

Only a subset of the commands is sanitized & covered: vLPI commands are
not covered at all and the Redistributor hardeninig is not covered. 


Changelog
=========

v1 -> v2:
 - fixed a bunch of functional errors in the parsing of the MAPD command
 - renamed the functions exported from the GIC driver, made the
   allocation GFP_ATOMIC, make it accepts a flags arg
 - renamed the entry point in the emulation and make it so that it
   accepts the size of pages that we use for the private state
 - added support for rollback on certain commands that timeout on
   hardware
 - clarified commit messages 

Previous posting:
################

v1: https://lore.kernel.org/all/20260310124933.830025-1-sebastianene@google.com/


 Mostafa Saleh (1):
  KVM: arm64: Donate MMIO to the hypervisor

Sebastian Ene (12):
  KVM: arm64: Track host-unmapped MMIO regions in a static array
  KVM: arm64: Support host MMIO trap handlers for unmapped devices
  KVM: Parse the device tree and register the ITS region with pKVM
  irqchip/gic-v3-its: Add support for the ITS emulation setup
  KVM: arm64: Shadow the ITS command queue and setup emulation
  KVM: arm64: Restrict host access to the private ITS tables
  KVM: arm64: Trap & emulate the ITS MAPD command
  KVM: arm64: Trap & emulate the ITS MAPC command
  KVM: arm64: Restrict host updates to GITS_CTLR
  KVM: arm64: Prevent the host from specifying a different command queue
  KVM: arm64: Prevent the host from programming new GITS_BASER tables
  KVM: arm64: Implement HVC interface for ITS emulation setup

 arch/arm64/include/asm/kvm_arm.h              |   2 +
 arch/arm64/include/asm/kvm_asm.h              |   1 +
 arch/arm64/include/asm/kvm_pkvm.h             |  18 +
 arch/arm64/kvm/hyp/include/nvhe/its_emulate.h |  14 +
 arch/arm64/kvm/hyp/include/nvhe/mem_protect.h |   7 +
 arch/arm64/kvm/hyp/nvhe/Makefile              |   3 +-
 arch/arm64/kvm/hyp/nvhe/hyp-main.c            |  16 +
 arch/arm64/kvm/hyp/nvhe/its_emulate.c         | 827 ++++++++++++++++++
 arch/arm64/kvm/hyp/nvhe/mem_protect.c         | 189 +++-
 arch/arm64/kvm/hyp/nvhe/setup.c               |  27 +
 arch/arm64/kvm/hyp/pgtable.c                  |  11 +-
 arch/arm64/kvm/pkvm.c                         |  76 +-
 drivers/irqchip/irq-gic-v3-its.c              | 169 +++-
 include/linux/irqchip/arm-gic-v3.h            |  51 ++
 14 files changed, 1379 insertions(+), 32 deletions(-)
 create mode 100644 arch/arm64/kvm/hyp/include/nvhe/its_emulate.h
 create mode 100644 arch/arm64/kvm/hyp/nvhe/its_emulate.c

-- 
2.55.0.654.g21b8a5bc05-goog

[PATCH v2 01/13] KVM: arm64: Donate MMIO to the hypervisor

[Sebastian Ene]

From: Mostafa Saleh <smostafa@google.com>

Extend the pKVM API to allow the donation of MMIO from the host
address space to the hypervisor linear map.
Initialize the host s2 page table with an invalid leaf with the owner ID
of the hypervisor to prevent the host from mapping the page on faults.
Prevent kvm_pgtable_stage2_unmap() from removing owner ID from
stage-2 PTEs, as this can be triggered from recycle logic under memory
pressure.

Signed-off-by: Mostafa Saleh <smostafa@google.com>
Signed-off-by: Sebastian Ene <sebastianene@google.com>
---
 arch/arm64/kvm/hyp/include/nvhe/mem_protect.h |   7 +
 arch/arm64/kvm/hyp/nvhe/mem_protect.c         | 137 +++++++++++++++++-
 arch/arm64/kvm/hyp/pgtable.c                  |  11 +-
 3 files changed, 148 insertions(+), 7 deletions(-)

diff --git a/arch/arm64/kvm/hyp/include/nvhe/mem_protect.h b/arch/arm64/kvm/hyp/include/nvhe/mem_protect.h
index 29935c7da1de..6aa83b129e61 100644
--- a/arch/arm64/kvm/hyp/include/nvhe/mem_protect.h
+++ b/arch/arm64/kvm/hyp/include/nvhe/mem_protect.h
@@ -36,6 +36,13 @@ int __pkvm_guest_share_host(struct pkvm_hyp_vcpu *vcpu, u64 gfn);
 int __pkvm_guest_unshare_host(struct pkvm_hyp_vcpu *vcpu, u64 gfn);
 int __pkvm_host_unshare_hyp(u64 pfn);
 int __pkvm_host_donate_hyp(u64 pfn, u64 nr_pages);
+/*
+ * Donate MMIO range to the hypervisor, it will be mapped in the hypervisor's
+ * linea map and unmapped from the host stage-2.
+ */
+int __pkvm_host_donate_hyp_mmio(phys_addr_t addr, size_t size);
+/* Remaps MMIO range in the host, typically used in error path. */
+int __pkvm_hyp_donate_host_mmio(phys_addr_t addr, size_t size);
 int __pkvm_hyp_donate_host(u64 pfn, u64 nr_pages);
 int __pkvm_host_share_ffa(u64 pfn, u64 nr_pages);
 int __pkvm_host_unshare_ffa(u64 pfn, u64 nr_pages);
diff --git a/arch/arm64/kvm/hyp/nvhe/mem_protect.c b/arch/arm64/kvm/hyp/nvhe/mem_protect.c
index 4e329e39a695..5cf7c4a0ed20 100644
--- a/arch/arm64/kvm/hyp/nvhe/mem_protect.c
+++ b/arch/arm64/kvm/hyp/nvhe/mem_protect.c
@@ -378,7 +378,11 @@ static int host_stage2_unmap_dev_all(void)
 	u64 addr = 0;
 	int i, ret;
 
-	/* Unmap all non-memory regions to recycle the pages */
+	/*
+	 * Unmap all non-memory regions to recycle the pages.
+	 * That relies on kvm_pgtable_stage2_unmap() not clearing
+	 * counted PTEs which include hypervisor MMIO.
+	 */
 	for (i = 0; i < hyp_memblock_nr; i++, addr = reg->base + reg->size) {
 		reg = &hyp_memory[i];
 		ret = kvm_pgtable_stage2_unmap(pgt, addr, reg->base - addr);
@@ -1119,6 +1123,137 @@ int __pkvm_host_donate_hyp(u64 pfn, u64 nr_pages)
 	return ret;
 }
 
+int __pkvm_host_donate_hyp_mmio(phys_addr_t addr, size_t size)
+{
+	kvm_pte_t pte;
+	u64 offset;
+	void *virt;
+	int ret;
+
+	/* Only before de-privilege. */
+	if (static_branch_unlikely(&kvm_protected_mode_initialized))
+		return -EPERM;
+
+	if (!PAGE_ALIGNED(addr | size) ||
+	    !pfn_range_is_valid(hyp_phys_to_pfn(addr), size >> PAGE_SHIFT))
+		return -EINVAL;
+
+	host_lock_component();
+	hyp_lock_component();
+
+	for (offset = 0; offset < size; offset += PAGE_SIZE) {
+		if (addr_is_memory(addr + offset)) {
+			ret = -EINVAL;
+			goto err_with_mapping;
+		}
+
+		ret = kvm_pgtable_get_leaf(&host_mmu.pgt, addr + offset, &pte, NULL);
+		if (ret)
+			goto err_with_mapping;
+
+		if (pte && !kvm_pte_valid(pte)) {
+			ret = -EPERM;
+			goto err_with_mapping;
+		}
+
+		virt = __hyp_va(addr + offset);
+		ret = kvm_pgtable_get_leaf(&pkvm_pgtable, (u64)virt, &pte, NULL);
+		if (ret)
+			goto err_with_mapping;
+		if (pte) {
+			ret = -EBUSY;
+			goto err_with_mapping;
+		}
+
+		ret = pkvm_create_mappings_locked(virt, virt + PAGE_SIZE, PAGE_HYP_DEVICE);
+		if (ret)
+			goto err_with_mapping;
+	}
+
+	/*
+	 * We set HYP as the owner of the MMIO pages in the host stage-2, for:
+	 * - host aborts: host_stage2_adjust_range() would fail for invalid non zero PTEs.
+	 * - recycle under memory pressure: host_stage2_unmap_dev_all() would call
+	 *   kvm_pgtable_stage2_unmap() which will not clear non zero invalid ptes (counted).
+	 * - other MMIO donation: Would fail as we check that the PTE is valid or empty.
+	 */
+	ret = host_stage2_try(kvm_pgtable_stage2_annotate, &host_mmu.pgt,
+			      addr, size, &host_s2_pool,
+			      KVM_HOST_INVALID_PTE_TYPE_DONATION,
+			      FIELD_PREP(KVM_HOST_DONATION_PTE_OWNER_MASK, PKVM_ID_HYP));
+	if (ret)
+		goto err_with_mapping;
+unlock:
+	hyp_unlock_component();
+	host_unlock_component();
+	return ret;
+err_with_mapping:
+	if (!offset)
+		goto unlock;
+
+	while (offset) {
+		offset -= PAGE_SIZE;
+		virt = __hyp_va(addr + offset);
+		WARN_ON(kvm_pgtable_hyp_unmap(&pkvm_pgtable, (u64)virt, PAGE_SIZE) != PAGE_SIZE);
+	}
+	goto unlock;
+}
+
+int __pkvm_hyp_donate_host_mmio(phys_addr_t addr, size_t size)
+{
+	kvm_pte_t pte;
+	u64 offset;
+	int ret = 0;
+	void *virt;
+
+	if (static_branch_unlikely(&kvm_protected_mode_initialized))
+		return -EPERM;
+
+	if (!PAGE_ALIGNED(addr | size) ||
+	    !pfn_range_is_valid(hyp_phys_to_pfn(addr), size >> PAGE_SHIFT))
+		return -EINVAL;
+
+	host_lock_component();
+	hyp_lock_component();
+
+	for (offset = 0; offset < size; offset += PAGE_SIZE) {
+		if (addr_is_memory(addr + offset)) {
+			ret = -EINVAL;
+			goto err_with_unmap;
+		}
+		ret = kvm_pgtable_get_leaf(&host_mmu.pgt, addr + offset, &pte, NULL);
+		if (ret)
+			goto err_with_unmap;
+		if (!pte || kvm_pte_valid(pte)) {
+			ret = -EINVAL;
+			goto err_with_unmap;
+		}
+		if (FIELD_GET(KVM_HOST_DONATION_PTE_OWNER_MASK, pte) != PKVM_ID_HYP) {
+			ret = -EPERM;
+			goto err_with_unmap;
+		}
+
+		virt = __hyp_va(addr + offset);
+		if (kvm_pgtable_hyp_unmap(&pkvm_pgtable, (u64)virt, PAGE_SIZE) != PAGE_SIZE)
+			goto err_with_unmap;
+	}
+	WARN_ON(host_stage2_idmap_locked(addr, size, PKVM_HOST_MMIO_PROT));
+unlock:
+	hyp_unlock_component();
+	host_unlock_component();
+	return ret;
+err_with_unmap:
+	if (!offset)
+		goto unlock;
+
+	while (offset) {
+		offset -= PAGE_SIZE;
+		virt = __hyp_va(addr + offset);
+		WARN_ON(pkvm_create_mappings_locked(virt, virt + PAGE_SIZE, PAGE_HYP_DEVICE));
+	}
+	goto unlock;
+}
+
 int __pkvm_hyp_donate_host(u64 pfn, u64 nr_pages)
 {
 	u64 phys = hyp_pfn_to_phys(pfn);
diff --git a/arch/arm64/kvm/hyp/pgtable.c b/arch/arm64/kvm/hyp/pgtable.c
index b74dd5ce1efd..7638213bd893 100644
--- a/arch/arm64/kvm/hyp/pgtable.c
+++ b/arch/arm64/kvm/hyp/pgtable.c
@@ -1161,13 +1161,12 @@ static int stage2_unmap_walker(const struct kvm_pgtable_visit_ctx *ctx,
 	kvm_pte_t *childp = NULL;
 	bool need_flush = false;
 
-	if (!kvm_pte_valid(ctx->old)) {
-		if (stage2_pte_is_counted(ctx->old)) {
-			kvm_clear_pte(ctx->ptep);
-			mm_ops->put_page(ctx->ptep);
-		}
+	/*
+	 * That also ignores stage2_pte_is_counted() instead of clearing
+	 * the PTE as the MMIO can be owned by the hypervisor.
+	 */
+	if (!kvm_pte_valid(ctx->old))
 		return 0;
-	}
 
 	if (kvm_pte_table(ctx->old, ctx->level)) {
 		childp = kvm_pte_follow(ctx->old, mm_ops);
-- 
2.55.0.654.g21b8a5bc05-goog

[PATCH v2 02/13] KVM: arm64: Track host-unmapped MMIO regions in a static array

[Sebastian Ene]

Introduce a registry to track protected MMIO regions that are unmapped
from the host stage-2 page tables. These regions are stored in a
fixed-size array and their ownership is donated to the hypervisor during
initialization to ensure host-exclusion and persistent tracking.

Signed-off-by: Sebastian Ene <sebastianene@google.com>
---
 arch/arm64/include/asm/kvm_pkvm.h     | 11 +++++++++++
 arch/arm64/kvm/hyp/nvhe/mem_protect.c |  3 +++
 arch/arm64/kvm/hyp/nvhe/setup.c       | 24 ++++++++++++++++++++++++
 3 files changed, 38 insertions(+)

diff --git a/arch/arm64/include/asm/kvm_pkvm.h b/arch/arm64/include/asm/kvm_pkvm.h
index 74fedd9c5ff0..ab26bec079d6 100644
--- a/arch/arm64/include/asm/kvm_pkvm.h
+++ b/arch/arm64/include/asm/kvm_pkvm.h
@@ -17,6 +17,17 @@
 
 #define HYP_MEMBLOCK_REGIONS 128
 
+/* The maximum number of hypervisor protected regions from the host */
+#define PKVM_PROTECTED_REGS_NUM	8
+
+struct pkvm_protected_reg {
+	u64 pfn;
+	u64 nr_pages;
+};
+
+extern struct pkvm_protected_reg kvm_nvhe_sym(pkvm_protected_regs)[];
+extern unsigned int kvm_nvhe_sym(num_protected_reg);
+
 int pkvm_init_host_vm(struct kvm *kvm, unsigned long type);
 int pkvm_create_hyp_vm(struct kvm *kvm);
 bool pkvm_hyp_vm_is_created(struct kvm *kvm);
diff --git a/arch/arm64/kvm/hyp/nvhe/mem_protect.c b/arch/arm64/kvm/hyp/nvhe/mem_protect.c
index 5cf7c4a0ed20..500c18c2fd48 100644
--- a/arch/arm64/kvm/hyp/nvhe/mem_protect.c
+++ b/arch/arm64/kvm/hyp/nvhe/mem_protect.c
@@ -28,6 +28,9 @@ struct host_mmu host_mmu;
 
 static struct hyp_pool host_s2_pool;
 
+unsigned int num_protected_reg;
+struct pkvm_protected_reg pkvm_protected_regs[PKVM_PROTECTED_REGS_NUM];
+
 static DEFINE_PER_CPU(struct pkvm_hyp_vm *, __current_vm);
 #define current_vm (*this_cpu_ptr(&__current_vm))
 
diff --git a/arch/arm64/kvm/hyp/nvhe/setup.c b/arch/arm64/kvm/hyp/nvhe/setup.c
index 75b00c323310..64c0290da888 100644
--- a/arch/arm64/kvm/hyp/nvhe/setup.c
+++ b/arch/arm64/kvm/hyp/nvhe/setup.c
@@ -284,6 +284,26 @@ static int fix_hyp_pgtable_refcnt(void)
 				&walker);
 }
 
+static int donate_protected_mmio_regions(void)
+{
+	int ret;
+	int i;
+
+	for (i = 0; i < num_protected_reg; i++) {
+		ret = __pkvm_host_donate_hyp_mmio(hyp_pfn_to_phys(pkvm_protected_regs[i].pfn),
+						  pkvm_protected_regs[i].nr_pages << PAGE_SHIFT);
+		if (ret)
+			goto err_setup;
+	}
+
+	return 0;
+err_setup:
+	while (--i >= 0)
+		__pkvm_hyp_donate_host_mmio(hyp_pfn_to_phys(pkvm_protected_regs[i].pfn),
+					    pkvm_protected_regs[i].nr_pages << PAGE_SHIFT);
+	return ret;
+}
+
 void __noreturn __pkvm_init_finalise(void)
 {
 	struct kvm_cpu_context *host_ctxt = host_data_ptr(host_ctxt);
@@ -324,6 +344,10 @@ void __noreturn __pkvm_init_finalise(void)
 	if (ret)
 		goto out;
 
+	ret = donate_protected_mmio_regions();
+	if (ret)
+		goto out;
+
 	ret = hyp_ffa_init(ffa_proxy_pages);
 	if (ret)
 		goto out;
-- 
2.55.0.654.g21b8a5bc05-goog

[PATCH v2 03/13] KVM: arm64: Support host MMIO trap handlers for unmapped devices

[Sebastian Ene]

Hook a handler to the host mem abort so that the hypervisor can
intercept host accesses to unmapped memory regions.
When a Stage-2 fault occurs on a registered device region, the
hypervisor will look if there is any registered function that
can handle the access. On the back of this, mediate host accesses
to devices and emulate them in pKVM.

Signed-off-by: Sebastian Ene <sebastianene@google.com>
Signed-off-by: Bartłomiej Grzesik <bgrzesik@google.com>
---
 arch/arm64/include/asm/kvm_arm.h      |  2 ++
 arch/arm64/include/asm/kvm_pkvm.h     |  4 +++
 arch/arm64/kvm/hyp/nvhe/mem_protect.c | 49 +++++++++++++++++++++++++++
 arch/arm64/kvm/hyp/nvhe/setup.c       |  3 ++
 4 files changed, 58 insertions(+)

diff --git a/arch/arm64/include/asm/kvm_arm.h b/arch/arm64/include/asm/kvm_arm.h
index 3f9233b5a130..6360c90f9855 100644
--- a/arch/arm64/include/asm/kvm_arm.h
+++ b/arch/arm64/include/asm/kvm_arm.h
@@ -304,6 +304,8 @@
 
 /* Hyp Prefetch Fault Address Register (HPFAR/HDFAR) */
 #define HPFAR_MASK	(~UL(0xf))
+#define FAR_MASK	GENMASK_ULL(11, 0)
+
 /*
  * We have
  *	PAR	[PA_Shift - 1	: 12] = PA	[PA_Shift - 1 : 12]
diff --git a/arch/arm64/include/asm/kvm_pkvm.h b/arch/arm64/include/asm/kvm_pkvm.h
index ab26bec079d6..0a471564be00 100644
--- a/arch/arm64/include/asm/kvm_pkvm.h
+++ b/arch/arm64/include/asm/kvm_pkvm.h
@@ -20,9 +20,13 @@
 /* The maximum number of hypervisor protected regions from the host */
 #define PKVM_PROTECTED_REGS_NUM	8
 
+struct pkvm_protected_reg;
+typedef void(pkvm_emulate_handler)(struct pkvm_protected_reg *region, u64 offset,
+				   bool write, u64 *reg, u8 reg_size);
 struct pkvm_protected_reg {
 	u64 pfn;
 	u64 nr_pages;
+	pkvm_emulate_handler *cb;
 };
 
 extern struct pkvm_protected_reg kvm_nvhe_sym(pkvm_protected_regs)[];
diff --git a/arch/arm64/kvm/hyp/nvhe/mem_protect.c b/arch/arm64/kvm/hyp/nvhe/mem_protect.c
index 500c18c2fd48..7e978e0c44b9 100644
--- a/arch/arm64/kvm/hyp/nvhe/mem_protect.c
+++ b/arch/arm64/kvm/hyp/nvhe/mem_protect.c
@@ -14,6 +14,7 @@
 #include <asm/stage2_pgtable.h>
 
 #include <hyp/fault.h>
+#include <hyp/adjust_pc.h>
 
 #include <nvhe/arm-smccc.h>
 #include <nvhe/gfp.h>
@@ -752,6 +753,50 @@ static void host_inject_mem_abort(struct kvm_cpu_context *host_ctxt)
 	inject_host_exception(esr);
 }
 
+static bool handle_host_mmio_trap(struct kvm_cpu_context *host_ctxt, u64 esr, u64 addr)
+{
+	u64 offset, reg_value = 0, start, end;
+	u8 reg_size, reg_index;
+	bool write;
+	int i;
+
+	for (i = 0; i < num_protected_reg; i++) {
+		if (!pkvm_protected_regs[i].pfn || !pkvm_protected_regs[i].nr_pages ||
+		    !pkvm_protected_regs[i].cb)
+			continue;
+
+		start = PFN_PHYS(pkvm_protected_regs[i].pfn);
+		end = start + PFN_PHYS(pkvm_protected_regs[i].nr_pages);
+		reg_size = BIT((esr & ESR_ELx_SAS) >> ESR_ELx_SAS_SHIFT);
+
+		if (start > addr || addr + reg_size > end)
+			continue;
+
+		reg_index = (esr & ESR_ELx_SRT_MASK) >> ESR_ELx_SRT_SHIFT;
+		write = (esr & ESR_ELx_WNR) == ESR_ELx_WNR;
+		offset = addr - start;
+
+		if (write && reg_index != 31)
+			reg_value = host_ctxt->regs.regs[reg_index];
+
+		pkvm_protected_regs[i].cb(&pkvm_protected_regs[i], offset, write,
+					  &reg_value, reg_size);
+
+		if (!write && reg_index != 31)
+			host_ctxt->regs.regs[reg_index] = reg_value;
+
+		kvm_skip_host_instr();
+		return true;
+	}
+
+	return false;
+}
+
+static bool is_dabt(u64 esr)
+{
+	return (ESR_ELx_EC(esr) == ESR_ELx_EC_DABT_LOW) && (esr & ESR_ELx_ISV);
+}
+
 void handle_host_mem_abort(struct kvm_cpu_context *host_ctxt)
 {
 	struct kvm_vcpu_fault_info fault;
@@ -774,6 +819,10 @@ void handle_host_mem_abort(struct kvm_cpu_context *host_ctxt)
 	BUG_ON(!(fault.hpfar_el2 & HPFAR_EL2_NS));
 	addr = FIELD_GET(HPFAR_EL2_FIPA, fault.hpfar_el2) << 12;
 
+	if (is_dabt(esr) && !addr_is_memory(addr) &&
+	    handle_host_mmio_trap(host_ctxt, esr, addr | (fault.far_el2 & FAR_MASK)))
+		return;
+
 	switch (host_stage2_idmap(addr)) {
 	case -EPERM:
 		host_inject_mem_abort(host_ctxt);
diff --git a/arch/arm64/kvm/hyp/nvhe/setup.c b/arch/arm64/kvm/hyp/nvhe/setup.c
index 64c0290da888..4395595b7f7e 100644
--- a/arch/arm64/kvm/hyp/nvhe/setup.c
+++ b/arch/arm64/kvm/hyp/nvhe/setup.c
@@ -294,6 +294,9 @@ static int donate_protected_mmio_regions(void)
 						  pkvm_protected_regs[i].nr_pages << PAGE_SHIFT);
 		if (ret)
 			goto err_setup;
+
+		if (pkvm_protected_regs[i].cb)
+			pkvm_protected_regs[i].cb = kern_hyp_va(pkvm_protected_regs[i].cb);
 	}
 
 	return 0;
-- 
2.55.0.654.g21b8a5bc05-goog

[PATCH v2 04/13] KVM: Parse the device tree and register the ITS region with pKVM

[Sebastian Ene]

Identify the ITS base address from the device tree and store it in the
pkvm_protected_regs array so that it will be unmapped from the host
address space.
Register a callback to forward all the MMIO requests to the device to
prevent breaking ITS functionality in this patch. The patch by itself
shouldn't break any existing functionality even though all the accesses
from the gic-ITS driver are now mediated inside pKVM.

Signed-off-by: Sebastian Ene <sebastianene@google.com>
---
 arch/arm64/include/asm/kvm_pkvm.h     |  2 ++
 arch/arm64/kvm/hyp/nvhe/Makefile      |  3 +-
 arch/arm64/kvm/hyp/nvhe/its_emulate.c | 37 +++++++++++++++++++
 arch/arm64/kvm/pkvm.c                 | 52 +++++++++++++++++++++++++++
 4 files changed, 93 insertions(+), 1 deletion(-)
 create mode 100644 arch/arm64/kvm/hyp/nvhe/its_emulate.c

diff --git a/arch/arm64/include/asm/kvm_pkvm.h b/arch/arm64/include/asm/kvm_pkvm.h
index 0a471564be00..370225f0e72c 100644
--- a/arch/arm64/include/asm/kvm_pkvm.h
+++ b/arch/arm64/include/asm/kvm_pkvm.h
@@ -31,6 +31,8 @@ struct pkvm_protected_reg {
 
 extern struct pkvm_protected_reg kvm_nvhe_sym(pkvm_protected_regs)[];
 extern unsigned int kvm_nvhe_sym(num_protected_reg);
+extern void kvm_nvhe_sym(its_emulate_forward_req)(struct pkvm_protected_reg *region, u64 offset,
+						  bool write, u64 *reg, u8 reg_size);
 
 int pkvm_init_host_vm(struct kvm *kvm, unsigned long type);
 int pkvm_create_hyp_vm(struct kvm *kvm);
diff --git a/arch/arm64/kvm/hyp/nvhe/Makefile b/arch/arm64/kvm/hyp/nvhe/Makefile
index f57450ebcb49..70fbca325852 100644
--- a/arch/arm64/kvm/hyp/nvhe/Makefile
+++ b/arch/arm64/kvm/hyp/nvhe/Makefile
@@ -24,7 +24,8 @@ CFLAGS_switch.nvhe.o += -Wno-override-init
 
 hyp-obj-y := timer-sr.o sysreg-sr.o debug-sr.o switch.o tlb.o hyp-init.o host.o \
 	 hyp-main.o hyp-smp.o psci-relay.o early_alloc.o page_alloc.o \
-	 cache.o setup.o mm.o mem_protect.o sys_regs.o pkvm.o stacktrace.o ffa.o
+	 cache.o setup.o mm.o mem_protect.o sys_regs.o pkvm.o stacktrace.o ffa.o \
+	 its_emulate.o
 hyp-obj-y += ../vgic-v3-sr.o ../aarch32.o ../vgic-v2-cpuif-proxy.o ../entry.o \
 	 ../hyp-entry.o ../exception.o ../pgtable.o ../vgic-v5-sr.o
 hyp-obj-y += ../../../kernel/smccc-call.o
diff --git a/arch/arm64/kvm/hyp/nvhe/its_emulate.c b/arch/arm64/kvm/hyp/nvhe/its_emulate.c
new file mode 100644
index 000000000000..63a42f520ed2
--- /dev/null
+++ b/arch/arm64/kvm/hyp/nvhe/its_emulate.c
@@ -0,0 +1,37 @@
+// SPDX-License-Identifier: GPL-2.0-only
+
+#include <asm/kvm_pkvm.h>
+#include <nvhe/mem_protect.h>
+
+void its_emulate_forward_req(struct pkvm_protected_reg *region, u64 offset, bool write, u64 *reg,
+			     u8 reg_size)
+{
+	void __iomem *addr = __hyp_va(PFN_PHYS(region->pfn) + offset);
+
+	switch (reg_size) {
+	case 1:
+		if (!write)
+			*reg = readb_relaxed(addr);
+		else
+			writeb_relaxed(*reg, addr);
+		break;
+	case 2:
+		if (!write)
+			*reg = readw_relaxed(addr);
+		else
+			writew_relaxed(*reg, addr);
+		break;
+	case 4:
+		if (!write)
+			*reg = readl_relaxed(addr);
+		else
+			writel_relaxed(*reg, addr);
+		break;
+	case 8:
+		if (!write)
+			*reg = readq_relaxed(addr);
+		else
+			writeq_relaxed(*reg, addr);
+		break;
+	}
+}
diff --git a/arch/arm64/kvm/pkvm.c b/arch/arm64/kvm/pkvm.c
index 428723b1b0f5..4bfffbedac4c 100644
--- a/arch/arm64/kvm/pkvm.c
+++ b/arch/arm64/kvm/pkvm.c
@@ -9,8 +9,13 @@
 #include <linux/kmemleak.h>
 #include <linux/kvm_host.h>
 #include <asm/kvm_mmu.h>
+#include <asm/kvm_pkvm.h>
 #include <linux/memblock.h>
 #include <linux/mutex.h>
+#include <linux/of_address.h>
+#include <linux/of_reserved_mem.h>
+#include <linux/platform_device.h>
+#include <linux/irqchip/arm-gic-v3.h>
 
 #include <asm/kvm_pkvm.h>
 
@@ -39,6 +44,47 @@ static int __init register_memblock_regions(void)
 	return 0;
 }
 
+static int __init register_its_emulated_region(void)
+{
+	struct device_node *np;
+	struct resource res;
+	int i = 0;
+	int ret;
+
+	for_each_compatible_node(np, NULL, "arm,gic-v3-its") {
+		ret = of_address_to_resource(np, 0, &res);
+		if (ret)
+			goto out_fail;
+
+		if (i >= PKVM_PROTECTED_REGS_NUM) {
+			kvm_err("Out of protected region slots\n");
+			ret = -ENOSPC;
+			goto out_fail;
+		}
+
+		/*
+		 * Note: don't unmap the entire animal from the host because devices need
+		 * to be able to access GITS_TRANSLATER to raise MSIs. If the
+		 * page where GITS_TRANSLATER is given to HYP, devices won't be
+		 * able to map it in their IOMMU when the IOMMU is managed by
+		 * pKVM.
+		 */
+		kvm_nvhe_sym(pkvm_protected_regs)[i].pfn = PHYS_PFN(res.start);
+		kvm_nvhe_sym(pkvm_protected_regs)[i].cb =
+			lm_alias(&kvm_nvhe_sym(its_emulate_forward_req));
+		kvm_nvhe_sym(pkvm_protected_regs)[i].nr_pages =
+			PFN_DOWN(min_t(u64, resource_size(&res), PAGE_ALIGN_DOWN(GITS_TRANSLATER)));
+
+		i++;
+	}
+
+	kvm_nvhe_sym(num_protected_reg) = i;
+	return 0;
+out_fail:
+	of_node_put(np);
+	return ret;
+}
+
 void __init kvm_hyp_reserve(void)
 {
 	u64 hyp_mem_pages = 0;
@@ -57,6 +103,12 @@ void __init kvm_hyp_reserve(void)
 		return;
 	}
 
+	ret = register_its_emulated_region();
+	if (ret) {
+		kvm_err("Failed to register ITS region %d\n", ret);
+		return;
+	}
+
 	hyp_mem_pages += hyp_s1_pgtable_pages();
 	hyp_mem_pages += host_s2_pgtable_pages();
 	hyp_mem_pages += hyp_vm_table_pages();
-- 
2.55.0.654.g21b8a5bc05-goog

[PATCH v2 05/13] irqchip/gic-v3-its: Add support for the ITS emulation setup

[Sebastian Ene]

Introduce two new helper functions to allow locking the ITS and setting
up a copy of the host ITS state that will be given to the pKVM
emulation. The caller of these functions is responsible to implement a
callback which will be used to setup the emulation layer. The calling
flow is expected to do the following:

pkvm_its_emulate_setup(its_phys, host)
	// allocate memory for the priv state of the ITS emulation
	// call the its emulation setup(its_phys, host, priv_state);

pkvm_drop_host_privileges()
	its_emulate_acquire_locks(&flags);
		on_each_cpu(_kvm_host_prot_finalize, &ret, 1);
	its_emulate_release_locks(ret, &flags, pkvm_its_emulate_setup);

Augment the its_baser structure with a new fiels that will hold a
pointer to the base table copy. The gic ITS driver will use the pointer
to the base table copy when emulation is enabled, as this allows us to
hide away the original first level of an indirect table to prevent the
following:

// assumming an indirect Device Table layout
1. malicious host patches an entry in the 1st level table with an
   address that it wants to write to.
2. malicious host issues MAPD to install a DTE in the table pointed by
   the address from (1).

As the driver only manipulates a copy of the table, the emulation is
responsible for looking at the updates from the copy table, sanitizing
them and updating the original table before talking to the hardware.

In a simillar fashion, when emulation is in place we no longer let the
gic ITS driver use the original command queue but we present the driver
a copy of it and we hide away the original command queue from the driver
as this will be used entirely by the emulation layer.

Co-authored-by: Bartłomiej Grzesik <bgrzesik@google.com>
Signed-off-by: Sebastian Ene <sebastianene@google.com>
---
 drivers/irqchip/irq-gic-v3-its.c   | 157 +++++++++++++++++++++++++++--
 include/linux/irqchip/arm-gic-v3.h |  39 +++++++
 2 files changed, 185 insertions(+), 11 deletions(-)

diff --git a/drivers/irqchip/irq-gic-v3-its.c b/drivers/irqchip/irq-gic-v3-its.c
index 6f5811aae59c..e74ae9220af5 100644
--- a/drivers/irqchip/irq-gic-v3-its.c
+++ b/drivers/irqchip/irq-gic-v3-its.c
@@ -78,17 +78,6 @@ struct its_collection {
 	u16			col_id;
 };
 
-/*
- * The ITS_BASER structure - contains memory information, cached
- * value of BASER register configuration and ITS page size.
- */
-struct its_baser {
-	void		*base;
-	u64		val;
-	u32		order;
-	u32		psz;
-};
-
 struct its_device;
 
 /*
@@ -5226,6 +5215,152 @@ static int __init its_compute_its_list_map(struct its_node *its)
 	return its_number;
 }
 
+static void its_free_snapshot(struct its_host_state *snapshot)
+{
+	int i;
+
+	if (snapshot->cmd_host_copy)
+		its_free_pages(snapshot->cmd_host_copy, get_order(ITS_CMD_QUEUE_SZ));
+
+	for (i = 0; i < GITS_BASER_NR_REGS; i++) {
+		if (!snapshot->tables[i].base_snapshot)
+			continue;
+
+		its_free_pages(snapshot->tables[i].base_snapshot, snapshot->tables[i].order);
+	}
+
+	its_free_pages(snapshot, 0);
+}
+
+static struct its_host_state *its_snapshot_host_state(struct its_node *its)
+{
+	void *page;
+	struct its_host_state *snapshot;
+	int i;
+
+	page = its_alloc_pages_node(its->numa_node, GFP_ATOMIC | __GFP_ZERO, 0);
+	if (!page)
+		return NULL;
+
+	snapshot = (void *)page_address(page);
+	page = its_alloc_pages_node(its->numa_node, GFP_ATOMIC | __GFP_ZERO,
+				    get_order(ITS_CMD_QUEUE_SZ));
+	if (!page)
+		goto err_alloc;
+
+	snapshot->cmd_host_copy	= page_address(page);
+	snapshot->cmdq_len	= ITS_CMD_QUEUE_SZ;
+	snapshot->cmd_original	= its->cmd_base;
+	snapshot->cmd_write	= its->cmd_write;
+
+	memcpy(snapshot->tables, its->tables, sizeof(struct its_baser) * GITS_BASER_NR_REGS);
+
+	for (i = 0; i < GITS_BASER_NR_REGS; i++) {
+		if (!(snapshot->tables[i].val & GITS_BASER_VALID))
+			continue;
+
+		if (!(snapshot->tables[i].val & GITS_BASER_INDIRECT))
+			continue;
+
+		page = its_alloc_pages_node(its->numa_node,
+					    GFP_ATOMIC | __GFP_ZERO,
+					    snapshot->tables[i].order);
+		if (!page)
+			goto err_alloc;
+
+		snapshot->tables[i].base_snapshot = page_address(page);
+
+		memcpy(snapshot->tables[i].base_snapshot, snapshot->tables[i].base,
+		       PAGE_ORDER_TO_SIZE(snapshot->tables[i].order));
+	}
+
+	return snapshot;
+
+err_alloc:
+	its_free_snapshot(snapshot);
+	return NULL;
+}
+
+static int its_emulate_switch_queues_locked(struct its_node *its, its_emulate_setup cb)
+{
+	struct its_host_state *host_snaphsot, host;
+	int i, ret;
+	u64 baser_phys;
+
+	host_snaphsot = its_snapshot_host_state(its);
+	if (!host_snaphsot)
+		return -ENOMEM;
+
+	/*
+	 * The snapshot of the ITS state will be given to the emulation, make a copy of it
+	 * so that we don't go in weeds.
+	 */
+	memcpy(&host, host_snaphsot, sizeof(host));
+
+	ret = cb(its->phys_base, host_snaphsot);
+	if (ret) {
+		its_free_snapshot(host_snaphsot);
+		return ret;
+	}
+
+	/* Switch the driver command queue to use the host copy and update the write index */
+	its->cmd_write = (its->cmd_write - its->cmd_base) +
+		(struct its_cmd_block *)host.cmd_host_copy;
+	its->cmd_base = host.cmd_host_copy;
+
+	/*
+	 * Replace the first level of the indirect tables with the snapshot table as the
+	 * emulation layer will make it innaccessible to the host.
+	 */
+	for (i = 0; i < GITS_BASER_NR_REGS; i++) {
+		if (!(host.tables[i].val & GITS_BASER_INDIRECT))
+			continue;
+
+		baser_phys = virt_to_phys(host.tables[i].base_snapshot);
+		if (IS_ENABLED(CONFIG_ARM64_64K_PAGES) && (baser_phys >> 48))
+			baser_phys = GITS_BASER_PHYS_52_to_48(baser_phys);
+
+		its->tables[i].val &= ~GENMASK(47, 12);
+		its->tables[i].val |= baser_phys;
+		its->tables[i].base = host.tables[i].base_snapshot;
+	}
+
+	return 0;
+}
+
+void its_emulate_acquire_locks(unsigned long *flags)
+{
+	struct its_node *its;
+
+	if (WARN_ON(!flags))
+		return;
+
+	raw_spin_lock_irqsave(&its_lock, *flags);
+
+	list_for_each_entry(its, &its_nodes, entry)
+		raw_spin_lock(&its->lock);
+}
+
+int its_emulate_release_locks(int ret_pkvm_finalize, unsigned long *flags, its_emulate_setup cb)
+{
+	struct its_node *its;
+	int ret = 0;
+
+	if (WARN_ON(!flags || !cb))
+		ret = -EINVAL;
+
+	list_for_each_entry(its, &its_nodes, entry) {
+		if (!ret_pkvm_finalize && !ret)
+			ret = its_emulate_switch_queues_locked(its, cb);
+
+		raw_spin_unlock(&its->lock);
+	}
+
+	raw_spin_unlock_irqrestore(&its_lock, *flags);
+
+	return ret;
+}
+
 static int __init its_probe_one(struct its_node *its)
 {
 	u64 baser, tmp;
diff --git a/include/linux/irqchip/arm-gic-v3.h b/include/linux/irqchip/arm-gic-v3.h
index ea5fd2374ebe..b75f82cef4bf 100644
--- a/include/linux/irqchip/arm-gic-v3.h
+++ b/include/linux/irqchip/arm-gic-v3.h
@@ -657,6 +657,45 @@ static inline bool gic_enable_sre(void)
 	return !!(val & ICC_SRE_EL1_SRE);
 }
 
+/*
+ * The ITS_BASER structure - contains memory information, cached
+ * value of BASER register configuration and ITS page size.
+ */
+struct its_baser {
+	void		*base;
+
+	/*
+	 * The table used when emulation is in place and indirect layout is
+	 * configured.
+	 */
+	void		*base_snapshot;
+	u64		val;
+	u32		order;
+	u32		psz;
+};
+
+struct its_host_state {
+	struct its_baser	tables[GITS_BASER_NR_REGS];
+
+	/* The command queue used after the emulation is in place */
+	void			*cmd_host_copy;
+
+	/* The command queue configured by the ITS driver at boot */
+	void			*cmd_original;
+	void			*cmd_write;
+	size_t			cmdq_len;
+};
+
+/*
+ * Callback used to initialize the emulation. It is expected to allocate memory for the private
+ * state of the emulation and receive as arguments copy of the host ITS driver state along
+ * with the address of the ITS.
+ */
+typedef int (*its_emulate_setup)(phys_addr_t its_phys_base, struct its_host_state *host);
+
+void its_emulate_acquire_locks(unsigned long *flags);
+int its_emulate_release_locks(int ret_pkvm_finalize, unsigned long *flags, its_emulate_setup cb);
+
 #endif
 
 #endif
-- 
2.55.0.654.g21b8a5bc05-goog

[PATCH v2 06/13] KVM: arm64: Shadow the ITS command queue and setup emulation

[Sebastian Ene]

Expose two functions that will be used to setup the entry point into the
pKVM ITS emulation. One will be called from an hvc to setup the ITS
structures and the other one will be called from a data abort to handle
the emulation. The later one will be stored in a pkvm_protected_reg as
part of the register_its_emulated_region once the command emulation is in
place.

Donate two memory regions as part of the emulation setup phase. One
holds GIC ITS driver state information and the other is used to store
private state information for the emulation and it is zeroed out.
Shadow the command queue by sharing a copy of it from the GIC ITS driver
and donate the original queue to the hypervisor. The host will use a
copy, while the emulation will use the original queue programmed in
hardware. This makes sure that the original queue is not accessible to
the host.
When the GIC ITS driver writes a command, the emulation will trap the
access to the CWRITER register and it will validate the
command before copying it to the original queue.

Re-use some of the definitions for command format and move them
from the GIC ITS driver to the public header.

Co-authored-by: Bartłomiej Grzesik <bgrzesik@google.com>
Signed-off-by: Sebastian Ene <sebastianene@google.com>
---
 arch/arm64/include/asm/kvm_pkvm.h             |   1 +
 arch/arm64/kvm/hyp/include/nvhe/its_emulate.h |  14 +
 arch/arm64/kvm/hyp/nvhe/its_emulate.c         | 285 ++++++++++++++++++
 drivers/irqchip/irq-gic-v3-its.c              |  12 -
 include/linux/irqchip/arm-gic-v3.h            |  12 +
 5 files changed, 312 insertions(+), 12 deletions(-)
 create mode 100644 arch/arm64/kvm/hyp/include/nvhe/its_emulate.h

diff --git a/arch/arm64/include/asm/kvm_pkvm.h b/arch/arm64/include/asm/kvm_pkvm.h
index 370225f0e72c..78597210a53c 100644
--- a/arch/arm64/include/asm/kvm_pkvm.h
+++ b/arch/arm64/include/asm/kvm_pkvm.h
@@ -27,6 +27,7 @@ struct pkvm_protected_reg {
 	u64 pfn;
 	u64 nr_pages;
 	pkvm_emulate_handler *cb;
+	void *priv;
 };
 
 extern struct pkvm_protected_reg kvm_nvhe_sym(pkvm_protected_regs)[];
diff --git a/arch/arm64/kvm/hyp/include/nvhe/its_emulate.h b/arch/arm64/kvm/hyp/include/nvhe/its_emulate.h
new file mode 100644
index 000000000000..29429feb30a9
--- /dev/null
+++ b/arch/arm64/kvm/hyp/include/nvhe/its_emulate.h
@@ -0,0 +1,14 @@
+/* SPDX-License-Identifier: GPL-2.0-only */
+
+#ifndef __NVHE_ITS_EMULATE_H
+#define __NVHE_ITS_EMULATE_H
+
+#include <asm/kvm_pkvm.h>
+
+struct its_host_state;
+
+int pkvm_its_emulate_setup(phys_addr_t dev_addr, struct its_host_state *host_state, void *priv,
+			   size_t priv_num_pages);
+void pkvm_its_emulate_handler(struct pkvm_protected_reg *region, u64 offset, bool write, u64 *reg,
+			      u8 reg_size);
+#endif /* __NVHE_ITS_EMULATE_H */
diff --git a/arch/arm64/kvm/hyp/nvhe/its_emulate.c b/arch/arm64/kvm/hyp/nvhe/its_emulate.c
index 63a42f520ed2..e943ab972aa5 100644
--- a/arch/arm64/kvm/hyp/nvhe/its_emulate.c
+++ b/arch/arm64/kvm/hyp/nvhe/its_emulate.c
@@ -2,6 +2,9 @@
 
 #include <asm/kvm_pkvm.h>
 #include <nvhe/mem_protect.h>
+#include <nvhe/its_emulate.h>
+
+#include <linux/irqchip/arm-gic-v3.h>
 
 void its_emulate_forward_req(struct pkvm_protected_reg *region, u64 offset, bool write, u64 *reg,
 			     u8 reg_size)
@@ -35,3 +38,285 @@ void its_emulate_forward_req(struct pkvm_protected_reg *region, u64 offset, bool
 		break;
 	}
 }
+
+struct its_handler {
+	u64	offset;
+	u8	access_size;
+	void	(*write)(struct pkvm_protected_reg *region, u64 offset, u64 value);
+	void	(*read)(struct pkvm_protected_reg *region, u64 offset, u64 *read);
+};
+
+#define ITS_HANDLER(off, sz, write_cb, read_cb)		\
+{							\
+	.offset = (off),				\
+	.access_size = (sz),				\
+	.write = (write_cb),				\
+	.read = (read_cb),				\
+}
+
+struct its_priv_state {
+	/* The location of the ITS in the hypervisor VA */
+	void __iomem	*base;
+
+	/* ITS command queue use by the hardware */
+	void		*cmd_original;
+	void		*cmd_host_copy;
+	u64		cmd_offset;
+	bool		needs_flush;
+	hyp_spinlock_t	its_lock;
+
+	struct its_host_state	*host_state;
+};
+
+#define GITS_CWRITER_RETRY	BIT_ULL(0)
+#define GITS_CWRITER_OFFSET	GENMASK_ULL(19, 5)
+
+#define GITS_CREADR_STALLED	BIT_ULL(0)
+#define GITS_CREADR_OFFSET	GENMASK_ULL(19, 5)
+
+static int submit_single_cmd(struct its_priv_state *its, bool retry)
+{
+	size_t cmdq_sz = its->host_state->cmdq_len;
+	u64 timeout = 1000;
+	u64 offset, cwriter, creadr;
+
+	offset = (its->cmd_offset + sizeof(struct its_cmd_block)) % cmdq_sz;
+
+	cwriter = offset & GITS_CWRITER_OFFSET;
+	cwriter |= FIELD_PREP(GITS_CWRITER_RETRY, retry);
+	writeq_relaxed(cwriter, its->base + GITS_CWRITER);
+
+	while (its->cmd_offset != offset) {
+		creadr = readq_relaxed(its->base + GITS_CREADR);
+
+		/* Command failed. */
+		if (FIELD_GET(GITS_CREADR_STALLED, creadr))
+			return -EIO;
+
+		its->cmd_offset = creadr & GITS_CREADR_OFFSET;
+		if (its->cmd_offset == offset)
+			return 0;
+
+		/*
+		 * We can't spin here forever and we can't roll back
+		 * the cmd queue pointer. Let's revert the cmd effects in the
+		 * emulation layer and then go back to the driver to let it
+		 * decide what to do next.
+		 */
+		if (!timeout--)
+			return -EBUSY;
+	}
+
+	return 0;
+}
+
+static int process_cmd(struct its_priv_state *its, struct its_cmd_block *cmd,
+		       bool rollback)
+{
+	/* Passthrough everything for now */
+	return 0;
+}
+
+static void cwriter_write(struct pkvm_protected_reg *region, u64 offset, u64 value)
+{
+	struct its_priv_state *its = region->priv;
+	struct its_cmd_block cmd, raw;
+	u64 new_offset;
+	bool retry;
+	int i;
+
+	new_offset = value & GITS_CWRITER_OFFSET;
+	if (new_offset >= its->host_state->cmdq_len)
+		return;
+
+	retry = FIELD_GET(GITS_CWRITER_RETRY, value);
+	while (its->cmd_offset != new_offset) {
+		memcpy(&raw, its->cmd_host_copy + its->cmd_offset, sizeof(raw));
+
+		for (i = 0; i < ARRAY_SIZE(cmd.raw_cmd); i++)
+			cmd.raw_cmd[i] = le64_to_cpu(raw.raw_cmd_le[i]);
+
+		if (process_cmd(its, &cmd, /* rollback */ false))
+			return;
+
+		memcpy(its->cmd_original + its->cmd_offset, &raw, sizeof(struct its_cmd_block));
+
+		if (its->needs_flush)
+			gic_flush_dcache_to_poc(its->cmd_original + its->cmd_offset, sizeof(cmd));
+		else
+			dsb(ishst);
+
+		if (submit_single_cmd(its, retry)) {
+			WARN_ON(process_cmd(its, &cmd, /* rollback */ true));
+			return;
+		}
+	}
+}
+
+static void cwriter_read(struct pkvm_protected_reg *region, u64 offset, u64 *read)
+{
+	struct its_priv_state *its = region->priv;
+	*read = readq_relaxed(its->base + GITS_CWRITER);
+}
+
+static struct its_handler its_handlers[] = {
+	ITS_HANDLER(GITS_CWRITER, sizeof(u64), cwriter_write, cwriter_read),
+	{},
+};
+
+void pkvm_its_emulate_handler(struct pkvm_protected_reg *region, u64 offset, bool write, u64 *reg,
+			      u8 reg_size)
+{
+	struct its_priv_state *priv = region->priv;
+	struct its_handler *reg_handler;
+
+	if (!priv || !IS_ALIGNED(offset, reg_size))
+		return;
+
+	for (reg_handler = its_handlers; reg_handler->access_size; reg_handler++) {
+		if (reg_handler->offset > offset ||
+		    reg_handler->offset + reg_handler->access_size <= offset)
+			continue;
+
+		if (reg_handler->access_size < reg_size)
+			return;
+
+		if (write && reg_handler->write) {
+			hyp_spin_lock(&priv->its_lock);
+			reg_handler->write(region, offset, *reg);
+			hyp_spin_unlock(&priv->its_lock);
+			return;
+		}
+
+		if (!write && reg_handler->read) {
+			hyp_spin_lock(&priv->its_lock);
+			reg_handler->read(region, offset, reg);
+			hyp_spin_unlock(&priv->its_lock);
+			return;
+		}
+
+		return;
+	}
+
+	its_emulate_forward_req(region, offset, write, reg, reg_size);
+}
+
+static int pkvm_setup_its_shadow_cmdq(struct its_host_state *host_state)
+{
+	u64 start_pfn, num_pages, i;
+	int ret;
+
+	start_pfn = hyp_virt_to_pfn(host_state->cmd_host_copy);
+	num_pages = host_state->cmdq_len >> PAGE_SHIFT;
+
+	for (i = 0; i < num_pages; i++) {
+		ret = __pkvm_host_share_hyp(start_pfn + i);
+		if (ret)
+			goto unshare_cmd_host;
+	}
+
+	ret = hyp_pin_shared_mem(host_state->cmd_host_copy,
+				 host_state->cmd_host_copy + host_state->cmdq_len);
+	if (ret)
+		goto unshare_cmd_host;
+
+	ret = __pkvm_host_donate_hyp(hyp_virt_to_pfn(host_state->cmd_original), num_pages);
+	if (ret) {
+		hyp_unpin_shared_mem(host_state->cmd_host_copy,
+				     host_state->cmd_host_copy + host_state->cmdq_len);
+		goto unshare_cmd_host;
+	}
+
+	return ret;
+unshare_cmd_host:
+	if (i == 0)
+		return ret;
+
+	for (i = i - 1; i >= 0; i--)
+		__pkvm_host_unshare_hyp(start_pfn + i);
+	return ret;
+}
+
+static struct pkvm_protected_reg *get_region(phys_addr_t dev_addr)
+{
+	int i;
+
+	for (i = 0; i < num_protected_reg; i++) {
+		if (PFN_PHYS(pkvm_protected_regs[i].pfn) == dev_addr)
+			return &pkvm_protected_regs[i];
+	}
+
+	return NULL;
+}
+
+DEFINE_HYP_SPINLOCK(its_setup_lock);
+
+int pkvm_its_emulate_setup(phys_addr_t dev_addr, struct its_host_state *host_state, void *priv,
+			   size_t priv_num_pages)
+{
+	struct pkvm_protected_reg *its_reg;
+	struct its_priv_state *priv_state;
+	int ret;
+
+	if (!PAGE_ALIGNED(host_state) || !PAGE_ALIGNED(priv) || !priv_num_pages)
+		return -EINVAL;
+
+	host_state = kern_hyp_va(host_state);
+	priv = kern_hyp_va(priv);
+
+	hyp_spin_lock(&its_setup_lock);
+	its_reg = get_region(dev_addr);
+	if (!its_reg) {
+		ret = -ENODEV;
+		goto err_unlock;
+	}
+
+	if (its_reg->priv) {
+		ret = -EOPNOTSUPP;
+		goto err_unlock;
+	}
+
+	ret = __pkvm_host_donate_hyp(hyp_virt_to_pfn(priv), priv_num_pages);
+	if (ret)
+		goto err_unlock;
+
+	priv_state = priv;
+	memset(priv_state, 0, priv_num_pages << PAGE_SHIFT);
+
+	ret = __pkvm_host_donate_hyp(hyp_virt_to_pfn(host_state), 1);
+	if (ret)
+		goto err_with_priv;
+
+	host_state->cmd_original = kern_hyp_va(host_state->cmd_original);
+	host_state->cmd_host_copy = kern_hyp_va(host_state->cmd_host_copy);
+
+	ret = pkvm_setup_its_shadow_cmdq(host_state);
+	if (ret)
+		goto err_with_host_state;
+
+	hyp_spin_lock_init(&priv_state->its_lock);
+
+	priv_state->host_state = host_state;
+	priv_state->base = (void __iomem *)__hyp_va(dev_addr);
+	priv_state->cmd_original = host_state->cmd_original;
+	priv_state->cmd_host_copy = host_state->cmd_host_copy;
+
+	priv_state->cmd_offset = readq_relaxed(priv_state->base + GITS_CREADR) &
+		GITS_CREADR_OFFSET;
+	priv_state->needs_flush =
+		(readq_relaxed(priv_state->base + GITS_CBASER) & GITS_CBASER_SHAREABILITY_MASK) !=
+		GITS_CBASER_InnerShareable;
+
+	its_reg->priv = priv_state;
+
+	hyp_spin_unlock(&its_setup_lock);
+
+	return 0;
+err_with_host_state:
+	WARN_ON(__pkvm_hyp_donate_host(hyp_virt_to_pfn(host_state), 1));
+err_with_priv:
+	WARN_ON(__pkvm_hyp_donate_host(hyp_virt_to_pfn(priv_state), 1));
+err_unlock:
+	hyp_spin_unlock(&its_setup_lock);
+	return ret;
+}
diff --git a/drivers/irqchip/irq-gic-v3-its.c b/drivers/irqchip/irq-gic-v3-its.c
index e74ae9220af5..4736e49e3f2d 100644
--- a/drivers/irqchip/irq-gic-v3-its.c
+++ b/drivers/irqchip/irq-gic-v3-its.c
@@ -121,8 +121,6 @@ static DEFINE_PER_CPU(struct its_node *, local_4_1_its);
 #define is_v4_1(its)		(!!((its)->typer & GITS_TYPER_VMAPP))
 #define device_ids(its)		(FIELD_GET(GITS_TYPER_DEVBITS, (its)->typer) + 1)
 
-#define ITS_ITT_ALIGN		SZ_256
-
 /* The maximum number of VPEID bits supported by VLPI commands */
 #define ITS_MAX_VPEID_BITS						\
 	({								\
@@ -515,16 +513,6 @@ struct its_cmd_desc {
 	};
 };
 
-/*
- * The ITS command block, which is what the ITS actually parses.
- */
-struct its_cmd_block {
-	union {
-		u64	raw_cmd[4];
-		__le64	raw_cmd_le[4];
-	};
-};
-
 #define ITS_CMD_QUEUE_SZ		SZ_64K
 #define ITS_CMD_QUEUE_NR_ENTRIES	(ITS_CMD_QUEUE_SZ / sizeof(struct its_cmd_block))
 
diff --git a/include/linux/irqchip/arm-gic-v3.h b/include/linux/irqchip/arm-gic-v3.h
index b75f82cef4bf..7f72632115b8 100644
--- a/include/linux/irqchip/arm-gic-v3.h
+++ b/include/linux/irqchip/arm-gic-v3.h
@@ -524,6 +524,8 @@
 #define GITS_CMD_VSGI			GITS_CMD_GICv4(3)
 #define GITS_CMD_INVDB			GITS_CMD_GICv4(0xe)
 
+#define ITS_ITT_ALIGN		SZ_256
+
 /*
  * ITS error numbers
  */
@@ -686,6 +688,16 @@ struct its_host_state {
 	size_t			cmdq_len;
 };
 
+/*
+ * The ITS command block, which is what the ITS actually parses.
+ */
+struct its_cmd_block {
+	union {
+		u64	raw_cmd[4];
+		__le64	raw_cmd_le[4];
+	};
+};
+
 /*
  * Callback used to initialize the emulation. It is expected to allocate memory for the private
  * state of the emulation and receive as arguments copy of the host ITS driver state along
-- 
2.55.0.654.g21b8a5bc05-goog

[PATCH v2 07/13] KVM: arm64: Restrict host access to the private ITS tables

[Sebastian Ene]

Make the last level of the tables(DeviceTable, Collection and vPE)
inaccessible to the host by donating them to the hypervisor.
This prevents a compromised host from patching an entry with an
address that it wants to write to and then using an ITS command to
write over the memory content from that address.

When tables are configured with indirect layout, shadow the first
layer by copying it to a separate table, update the gic ITS host
driver to use the copy instead of the original table and share the copy
between the host and the hypervisor. Make the original layer
innaccessible to the host by donating the table memory from the host to
the hypervisor.
This ensures that the pKVM ITS emulation mediates the
configuration written by the driver in the first layer of the table and
sanitizes the entries before writing to the original table programmed
in hardware. The update phase of the original table from the copy will
be done when commands are sent to the ITS.

Signed-off-by: Sebastian Ene <sebastianene@google.com>
---
 arch/arm64/kvm/hyp/nvhe/its_emulate.c | 161 ++++++++++++++++++++++++++
 1 file changed, 161 insertions(+)

diff --git a/arch/arm64/kvm/hyp/nvhe/its_emulate.c b/arch/arm64/kvm/hyp/nvhe/its_emulate.c
index e943ab972aa5..1ce2f9d8fcf9 100644
--- a/arch/arm64/kvm/hyp/nvhe/its_emulate.c
+++ b/arch/arm64/kvm/hyp/nvhe/its_emulate.c
@@ -237,6 +237,20 @@ static int pkvm_setup_its_shadow_cmdq(struct its_host_state *host_state)
 	return ret;
 }
 
+static void pkvm_teardown_its_shadow_cmdq(struct its_host_state *host_state)
+{
+	u64 i, start_pfn, num_pages = host_state->cmdq_len >> PAGE_SHIFT;
+
+	start_pfn = hyp_virt_to_pfn(host_state->cmd_host_copy);
+	hyp_unpin_shared_mem(host_state->cmd_host_copy,
+			     host_state->cmd_host_copy + host_state->cmdq_len);
+
+	for (i = 0; i < num_pages; i++)
+		WARN_ON(__pkvm_host_unshare_hyp(start_pfn + i));
+
+	WARN_ON(__pkvm_hyp_donate_host(hyp_virt_to_pfn(host_state->cmd_original), num_pages));
+}
+
 static struct pkvm_protected_reg *get_region(phys_addr_t dev_addr)
 {
 	int i;
@@ -249,6 +263,147 @@ static struct pkvm_protected_reg *get_region(phys_addr_t dev_addr)
 	return NULL;
 }
 
+static void pkvm_unshare_shadow_table(void *shadow, u64 nr_pages)
+{
+	u64 i, start_pfn = hyp_virt_to_pfn(shadow);
+
+	hyp_unpin_shared_mem(shadow, shadow + (nr_pages << PAGE_SHIFT));
+
+	for (i = 0; i < nr_pages; i++)
+		WARN_ON(__pkvm_host_unshare_hyp(start_pfn + i));
+}
+
+static int pkvm_host_unmap_last_level(void *shadow, size_t num_pages, u32 psz)
+{
+	phys_addr_t table_addr;
+	u64 *table = shadow;
+	int i, end;
+	int ret;
+
+	end = (num_pages << PAGE_SHIFT) / sizeof(*table);
+	for (i = 0; i < end; i++) {
+		if (!(table[i] & GITS_BASER_VALID))
+			continue;
+
+		table_addr = table[i] & PHYS_MASK;
+		ret = __pkvm_host_donate_hyp(hyp_phys_to_pfn(table_addr), psz >> PAGE_SHIFT);
+		if (ret)
+			goto err_donate;
+	}
+
+	return 0;
+err_donate:
+	for (i = i - 1; i >= 0; i--) {
+		if (!(table[i] & GITS_BASER_VALID))
+			continue;
+
+		table_addr = table[i] & PHYS_MASK;
+		__pkvm_hyp_donate_host(hyp_phys_to_pfn(table_addr), psz >> PAGE_SHIFT);
+	}
+	return ret;
+}
+
+static int pkvm_share_shadow_table(void *shadow, u64 nr_pages)
+{
+	u64 i, ret, start_pfn = hyp_virt_to_pfn(shadow);
+
+	for (i = 0; i < nr_pages; i++) {
+		ret = __pkvm_host_share_hyp(start_pfn + i);
+		if (ret)
+			goto unshare;
+	}
+
+	ret = hyp_pin_shared_mem(shadow, shadow + (nr_pages << PAGE_SHIFT));
+	if (ret)
+		goto unshare;
+
+	return ret;
+unshare:
+	while (i--)
+		__pkvm_host_unshare_hyp(start_pfn + i);
+	return ret;
+}
+
+static void pkvm_host_map_last_level(void *shadow, size_t num_pages, u32 psz)
+{
+	u64 *table = shadow;
+	int i, end = (num_pages << PAGE_SHIFT) / sizeof(*table);
+	phys_addr_t table_addr;
+
+	for (i = 0; i < end; i++) {
+		if (!(table[i] & GITS_BASER_VALID))
+			continue;
+
+		table_addr = table[i] & PHYS_MASK;
+		WARN_ON(__pkvm_hyp_donate_host(hyp_phys_to_pfn(table_addr), psz >> PAGE_SHIFT));
+	}
+}
+
+static int pkvm_setup_its_shadow_baser(struct its_host_state *host_state)
+{
+	u64 baser_val, num_pages;
+	void *original_table, *snapshot_table;
+	int ret;
+	int i;
+
+	for (i = 0; i < GITS_BASER_NR_REGS; i++) {
+		baser_val = host_state->tables[i].val;
+		if (!(baser_val & GITS_BASER_VALID))
+			continue;
+
+		original_table = kern_hyp_va(host_state->tables[i].base);
+		num_pages = (1 << host_state->tables[i].order);
+
+		ret = __pkvm_host_donate_hyp(hyp_virt_to_pfn(original_table), num_pages);
+		if (ret)
+			goto err_donate;
+
+		if (baser_val & GITS_BASER_INDIRECT) {
+			if (!host_state->tables[i].base_snapshot) {
+				ret = -EINVAL;
+				goto err_with_donation;
+			}
+
+			snapshot_table = kern_hyp_va(host_state->tables[i].base_snapshot);
+			ret = pkvm_share_shadow_table(snapshot_table, num_pages);
+			if (ret)
+				goto err_with_donation;
+
+			ret = pkvm_host_unmap_last_level(original_table, num_pages,
+							 host_state->tables[i].psz);
+			if (ret)
+				goto err_with_share;
+		}
+	}
+
+	return 0;
+err_with_share:
+	pkvm_unshare_shadow_table(snapshot_table, num_pages);
+err_with_donation:
+	__pkvm_hyp_donate_host(hyp_virt_to_pfn(original_table), num_pages);
+err_donate:
+	for (i = i - 1; i >= 0; i--) {
+		baser_val = host_state->tables[i].val;
+		if (!(baser_val & GITS_BASER_VALID))
+			continue;
+
+		original_table = kern_hyp_va(host_state->tables[i].base);
+		num_pages = (1 << host_state->tables[i].order);
+
+		if (baser_val & GITS_BASER_INDIRECT) {
+			snapshot_table = kern_hyp_va(host_state->tables[i].base_snapshot);
+			pkvm_unshare_shadow_table(snapshot_table, num_pages);
+
+			pkvm_host_map_last_level(original_table, num_pages,
+						 host_state->tables[i].psz);
+		}
+
+		WARN_ON(__pkvm_hyp_donate_host(hyp_virt_to_pfn(original_table), num_pages));
+	}
+
+	return ret;
+}
+
 DEFINE_HYP_SPINLOCK(its_setup_lock);
 
 int pkvm_its_emulate_setup(phys_addr_t dev_addr, struct its_host_state *host_state, void *priv,
@@ -294,6 +449,10 @@ int pkvm_its_emulate_setup(phys_addr_t dev_addr, struct its_host_state *host_sta
 	if (ret)
 		goto err_with_host_state;
 
+	ret = pkvm_setup_its_shadow_baser(host_state);
+	if (ret)
+		goto err_with_shadow_cmdq;
+
 	hyp_spin_lock_init(&priv_state->its_lock);
 
 	priv_state->host_state = host_state;
@@ -312,6 +471,8 @@ int pkvm_its_emulate_setup(phys_addr_t dev_addr, struct its_host_state *host_sta
 	hyp_spin_unlock(&its_setup_lock);
 
 	return 0;
+err_with_shadow_cmdq:
+	pkvm_teardown_its_shadow_cmdq(host_state);
 err_with_host_state:
 	WARN_ON(__pkvm_hyp_donate_host(hyp_virt_to_pfn(host_state), 1));
 err_with_priv:
-- 
2.55.0.654.g21b8a5bc05-goog

[PATCH v2 08/13] KVM: arm64: Trap & emulate the ITS MAPD command

[Sebastian Ene]

Parse the MAPD command and extract the ITT address to sanitize it. When
the command has the valid bit set, share and pin the memory that holds
the ITT table with the hypervisor to prevent it from being given to
someone else (eg. a VM). Use the pinning as a mechanism to get a grip
to the page and to prevent other users of the pKVM API from sharing or
donating the page for something else.
This is to prevent a a situation where a page is given to someone else
and then a MAPTI command is used to create an ITE entry in that page.

Implement shadow table updates for the first level of the indirect
tables when a MAPD command is issued. Compare the host view of the table
for the entry identified by the deviceId with the original table at the
same index and check if the valid bit is changed. If it didn't change,
don't update the original table. If it changed, verify if the new entry
has the valid bit set and donate the level2 table from the host to the
hypervisor (with the address of the table used from the new entry).
If the new entry has the valid bit cleared, donate the level2 table
from the hypervisor to the host with the address of the table extracted
from the original table managed by the hypervisor.

Signed-off-by: Sebastian Ene <sebastianene@google.com>
---
 arch/arm64/kvm/hyp/nvhe/its_emulate.c | 240 +++++++++++++++++++++++++-
 1 file changed, 238 insertions(+), 2 deletions(-)

diff --git a/arch/arm64/kvm/hyp/nvhe/its_emulate.c b/arch/arm64/kvm/hyp/nvhe/its_emulate.c
index 1ce2f9d8fcf9..071a08d3602d 100644
--- a/arch/arm64/kvm/hyp/nvhe/its_emulate.c
+++ b/arch/arm64/kvm/hyp/nvhe/its_emulate.c
@@ -54,6 +54,11 @@ struct its_handler {
 	.read = (read_cb),				\
 }
 
+struct dte_entry {
+	u32	device_id;
+	u64	itt_pfn;
+};
+
 struct its_priv_state {
 	/* The location of the ITS in the hypervisor VA */
 	void __iomem	*base;
@@ -66,6 +71,9 @@ struct its_priv_state {
 	hyp_spinlock_t	its_lock;
 
 	struct its_host_state	*host_state;
+	u16			empty_entry;
+	u16			num_tracked_entries;
+	struct dte_entry	tracked_entries[];
 };
 
 #define GITS_CWRITER_RETRY	BIT_ULL(0)
@@ -110,11 +118,236 @@ static int submit_single_cmd(struct its_priv_state *its, bool retry)
 	return 0;
 }
 
+static int get_num_itt_pages(struct its_priv_state *its, u8 num_bits)
+{
+	u64 gits_typer, nr_ites;
+	size_t sz;
+
+	gits_typer = readq_relaxed(its->base + GITS_TYPER);
+	if (num_bits > FIELD_GET(GITS_TYPER_IDBITS, gits_typer))
+		return -EINVAL;
+
+	nr_ites = BIT_ULL(num_bits + 1);
+	sz = nr_ites * (FIELD_GET(GITS_TYPER_ITT_ENTRY_SIZE, gits_typer) + 1);
+	sz = max(sz, ITS_ITT_ALIGN) + ITS_ITT_ALIGN - 1;
+
+	return PAGE_ALIGN(sz) >> PAGE_SHIFT;
+}
+
+static struct its_baser *get_table_from_snapshot(struct its_host_state *host, u64 baser_type)
+{
+	int i;
+
+	for (i = 0; i < GITS_BASER_NR_REGS; i++) {
+		if (GITS_BASER_TYPE(host->tables[i].val) == baser_type)
+			return &host->tables[i];
+	}
+
+	return NULL;
+}
+
+static int check_table_update(struct its_priv_state *its, u32 device_id, u64 type, bool rollback)
+{
+	struct its_baser *table = get_table_from_snapshot(its->host_state, type);
+	size_t lvl2_entry_sz, lvl1_table_sz, num_lvl2_entries, num_lvl1_entries;
+	u64 *snapshot_table, *original_table;
+	u64 prev_entry, new_entry;
+	u32 new_entry_index;
+	int ret;
+
+	if (!table)
+		return -EINVAL;
+
+	/* We only do shadow udates for the first level of indirect tables */
+	if (!(table->val & GITS_BASER_INDIRECT))
+		return 0;
+
+	lvl2_entry_sz = GITS_BASER_ENTRY_SIZE(table->val);
+	num_lvl2_entries = table->psz / lvl2_entry_sz;
+
+	lvl1_table_sz = (1 << table->order) << PAGE_SHIFT;
+	num_lvl1_entries = lvl1_table_sz / sizeof(u64);
+
+	new_entry_index = device_id / num_lvl2_entries;
+	if (new_entry_index >= num_lvl1_entries)
+		return -ENOSPC;
+
+	snapshot_table = kern_hyp_va(table->base_snapshot);
+	original_table = kern_hyp_va(table->base);
+
+	/*
+	 * Look at the host table copy and if the entry hasn't changed the valid
+	 * bit compared to the original table used by the hardwre, don't update anything.
+	 */
+	new_entry = snapshot_table[new_entry_index];
+	prev_entry = original_table[new_entry_index];
+	if (!((new_entry ^ prev_entry) & GITS_BASER_VALID))
+		return 0;
+
+	/*
+	 * The host can play nasty tricks with read-modify-write after a
+	 * rollback is triggered but we still hold on to the original tables
+	 * which are hyp managed and we don't give back any other page to the
+	 * host.
+	 */
+	if (rollback)
+		new_entry = new_entry ^ GITS_BASER_VALID;
+
+	if (new_entry & GITS_BASER_VALID)
+		ret = __pkvm_host_donate_hyp(hyp_phys_to_pfn(new_entry & PHYS_MASK),
+					     table->psz >> PAGE_SHIFT);
+	else
+		ret = __pkvm_hyp_donate_host(hyp_phys_to_pfn(prev_entry & PHYS_MASK),
+					     table->psz >> PAGE_SHIFT);
+	if (ret)
+		return ret;
+
+	original_table[new_entry_index] = new_entry;
+	return 0;
+}
+
+static int track_pfn_add(struct its_priv_state *its, u32 device_id, u64 pfn)
+{
+	void *virt = hyp_phys_to_virt(hyp_pfn_to_phys(pfn));
+	struct dte_entry *entries = &its->tracked_entries[0];
+	bool pfn_shared = false;
+	int ret;
+	int i;
+
+	for (i = 0; i < its->num_tracked_entries; i++) {
+		if (entries[i].itt_pfn == pfn) {
+			if (entries[i].device_id != device_id) {
+				pfn_shared = true;
+				break;
+			} else {
+				return hyp_pin_shared_mem(virt, virt + PAGE_SIZE);
+			}
+		}
+	}
+
+	if (its->empty_entry >= its->num_tracked_entries)
+		return -ENOSPC;
+
+	if (!pfn_shared) {
+		ret = __pkvm_host_share_hyp(pfn);
+		if (ret)
+			return ret;
+	}
+
+	ret = hyp_pin_shared_mem(virt, virt + PAGE_SIZE);
+	if (ret) {
+		__pkvm_host_unshare_hyp(pfn);
+		return ret;
+	}
+
+	entries[its->empty_entry].itt_pfn = pfn;
+	entries[its->empty_entry].device_id = device_id;
+
+	for (i = 0; i < its->num_tracked_entries; i++) {
+		if (!entries[i].itt_pfn && !entries[i].device_id)
+			break;
+	}
+	its->empty_entry = i;
+	return 0;
+}
+
+static int track_pfn_remove(struct its_priv_state *its, u32 device_id, u64 pfn)
+{
+	void *virt = hyp_phys_to_virt(hyp_pfn_to_phys(pfn));
+	struct dte_entry *entries = &its->tracked_entries[0];
+	int ret;
+	int i;
+
+	for (i = 0; i < its->num_tracked_entries; i++) {
+		if (entries[i].itt_pfn != pfn || entries[i].device_id != device_id)
+			continue;
+
+		/* To decrement the refcount, first try to unshare it */
+		ret = __pkvm_host_unshare_hyp(pfn);
+		if (ret == -EBUSY) {
+			hyp_unpin_shared_mem(virt, virt + PAGE_SIZE);
+			ret = __pkvm_host_unshare_hyp(pfn);
+			if (ret == -EBUSY)
+				return 0;
+
+			WARN_ON(ret);
+		}
+
+		memset(&entries[i], 0, sizeof(struct dte_entry));
+		its->empty_entry = i;
+		return 0;
+	}
+
+	return -EINVAL;
+}
+
+static int track_pfn(struct its_priv_state *its, u32 device_id, u64 pfn, int num_pages,
+		     bool remove)
+{
+	int ret;
+	int i;
+
+	for (i = 0; i < num_pages; i++) {
+		if (remove)
+			ret = track_pfn_remove(its, device_id, pfn + i);
+		else
+			ret = track_pfn_add(its, device_id, pfn + i);
+
+		if (ret)
+			goto err_track_pfn;
+	}
+
+	return 0;
+err_track_pfn:
+	for (i = i - 1; i >= 0; i--) {
+		if (remove)
+			WARN_ON(track_pfn_add(its, device_id, pfn + i));
+		else
+			WARN_ON(track_pfn_remove(its, device_id, pfn + i));
+	}
+	return ret;
+}
+
+static int process_its_mapd(struct its_priv_state *its, struct its_cmd_block *cmd, bool rollback)
+{
+	phys_addr_t itt_addr = cmd->raw_cmd[2] & GENMASK(51, 8);
+	bool remove = !(cmd->raw_cmd[2] & BIT(63));
+	u8 size = cmd->raw_cmd[1] & GENMASK(4, 0);
+	u32 device_id = cmd->raw_cmd[0] >> 32;
+	int num_pages, ret;
+	u64 itt_pfn;
+
+	if (rollback)
+		remove = !remove;
+
+	itt_pfn = hyp_phys_to_pfn(itt_addr);
+	num_pages = get_num_itt_pages(its, size);
+	if (num_pages < 0)
+		return num_pages;
+
+	ret = check_table_update(its, device_id, GITS_BASER_TYPE_DEVICE, rollback);
+	if (ret)
+		return ret;
+
+	return track_pfn(its, device_id, itt_pfn, num_pages, remove);
+}
+
 static int process_cmd(struct its_priv_state *its, struct its_cmd_block *cmd,
 		       bool rollback)
 {
-	/* Passthrough everything for now */
-	return 0;
+	u8 req_type = cmd->raw_cmd[0] & GENMASK_ULL(7, 0);
+	int ret = 0;
+
+	switch (req_type) {
+	case GITS_CMD_MAPD:
+		ret = process_its_mapd(its, cmd, rollback);
+		break;
+	default:
+		/* Passthrough everything for now */
+		break;
+	}
+
+	return ret;
 }
 
 static void cwriter_write(struct pkvm_protected_reg *region, u64 offset, u64 value)
@@ -459,6 +692,9 @@ int pkvm_its_emulate_setup(phys_addr_t dev_addr, struct its_host_state *host_sta
 	priv_state->base = (void __iomem *)__hyp_va(dev_addr);
 	priv_state->cmd_original = host_state->cmd_original;
 	priv_state->cmd_host_copy = host_state->cmd_host_copy;
+	priv_state->empty_entry = 0;
+	priv_state->num_tracked_entries = ((priv_num_pages << PAGE_SHIFT) -
+		offsetof(struct its_priv_state, tracked_entries)) / sizeof(struct dte_entry);
 
 	priv_state->cmd_offset = readq_relaxed(priv_state->base + GITS_CREADR) &
 		GITS_CREADR_OFFSET;
-- 
2.55.0.654.g21b8a5bc05-goog

[PATCH v2 09/13] KVM: arm64: Trap & emulate the ITS MAPC command

[Sebastian Ene]

Parse the MAPC command and verify if we need to do any updates to the
shadow collection table.

Signed-off-by: Sebastian Ene <sebastianene@google.com>
---
 arch/arm64/kvm/hyp/nvhe/its_emulate.c | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/arch/arm64/kvm/hyp/nvhe/its_emulate.c b/arch/arm64/kvm/hyp/nvhe/its_emulate.c
index 071a08d3602d..5629e2a070df 100644
--- a/arch/arm64/kvm/hyp/nvhe/its_emulate.c
+++ b/arch/arm64/kvm/hyp/nvhe/its_emulate.c
@@ -332,6 +332,13 @@ static int process_its_mapd(struct its_priv_state *its, struct its_cmd_block *cm
 	return track_pfn(its, device_id, itt_pfn, num_pages, remove);
 }
 
+static int process_its_mapc(struct its_priv_state *its, struct its_cmd_block *cmd, bool rollback)
+{
+	u32 icid = cmd->raw_cmd[2] & GENMASK(15, 0);
+
+	return check_table_update(its, icid, GITS_BASER_TYPE_COLLECTION, rollback);
+}
+
 static int process_cmd(struct its_priv_state *its, struct its_cmd_block *cmd,
 		       bool rollback)
 {
@@ -342,6 +349,10 @@ static int process_cmd(struct its_priv_state *its, struct its_cmd_block *cmd,
 	case GITS_CMD_MAPD:
 		ret = process_its_mapd(its, cmd, rollback);
 		break;
+
+	case GITS_CMD_MAPC:
+		ret = process_its_mapc(its, cmd, rollback);
+		break;
 	default:
 		/* Passthrough everything for now */
 		break;
-- 
2.55.0.654.g21b8a5bc05-goog

[PATCH v2 10/13] KVM: arm64: Restrict host updates to GITS_CTLR

[Sebastian Ene]

Prevent unpredictable hardware behavior when the host tries to enable
the ITS while it is not in quiescent state.

Signed-off-by: Sebastian Ene <sebastianene@google.com>
---
 arch/arm64/kvm/hyp/nvhe/its_emulate.c | 27 +++++++++++++++++++++++++++
 1 file changed, 27 insertions(+)

diff --git a/arch/arm64/kvm/hyp/nvhe/its_emulate.c b/arch/arm64/kvm/hyp/nvhe/its_emulate.c
index 5629e2a070df..b9b71aa18d48 100644
--- a/arch/arm64/kvm/hyp/nvhe/its_emulate.c
+++ b/arch/arm64/kvm/hyp/nvhe/its_emulate.c
@@ -403,8 +403,35 @@ static void cwriter_read(struct pkvm_protected_reg *region, u64 offset, u64 *rea
 	*read = readq_relaxed(its->base + GITS_CWRITER);
 }
 
+static void ctlr_read(struct pkvm_protected_reg *region, u64 offset, u64 *read)
+{
+	struct its_priv_state *its = region->priv;
+	*read = readl_relaxed(its->base + GITS_CTLR);
+}
+
+static void ctlr_write(struct pkvm_protected_reg *region, u64 offset, u64 value)
+{
+	struct its_priv_state *its = region->priv;
+	bool is_quiescent, is_enabled;
+	u32 ctlr;
+
+	ctlr = readl_relaxed(its->base + GITS_CTLR);
+	is_quiescent = !!(ctlr & GITS_CTLR_QUIESCENT);
+	is_enabled = !!(ctlr & GITS_CTLR_ENABLE);
+
+	/*
+	 * If it's disabled and not in quiescent state and it tries to enable
+	 * it, bail out.
+	 */
+	if (!is_enabled && (value & GITS_CTLR_ENABLE) && !is_quiescent)
+		return;
+
+	writel_relaxed(value, its->base + GITS_CTLR);
+}
+
 static struct its_handler its_handlers[] = {
 	ITS_HANDLER(GITS_CWRITER, sizeof(u64), cwriter_write, cwriter_read),
+	ITS_HANDLER(GITS_CTLR, sizeof(u32), ctlr_write, ctlr_read),
 	{},
 };
 
-- 
2.55.0.654.g21b8a5bc05-goog

[PATCH v2 11/13] KVM: arm64: Prevent the host from specifying a different command queue

[Sebastian Ene]

Don't let a malicious host re-program the command queue register with a
different address and size to bypass the commands sanitization.
Prevent unpredictable hardware behavior and restrict updates to the
GITS_CBASER while the ITS is enabled or not in a quiescent state.

Signed-off-by: Sebastian Ene <sebastianene@google.com>
---
 arch/arm64/kvm/hyp/nvhe/its_emulate.c | 32 +++++++++++++++++++++++++++
 1 file changed, 32 insertions(+)

diff --git a/arch/arm64/kvm/hyp/nvhe/its_emulate.c b/arch/arm64/kvm/hyp/nvhe/its_emulate.c
index b9b71aa18d48..97cfa31d90d1 100644
--- a/arch/arm64/kvm/hyp/nvhe/its_emulate.c
+++ b/arch/arm64/kvm/hyp/nvhe/its_emulate.c
@@ -429,9 +429,41 @@ static void ctlr_write(struct pkvm_protected_reg *region, u64 offset, u64 value)
 	writel_relaxed(value, its->base + GITS_CTLR);
 }
 
+static void cbaser_write(struct pkvm_protected_reg *region, u64 offset, u64 value)
+{
+	struct its_priv_state *its = region->priv;
+	int num_pages;
+	u64 ctlr;
+
+	ctlr = readl_relaxed(its->base + GITS_CTLR);
+	if ((ctlr & GITS_CTLR_ENABLE) || !(ctlr & GITS_CTLR_QUIESCENT))
+		return;
+
+	num_pages = its->host_state->cmdq_len / SZ_4K;
+
+	/* Don't let the host program a different command queue */
+	value &= ~(GENMASK(7, 0) | GENMASK_ULL(51, 12));
+	value |= (num_pages - 1) & GENMASK(7, 0);
+	value |= __hyp_pa(its->cmd_original) & GENMASK_ULL(51, 12);
+	its->needs_flush = (value & GITS_CBASER_SHAREABILITY_MASK) != GITS_CBASER_InnerShareable;
+
+	writeq_relaxed(value, its->base + GITS_CBASER);
+
+	/* Restart the CMDQ to read from 0 */
+	its->cmd_offset = 0;
+	writeq_relaxed(0, its->base + GITS_CWRITER);
+}
+
+static void cbaser_read(struct pkvm_protected_reg *region, u64 offset, u64 *read)
+{
+	struct its_priv_state *its = region->priv;
+	*read = readq_relaxed(its->base + GITS_CBASER);
+}
+
 static struct its_handler its_handlers[] = {
 	ITS_HANDLER(GITS_CWRITER, sizeof(u64), cwriter_write, cwriter_read),
 	ITS_HANDLER(GITS_CTLR, sizeof(u32), ctlr_write, ctlr_read),
+	ITS_HANDLER(GITS_CBASER, sizeof(u64), cbaser_write, cbaser_read),
 	{},
 };
 
-- 
2.55.0.654.g21b8a5bc05-goog

[PATCH v2 12/13] KVM: arm64: Prevent the host from programming new GITS_BASER tables

[Sebastian Ene]

Don't allow the host to change the layout of the tables or to modify the
address programmed in the GITS_BASER registers to point to new tables
and bypass the sanitization.
Prevent the host from updating the ITS tables while the ITS is enabled
and the tables are set to prevent undefined behavior.

Signed-off-by: Sebastian Ene <sebastianene@google.com>
---
 arch/arm64/kvm/hyp/nvhe/its_emulate.c | 44 +++++++++++++++++++++++++--
 1 file changed, 41 insertions(+), 3 deletions(-)

diff --git a/arch/arm64/kvm/hyp/nvhe/its_emulate.c b/arch/arm64/kvm/hyp/nvhe/its_emulate.c
index 97cfa31d90d1..82dc60dcde68 100644
--- a/arch/arm64/kvm/hyp/nvhe/its_emulate.c
+++ b/arch/arm64/kvm/hyp/nvhe/its_emulate.c
@@ -42,18 +42,23 @@ void its_emulate_forward_req(struct pkvm_protected_reg *region, u64 offset, bool
 struct its_handler {
 	u64	offset;
 	u8	access_size;
+	u8	num_registers;
 	void	(*write)(struct pkvm_protected_reg *region, u64 offset, u64 value);
 	void	(*read)(struct pkvm_protected_reg *region, u64 offset, u64 *read);
 };
 
-#define ITS_HANDLER(off, sz, write_cb, read_cb)		\
+#define ITS_HANDLER_REG_PAIR(off, sz, registers, write_cb, read_cb)	\
 {							\
 	.offset = (off),				\
 	.access_size = (sz),				\
+	.num_registers = (registers),			\
 	.write = (write_cb),				\
 	.read = (read_cb),				\
 }
 
+#define ITS_HANDLER(off, sz, write_cb, read_cb)		\
+	ITS_HANDLER_REG_PAIR(off, sz, 1, write_cb, read_cb)
+
 struct dte_entry {
 	u32	device_id;
 	u64	itt_pfn;
@@ -460,10 +465,42 @@ static void cbaser_read(struct pkvm_protected_reg *region, u64 offset, u64 *read
 	*read = readq_relaxed(its->base + GITS_CBASER);
 }
 
+static void baser_write(struct pkvm_protected_reg *region, u64 offset, u64 value)
+{
+	struct its_priv_state *its = region->priv;
+	u32 ctlr = readl_relaxed(its->base + GITS_CTLR);
+	int baser_idx;
+	u64 baser;
+
+	if ((ctlr & GITS_CTLR_ENABLE) || !(ctlr & GITS_CTLR_QUIESCENT))
+		return;
+
+	baser_idx = (offset - GITS_BASER) >> 3;
+	baser = its->host_state->tables[baser_idx].val;
+
+	/* Prevent if it tries to change from direct layout to indirect layout */
+	if ((value & GITS_BASER_INDIRECT) != (baser & GITS_BASER_INDIRECT))
+		return;
+
+	/* Don't allow the host to point to new tables or new attributes */
+	value &= ~(GENMASK_ULL(47, 12) | GENMASK_ULL(9, 0));
+	value |= (baser & GENMASK_ULL(47, 12)) | (baser & GENMASK_ULL(9, 0));
+
+	writeq_relaxed(value, its->base + offset);
+}
+
+static void baser_read(struct pkvm_protected_reg *region, u64 offset, u64 *read)
+{
+	struct its_priv_state *its = region->priv;
+	*read = readq_relaxed(its->base + offset);
+}
+
 static struct its_handler its_handlers[] = {
 	ITS_HANDLER(GITS_CWRITER, sizeof(u64), cwriter_write, cwriter_read),
 	ITS_HANDLER(GITS_CTLR, sizeof(u32), ctlr_write, ctlr_read),
 	ITS_HANDLER(GITS_CBASER, sizeof(u64), cbaser_write, cbaser_read),
+
+	ITS_HANDLER_REG_PAIR(GITS_BASER, sizeof(u64), 8, baser_write, baser_read),
 	{},
 };
 
@@ -472,13 +509,14 @@ void pkvm_its_emulate_handler(struct pkvm_protected_reg *region, u64 offset, boo
 {
 	struct its_priv_state *priv = region->priv;
 	struct its_handler *reg_handler;
+	u64 end;
 
 	if (!priv || !IS_ALIGNED(offset, reg_size))
 		return;
 
 	for (reg_handler = its_handlers; reg_handler->access_size; reg_handler++) {
-		if (reg_handler->offset > offset ||
-		    reg_handler->offset + reg_handler->access_size <= offset)
+		end = reg_handler->offset + reg_handler->access_size * reg_handler->num_registers;
+		if (reg_handler->offset > offset || end <= offset)
 			continue;
 
 		if (reg_handler->access_size < reg_size)
-- 
2.55.0.654.g21b8a5bc05-goog

[PATCH v2 13/13] KVM: arm64: Implement HVC interface for ITS emulation setup

[Sebastian Ene]

Introduce a new HVC to allow the host to trigger the ITS emulation
setup. Use the introduced API in the GIC ITS driver to call the driver
to lock the ITS before pKVM finalize and to prepare for emulation setup.
On the return path from the pKVM finalize, call into the driver to
release the ITS locks which performs a switch in the driver to use a
different command queue and a different set of level-1 indirect tables.
Allocate memory that will be used by the emulation to track the internal
state and send the snapshot state from the driver.
Replace the initial "trap-and-forward" MMIO handler with a full-featured
emulation handler.

Signed-off-by: Sebastian Ene <sebastianene@google.com>
---
 arch/arm64/include/asm/kvm_asm.h      |  1 +
 arch/arm64/include/asm/kvm_pkvm.h     |  4 ++--
 arch/arm64/kvm/hyp/nvhe/hyp-main.c    | 16 ++++++++++++++++
 arch/arm64/kvm/hyp/nvhe/its_emulate.c |  4 ++--
 arch/arm64/kvm/pkvm.c                 | 26 ++++++++++++++++++++++++--
 5 files changed, 45 insertions(+), 6 deletions(-)

diff --git a/arch/arm64/include/asm/kvm_asm.h b/arch/arm64/include/asm/kvm_asm.h
index 043495f7fc78..fcb2871b8a86 100644
--- a/arch/arm64/include/asm/kvm_asm.h
+++ b/arch/arm64/include/asm/kvm_asm.h
@@ -114,6 +114,7 @@ enum __kvm_host_smccc_func {
 	__KVM_HOST_SMCCC_FUNC___pkvm_vcpu_load,
 	__KVM_HOST_SMCCC_FUNC___pkvm_vcpu_put,
 	__KVM_HOST_SMCCC_FUNC___pkvm_tlb_flush_vmid,
+	__KVM_HOST_SMCCC_FUNC___pkvm_its_emulate_setup,
 
 	MARKER(__KVM_HOST_SMCCC_FUNC_MAX)
 };
diff --git a/arch/arm64/include/asm/kvm_pkvm.h b/arch/arm64/include/asm/kvm_pkvm.h
index 78597210a53c..cc89e2bde468 100644
--- a/arch/arm64/include/asm/kvm_pkvm.h
+++ b/arch/arm64/include/asm/kvm_pkvm.h
@@ -32,8 +32,8 @@ struct pkvm_protected_reg {
 
 extern struct pkvm_protected_reg kvm_nvhe_sym(pkvm_protected_regs)[];
 extern unsigned int kvm_nvhe_sym(num_protected_reg);
-extern void kvm_nvhe_sym(its_emulate_forward_req)(struct pkvm_protected_reg *region, u64 offset,
-						  bool write, u64 *reg, u8 reg_size);
+extern void kvm_nvhe_sym(pkvm_its_emulate_handler)(struct pkvm_protected_reg *region, u64 offset,
+						   bool write, u64 *reg, u8 reg_size);
 
 int pkvm_init_host_vm(struct kvm *kvm, unsigned long type);
 int pkvm_create_hyp_vm(struct kvm *kvm);
diff --git a/arch/arm64/kvm/hyp/nvhe/hyp-main.c b/arch/arm64/kvm/hyp/nvhe/hyp-main.c
index d3df96ed8ba4..ad57b2076eee 100644
--- a/arch/arm64/kvm/hyp/nvhe/hyp-main.c
+++ b/arch/arm64/kvm/hyp/nvhe/hyp-main.c
@@ -16,6 +16,7 @@
 #include <asm/kvm_mmu.h>
 
 #include <nvhe/ffa.h>
+#include <nvhe/its_emulate.h>
 #include <nvhe/mem_protect.h>
 #include <nvhe/mm.h>
 #include <nvhe/pkvm.h>
@@ -705,6 +706,20 @@ static void handle___vgic_v5_restore_vmcr_apr(struct kvm_cpu_context *host_ctxt)
 	__vgic_v5_restore_vmcr_apr(kern_hyp_va(cpu_if));
 }
 
+static void handle___pkvm_its_emulate_setup(struct kvm_cpu_context *host_ctxt)
+{
+	DECLARE_REG(phys_addr_t, dev_addr, host_ctxt, 1);
+	DECLARE_REG(struct its_host_state *, host_state, host_ctxt, 2);
+	DECLARE_REG(void *, priv_state, host_ctxt, 3);
+	DECLARE_REG(size_t, priv_state_num_pages, host_ctxt, 4);
+
+	if (!is_protected_kvm_enabled())
+		return;
+
+	cpu_reg(host_ctxt, 1) = pkvm_its_emulate_setup(dev_addr, host_state, priv_state,
+						       priv_state_num_pages);
+}
+
 typedef void (*hcall_t)(struct kvm_cpu_context *);
 
 #define HANDLE_FUNC(x)	[__KVM_HOST_SMCCC_FUNC_##x] = (hcall_t)handle_##x
@@ -762,6 +777,7 @@ static const hcall_t host_hcall[] = {
 	HANDLE_FUNC(__pkvm_vcpu_load),
 	HANDLE_FUNC(__pkvm_vcpu_put),
 	HANDLE_FUNC(__pkvm_tlb_flush_vmid),
+	HANDLE_FUNC(__pkvm_its_emulate_setup),
 };
 
 static void handle_host_hcall(struct kvm_cpu_context *host_ctxt)
diff --git a/arch/arm64/kvm/hyp/nvhe/its_emulate.c b/arch/arm64/kvm/hyp/nvhe/its_emulate.c
index 82dc60dcde68..8c8acaee4d2b 100644
--- a/arch/arm64/kvm/hyp/nvhe/its_emulate.c
+++ b/arch/arm64/kvm/hyp/nvhe/its_emulate.c
@@ -6,8 +6,8 @@
 
 #include <linux/irqchip/arm-gic-v3.h>
 
-void its_emulate_forward_req(struct pkvm_protected_reg *region, u64 offset, bool write, u64 *reg,
-			     u8 reg_size)
+static void its_emulate_forward_req(struct pkvm_protected_reg *region, u64 offset, bool write,
+				    u64 *reg, u8 reg_size)
 {
 	void __iomem *addr = __hyp_va(PFN_PHYS(region->pfn) + offset);
 
diff --git a/arch/arm64/kvm/pkvm.c b/arch/arm64/kvm/pkvm.c
index 4bfffbedac4c..a9ceb9ffe6a4 100644
--- a/arch/arm64/kvm/pkvm.c
+++ b/arch/arm64/kvm/pkvm.c
@@ -71,7 +71,7 @@ static int __init register_its_emulated_region(void)
 		 */
 		kvm_nvhe_sym(pkvm_protected_regs)[i].pfn = PHYS_PFN(res.start);
 		kvm_nvhe_sym(pkvm_protected_regs)[i].cb =
-			lm_alias(&kvm_nvhe_sym(its_emulate_forward_req));
+			lm_alias(&kvm_nvhe_sym(pkvm_its_emulate_handler));
 		kvm_nvhe_sym(pkvm_protected_regs)[i].nr_pages =
 			PFN_DOWN(min_t(u64, resource_size(&res), PAGE_ALIGN_DOWN(GITS_TRANSLATER)));
 
@@ -312,8 +312,28 @@ static void __init _kvm_host_prot_finalize(void *arg)
 		WRITE_ONCE(*err, -EINVAL);
 }
 
+#define ITS_PAGES	(2UL)
+
+static int pkvm_init_its_emulation(phys_addr_t dev_addr, struct its_host_state *host)
+{
+	size_t priv_state_sz = ITS_PAGES << PAGE_SHIFT;
+	void *priv_state;
+	int ret;
+
+	priv_state = alloc_pages_exact(priv_state_sz, GFP_ATOMIC);
+	if (!priv_state)
+		return -ENOMEM;
+
+	ret = kvm_call_hyp_nvhe(__pkvm_its_emulate_setup, dev_addr, host, priv_state, ITS_PAGES);
+	if (ret)
+		free_pages_exact(priv_state, priv_state_sz);
+
+	return ret;
+}
+
 static int __init pkvm_drop_host_privileges(void)
 {
+	unsigned long its_flags;
 	int ret = 0;
 
 	/*
@@ -321,8 +341,10 @@ static int __init pkvm_drop_host_privileges(void)
 	 * once the host stage 2 is installed.
 	 */
 	static_branch_enable(&kvm_protected_mode_initialized);
+
+	its_emulate_acquire_locks(&its_flags);
 	on_each_cpu(_kvm_host_prot_finalize, &ret, 1);
-	return ret;
+	return its_emulate_release_locks(ret, &its_flags, pkvm_init_its_emulation);
 }
 
 static int __init finalize_pkvm(void)
-- 
2.55.0.654.g21b8a5bc05-goog