#include <linux/mm.h>
 
 enum __kvm_host_smccc_func {
+       /* Hypercalls available only prior to pKVM finalisation */
        /* __KVM_HOST_SMCCC_FUNC___kvm_hyp_init */
-       __KVM_HOST_SMCCC_FUNC___kvm_vcpu_run = __KVM_HOST_SMCCC_FUNC___kvm_hyp_init + 1,
+       __KVM_HOST_SMCCC_FUNC___kvm_get_mdcr_el2 = __KVM_HOST_SMCCC_FUNC___kvm_hyp_init + 1,
+       __KVM_HOST_SMCCC_FUNC___pkvm_init,
+       __KVM_HOST_SMCCC_FUNC___pkvm_create_private_mapping,
+       __KVM_HOST_SMCCC_FUNC___pkvm_cpu_set_vector,
+       __KVM_HOST_SMCCC_FUNC___kvm_enable_ssbs,
+       __KVM_HOST_SMCCC_FUNC___vgic_v3_init_lrs,
+       __KVM_HOST_SMCCC_FUNC___vgic_v3_get_gic_config,
+       __KVM_HOST_SMCCC_FUNC___pkvm_prot_finalize,
+
+       /* Hypercalls available after pKVM finalisation */
+       __KVM_HOST_SMCCC_FUNC___pkvm_host_share_hyp,
+       __KVM_HOST_SMCCC_FUNC___kvm_adjust_pc,
+       __KVM_HOST_SMCCC_FUNC___kvm_vcpu_run,
        __KVM_HOST_SMCCC_FUNC___kvm_flush_vm_context,
        __KVM_HOST_SMCCC_FUNC___kvm_tlb_flush_vmid_ipa,
        __KVM_HOST_SMCCC_FUNC___kvm_tlb_flush_vmid,
        __KVM_HOST_SMCCC_FUNC___kvm_flush_cpu_context,
        __KVM_HOST_SMCCC_FUNC___kvm_timer_set_cntvoff,
-       __KVM_HOST_SMCCC_FUNC___kvm_enable_ssbs,
-       __KVM_HOST_SMCCC_FUNC___vgic_v3_get_gic_config,
        __KVM_HOST_SMCCC_FUNC___vgic_v3_read_vmcr,
        __KVM_HOST_SMCCC_FUNC___vgic_v3_write_vmcr,
-       __KVM_HOST_SMCCC_FUNC___vgic_v3_init_lrs,
-       __KVM_HOST_SMCCC_FUNC___kvm_get_mdcr_el2,
        __KVM_HOST_SMCCC_FUNC___vgic_v3_save_aprs,
        __KVM_HOST_SMCCC_FUNC___vgic_v3_restore_aprs,
-       __KVM_HOST_SMCCC_FUNC___pkvm_init,
-       __KVM_HOST_SMCCC_FUNC___pkvm_host_share_hyp,
-       __KVM_HOST_SMCCC_FUNC___pkvm_create_private_mapping,
-       __KVM_HOST_SMCCC_FUNC___pkvm_cpu_set_vector,
-       __KVM_HOST_SMCCC_FUNC___pkvm_prot_finalize,
-       __KVM_HOST_SMCCC_FUNC___kvm_adjust_pc,
 };
 
 #define DECLARE_KVM_VHE_SYM(sym)       extern char sym[]
 
 #define HANDLE_FUNC(x) [__KVM_HOST_SMCCC_FUNC_##x] = (hcall_t)handle_##x
 
 static const hcall_t host_hcall[] = {
-       HANDLE_FUNC(__kvm_vcpu_run),
+       /* ___kvm_hyp_init */
+       HANDLE_FUNC(__kvm_get_mdcr_el2),
+       HANDLE_FUNC(__pkvm_init),
+       HANDLE_FUNC(__pkvm_create_private_mapping),
+       HANDLE_FUNC(__pkvm_cpu_set_vector),
+       HANDLE_FUNC(__kvm_enable_ssbs),
+       HANDLE_FUNC(__vgic_v3_init_lrs),
+       HANDLE_FUNC(__vgic_v3_get_gic_config),
+       HANDLE_FUNC(__pkvm_prot_finalize),
+
+       HANDLE_FUNC(__pkvm_host_share_hyp),
        HANDLE_FUNC(__kvm_adjust_pc),
+       HANDLE_FUNC(__kvm_vcpu_run),
        HANDLE_FUNC(__kvm_flush_vm_context),
        HANDLE_FUNC(__kvm_tlb_flush_vmid_ipa),
        HANDLE_FUNC(__kvm_tlb_flush_vmid),
        HANDLE_FUNC(__kvm_flush_cpu_context),
        HANDLE_FUNC(__kvm_timer_set_cntvoff),
-       HANDLE_FUNC(__kvm_enable_ssbs),
-       HANDLE_FUNC(__vgic_v3_get_gic_config),
        HANDLE_FUNC(__vgic_v3_read_vmcr),
        HANDLE_FUNC(__vgic_v3_write_vmcr),
-       HANDLE_FUNC(__vgic_v3_init_lrs),
-       HANDLE_FUNC(__kvm_get_mdcr_el2),
        HANDLE_FUNC(__vgic_v3_save_aprs),
        HANDLE_FUNC(__vgic_v3_restore_aprs),
-       HANDLE_FUNC(__pkvm_init),
-       HANDLE_FUNC(__pkvm_cpu_set_vector),
-       HANDLE_FUNC(__pkvm_host_share_hyp),
-       HANDLE_FUNC(__pkvm_create_private_mapping),
-       HANDLE_FUNC(__pkvm_prot_finalize),
 };
 
 static void handle_host_hcall(struct kvm_cpu_context *host_ctxt)
 {
        DECLARE_REG(unsigned long, id, host_ctxt, 0);
+       unsigned long hcall_min = 0;
        hcall_t hfn;
 
+       /*
+        * If pKVM has been initialised then reject any calls to the
+        * early "privileged" hypercalls. Note that we cannot reject
+        * calls to __pkvm_prot_finalize for two reasons: (1) The static
+        * key used to determine initialisation must be toggled prior to
+        * finalisation and (2) finalisation is performed on a per-CPU
+        * basis. This is all fine, however, since __pkvm_prot_finalize
+        * returns -EPERM after the first call for a given CPU.
+        */
+       if (static_branch_unlikely(&kvm_protected_mode_initialized))
+               hcall_min = __KVM_HOST_SMCCC_FUNC___pkvm_prot_finalize;
+
        id -= KVM_HOST_SMCCC_ID(0);
 
-       if (unlikely(id >= ARRAY_SIZE(host_hcall)))
+       if (unlikely(id < hcall_min || id >= ARRAY_SIZE(host_hcall)))
                goto inval;
 
        hfn = host_hcall[id];