#define SATP_ASID_SHIFT                                44
 #define SATP_ASID_MASK                         _AC(0xFFFF, UL)
 
-#define SBI_EXT_EXPERIMENTAL_START     0x08000000
-#define SBI_EXT_EXPERIMENTAL_END       0x08FFFFFF
+#define SBI_EXT_EXPERIMENTAL_START             0x08000000
+#define SBI_EXT_EXPERIMENTAL_END               0x08FFFFFF
 
-#define KVM_RISCV_SELFTESTS_SBI_EXT    SBI_EXT_EXPERIMENTAL_END
+#define KVM_RISCV_SELFTESTS_SBI_EXT            SBI_EXT_EXPERIMENTAL_END
+#define KVM_RISCV_SELFTESTS_SBI_UCALL          0
+#define KVM_RISCV_SELFTESTS_SBI_UNEXP          1
 
 struct sbiret {
        long error;
 
                core.regs.t3, core.regs.t4, core.regs.t5, core.regs.t6);
 }
 
-static void __aligned(16) guest_hang(void)
+static void __aligned(16) guest_unexp_trap(void)
 {
-       while (1)
-               ;
+       sbi_ecall(KVM_RISCV_SELFTESTS_SBI_EXT,
+                 KVM_RISCV_SELFTESTS_SBI_UNEXP,
+                 0, 0, 0, 0, 0, 0);
 }
 
 void vm_vcpu_add_default(struct kvm_vm *vm, uint32_t vcpuid, void *guest_code)
 
        /* Setup default exception vector of guest */
        set_reg(vm, vcpuid, RISCV_CSR_REG(stvec),
-               (unsigned long)guest_hang);
+               (unsigned long)guest_unexp_trap);
 }
 
 void vcpu_args_set(struct kvm_vm *vm, uint32_t vcpuid, unsigned int num, ...)
 
                uc.args[i] = va_arg(va, uint64_t);
        va_end(va);
 
-       sbi_ecall(KVM_RISCV_SELFTESTS_SBI_EXT, 0, (vm_vaddr_t)&uc,
-                 0, 0, 0, 0, 0);
+       sbi_ecall(KVM_RISCV_SELFTESTS_SBI_EXT,
+                 KVM_RISCV_SELFTESTS_SBI_UCALL,
+                 (vm_vaddr_t)&uc, 0, 0, 0, 0, 0);
 }
 
 uint64_t get_ucall(struct kvm_vm *vm, uint32_t vcpu_id, struct ucall *uc)
                memset(uc, 0, sizeof(*uc));
 
        if (run->exit_reason == KVM_EXIT_RISCV_SBI &&
-           run->riscv_sbi.extension_id == KVM_RISCV_SELFTESTS_SBI_EXT &&
-           run->riscv_sbi.function_id == 0) {
-               memcpy(&ucall, addr_gva2hva(vm, run->riscv_sbi.args[0]),
-                       sizeof(ucall));
-
-               vcpu_run_complete_io(vm, vcpu_id);
-               if (uc)
-                       memcpy(uc, &ucall, sizeof(ucall));
+           run->riscv_sbi.extension_id == KVM_RISCV_SELFTESTS_SBI_EXT) {
+               switch (run->riscv_sbi.function_id) {
+               case KVM_RISCV_SELFTESTS_SBI_UCALL:
+                       memcpy(&ucall, addr_gva2hva(vm,
+                              run->riscv_sbi.args[0]), sizeof(ucall));
+
+                       vcpu_run_complete_io(vm, vcpu_id);
+                       if (uc)
+                               memcpy(uc, &ucall, sizeof(ucall));
+
+                       break;
+               case KVM_RISCV_SELFTESTS_SBI_UNEXP:
+                       vcpu_dump(stderr, vm, vcpu_id, 2);
+                       TEST_ASSERT(0, "Unexpected trap taken by guest");
+                       break;
+               default:
+                       break;
+               }
        }
 
        return ucall.cmd;