From 7b6f88ea5ab60a70b729624051b05e30931dd587 Mon Sep 17 00:00:00 2001 From: lisyarus Date: Tue, 31 Mar 2026 19:27:57 +0300 Subject: [PATCH] Add bounds checks for aarch64 instruction builder immediate arguments --- libs/jit/source/arch/aarch64/compiler.cpp | 1 + .../arch/aarch64/instruction_builder.cpp | 69 +++++++++++++++++++ 2 files changed, 70 insertions(+) diff --git a/libs/jit/source/arch/aarch64/compiler.cpp b/libs/jit/source/arch/aarch64/compiler.cpp index 320fe32..3b88356 100644 --- a/libs/jit/source/arch/aarch64/compiler.cpp +++ b/libs/jit/source/arch/aarch64/compiler.cpp @@ -985,6 +985,7 @@ namespace pslang::jit::aarch64 // TODO: inefficient for small types, maybe only round up to type alignment? // Need to make sure all read/write arm64 instructions used can handle offsets that // are not a multiple of 8 + // UPD: Simply rounding to alignment doesn't work here, need to investigate stack_size += ((size + 7) / 8) * 8; } stack_position[it] = stack_size; diff --git a/libs/jit/source/arch/aarch64/instruction_builder.cpp b/libs/jit/source/arch/aarch64/instruction_builder.cpp index af33297..3bde25f 100644 --- a/libs/jit/source/arch/aarch64/instruction_builder.cpp +++ b/libs/jit/source/arch/aarch64/instruction_builder.cpp @@ -1,8 +1,40 @@ #include +#include + namespace pslang::jit::aarch64 { + template + static void check_bits(T value, std::uint32_t bits, char const * message) + { + auto mask = (T(1) << bits) - T(1); + bool good = true; + if (value >= T(0)) + { + good = value == (value & mask); + } + else + { + good = (value & (~mask)) == (T(-1) & (~mask)); + } + if (!good) + { + T min, max; + if (std::is_unsigned_v) + { + min = 0; + max = mask; + } + else + { + max = (T(1) << (bits - 1)) - 1; + min = - (max + 1); + } + throw std::runtime_error(std::format("{}: {} not in range [{}..{})", message, value, min, max)); + } + } + static constexpr std::uint32_t REG_MASK = 0x1fu; void instruction_builder::nop() @@ -12,121 +44,145 @@ namespace pslang::jit::aarch64 void instruction_builder::movz(std::uint8_t reg, std::uint16_t val, std::uint8_t shift) { + check_bits(shift, 2, "Bad movz shift value"); do_push(0xd2800000u | (reg & REG_MASK) | (val << 5) | ((shift & 0x3u) << 21)); } void instruction_builder::movk(std::uint8_t reg, std::uint16_t val, std::uint8_t shift) { + check_bits(shift, 2, "Bad movk shift value"); do_push(0xf2800000u | (reg & REG_MASK) | (val << 5) | ((shift & 0x3u) << 21)); } void instruction_builder::str(std::uint8_t reg_src, std::uint8_t reg_addr, std::uint16_t offset) { + check_bits(offset, 12, "Bad str offset value"); do_push(0xf9000000u | (reg_src & REG_MASK) | ((reg_addr & REG_MASK) << 5) | ((std::uint16_t(offset) & 0xfffu) << 10)); } void instruction_builder::strw(std::uint8_t reg_src, std::uint8_t reg_addr, std::uint16_t offset) { + check_bits(offset, 12, "Bad strw offset value"); do_push(0xb9000000u | (reg_src & REG_MASK) | ((reg_addr & REG_MASK) << 5) | ((std::uint16_t(offset) & 0xfffu) << 10)); } void instruction_builder::strh(std::uint8_t reg_src, std::uint8_t reg_addr, std::uint16_t offset) { + check_bits(offset, 12, "Bad strh offset value"); do_push(0x79000000u | (reg_src & REG_MASK) | ((reg_addr & REG_MASK) << 5) | ((std::uint16_t(offset) & 0xfffu) << 10)); } void instruction_builder::strb(std::uint8_t reg_src, std::uint8_t reg_addr, std::uint16_t offset) { + check_bits(offset, 12, "Bad strb offset value"); do_push(0x39000000u | (reg_src & REG_MASK) | ((reg_addr & REG_MASK) << 5) | ((std::uint16_t(offset) & 0xfffu) << 10)); } void instruction_builder::str_pre(std::uint8_t reg_src, std::uint8_t reg_addr, std::int16_t offset) { + check_bits(offset, 9, "Bad str_pre offset value"); do_push(0xf8000c00u | (reg_src & REG_MASK) | ((reg_addr & REG_MASK) << 5) | ((std::uint16_t(offset) & 0x1ffu) << 12)); } void instruction_builder::stur(std::uint8_t reg_src, std::uint8_t reg_addr, std::int16_t offset) { + check_bits(offset, 9, "Bad stur offset value"); do_push(0xf8000000u | (reg_src & REG_MASK) | ((reg_addr & REG_MASK) << 5) | ((std::uint16_t(offset) & 0x1ffu) << 12)); } void instruction_builder::sturw(std::uint8_t reg_src, std::uint8_t reg_addr, std::int16_t offset) { + check_bits(offset, 9, "Bad sturw offset value"); do_push(0xb8000000u | (reg_src & REG_MASK) | ((reg_addr & REG_MASK) << 5) | ((std::uint16_t(offset) & 0x1ffu) << 12)); } void instruction_builder::sturh(std::uint8_t reg_src, std::uint8_t reg_addr, std::int16_t offset) { + check_bits(offset, 9, "Bad sturh offset value"); do_push(0x78000000u | (reg_src & REG_MASK) | ((reg_addr & REG_MASK) << 5) | ((std::uint16_t(offset) & 0x1ffu) << 12)); } void instruction_builder::sturb(std::uint8_t reg_src, std::uint8_t reg_addr, std::int16_t offset) { + check_bits(offset, 9, "Bad sturb offset value"); do_push(0x38000000u | (reg_src & REG_MASK) | ((reg_addr & REG_MASK) << 5) | ((std::uint16_t(offset) & 0x1ffu) << 12)); } void instruction_builder::ldr(std::uint8_t reg_dst, std::uint8_t reg_addr, std::uint16_t offset) { + check_bits(offset, 12, "Bad ldr offset value"); do_push(0xf9400000u | (reg_dst & REG_MASK) | ((reg_addr & REG_MASK) << 5) | ((std::uint16_t(offset) & 0xfffu) << 10)); } void instruction_builder::ldrw(std::uint8_t reg_dst, std::uint8_t reg_addr, std::uint16_t offset) { + check_bits(offset, 12, "Bad ldrw offset value"); do_push(0xb9400000u | (reg_dst & REG_MASK) | ((reg_addr & REG_MASK) << 5) | ((std::uint16_t(offset) & 0xfffu) << 10)); } void instruction_builder::ldrh(std::uint8_t reg_dst, std::uint8_t reg_addr, std::uint16_t offset) { + check_bits(offset, 12, "Bad ldrh offset value"); do_push(0x79400000u | (reg_dst & REG_MASK) | ((reg_addr & REG_MASK) << 5) | ((std::uint16_t(offset) & 0xfffu) << 10)); } void instruction_builder::ldrb(std::uint8_t reg_dst, std::uint8_t reg_addr, std::uint16_t offset) { + check_bits(offset, 12, "Bad ldrb offset value"); do_push(0x39400000u | (reg_dst & REG_MASK) | ((reg_addr & REG_MASK) << 5) | ((std::uint16_t(offset) & 0xfffu) << 10)); } void instruction_builder::ldr_pre(std::uint8_t reg_dst, std::uint8_t reg_addr, std::int16_t offset) { + check_bits(offset, 9, "Bad ldr_ptr offset value"); do_push(0xf8400c00u | (reg_dst & REG_MASK) | ((reg_addr & REG_MASK) << 5) | ((std::uint16_t(offset) & 0x1ffu) << 12)); } void instruction_builder::ldr_post(std::uint8_t reg_dst, std::uint8_t reg_addr, std::int16_t offset) { + check_bits(offset, 9, "Bad ldr_post offset value"); do_push(0xf8400400u | (reg_dst & REG_MASK) | ((reg_addr & REG_MASK) << 5) | ((std::uint16_t(offset) & 0x1ffu) << 12)); } void instruction_builder::ldur(std::uint8_t reg_dst, std::uint8_t reg_addr, std::int16_t offset) { + check_bits(offset, 9, "Bad ldur offset value"); do_push(0xf8400000u | (reg_dst & REG_MASK) | ((reg_addr & REG_MASK) << 5) | ((std::uint16_t(offset) & 0x1ffu) << 12)); } void instruction_builder::ldurw(std::uint8_t reg_dst, std::uint8_t reg_addr, std::int16_t offset) { + check_bits(offset, 9, "Bad ldurw offset value"); do_push(0xb8400000u | (reg_dst & REG_MASK) | ((reg_addr & REG_MASK) << 5) | ((std::uint16_t(offset) & 0x1ffu) << 12)); } void instruction_builder::ldurh(std::uint8_t reg_dst, std::uint8_t reg_addr, std::int16_t offset) { + check_bits(offset, 9, "Bad ldurh offset value"); do_push(0x78400000u | (reg_dst & REG_MASK) | ((reg_addr & REG_MASK) << 5) | ((std::uint16_t(offset) & 0x1ffu) << 12)); } void instruction_builder::ldurb(std::uint8_t reg_dst, std::uint8_t reg_addr, std::int16_t offset) { + check_bits(offset, 9, "Bad ldurb offset value"); do_push(0x38400000u | (reg_dst & REG_MASK) | ((reg_addr & REG_MASK) << 5) | ((std::uint16_t(offset) & 0x1ffu) << 12)); } void instruction_builder::ldr_pc(std::uint8_t reg_dst, std::int32_t offset) { + check_bits(offset, 19, "Bad ldr_pc offset value"); do_push(0x58000000u | (reg_dst & REG_MASK) | ((((std::uint32_t)offset) & 0x7ffffu) << 5)); } void instruction_builder::add_imm(std::uint8_t reg_src, std::uint8_t reg_dst, std::uint16_t value) { + check_bits(value, 12, "Bad add_imm value"); do_push(0x91000000u | (reg_dst & REG_MASK) | ((reg_src & REG_MASK) << 5) | ((value & 0xfffu) << 10)); } void instruction_builder::sub_imm(std::uint8_t reg_src, std::uint8_t reg_dst, std::uint16_t value) { + check_bits(value, 12, "Bad sub_imm value"); do_push(0xd1000000u | (reg_dst & REG_MASK) | ((reg_src & REG_MASK) << 5) | ((value & 0xfffu) << 10)); } @@ -227,16 +283,19 @@ namespace pslang::jit::aarch64 void instruction_builder::cbnz(std::uint8_t reg_src, std::int32_t offset) { + check_bits(offset, 19, "Bad cbnz offset value"); do_push(0xb5000000u | (reg_src & REG_MASK) | ((std::uint32_t(offset) & 0x7ffffu) << 5)); } void instruction_builder::cbz(std::uint8_t reg_src, std::int32_t offset) { + check_bits(offset, 19, "Bad cbz offset value"); do_push(0xb4000000u | (reg_src & REG_MASK) | ((std::uint32_t(offset) & 0x7ffffu) << 5)); } void instruction_builder::cb_inject(std::uint8_t * opcode, std::int32_t offset) { + check_bits(offset, 19, "Bad cb_inject offset value"); auto dst = (std::uint32_t *)opcode; *dst &= ~(0x7ffffu << 5); *dst |= (std::uint32_t(offset) & 0x7ffffu) << 5; @@ -244,11 +303,13 @@ namespace pslang::jit::aarch64 void instruction_builder::b(std::int32_t offset) { + check_bits(offset, 26, "Bad b offset value"); do_push(0x14000000u | (std::uint32_t(offset) & 0x3ffffffu)); } void instruction_builder::bl(std::int32_t offset) { + check_bits(offset, 26, "Bad bl offset value"); do_push(0x94000000u | (std::uint32_t(offset) & 0x3ffffffu)); } @@ -264,6 +325,7 @@ namespace pslang::jit::aarch64 void instruction_builder::b_inject(std::uint8_t * opcode, std::int32_t offset) { + check_bits(offset, 26, "Bad b_inject offset value"); auto dst = (std::uint32_t *)opcode; *dst &= ~0x3ffffffu; *dst |= (std::uint32_t(offset) & 0x3ffffffu); @@ -271,12 +333,14 @@ namespace pslang::jit::aarch64 void instruction_builder::adr(std::uint8_t reg_dst, std::int32_t offset) { + check_bits(offset, 21, "Bad adr offset value"); auto offset_val = std::uint32_t(offset) & 0x1fffffu; do_push(0x10000000u | (reg_dst & REG_MASK) | ((offset_val >> 2) << 5) | ((offset_val & 0x3u) << 29)); } void instruction_builder::adr_inject(std::uint8_t * opcode, std::int32_t offset) { + check_bits(offset, 21, "Bad adr_inject offset value"); auto offset_val = std::uint32_t(offset) & 0x1fffffu; auto dst = (std::uint32_t *)opcode; *dst &= 0x9f00001fu; @@ -286,26 +350,31 @@ namespace pslang::jit::aarch64 void instruction_builder::ldr_fp_pc(std::uint8_t reg_dst, std::uint8_t mode, std::int32_t offset) { + check_bits(offset, 19, "Bad ldr_fp_pc offset value"); do_push(0x1c000000u | ((mode & 0x1u) << 30) | (reg_dst & REG_MASK) | ((std::uint32_t(offset) & 0x7ffffu) << 5)); } void instruction_builder::ldr_fp(std::uint8_t reg_dst, std::uint8_t mode, std::uint8_t reg_addr, std::uint16_t offset) { + check_bits(offset, 12, "Bad ldr_fp offset value"); do_push(0x3d400000u | (reg_dst & REG_MASK) | ((reg_addr & REG_MASK) << 5) | ((offset & 0xfffu) << 10) | ((mode & 0x3u) << 30)); } void instruction_builder::ldur_fp(std::uint8_t reg_dst, std::uint8_t mode, std::uint8_t reg_addr, std::uint16_t offset) { + check_bits(offset, 9, "Bad ldur_fp offset value"); do_push(0x3c400000u | (reg_dst & REG_MASK) | ((reg_addr & REG_MASK) << 5) | ((std::uint16_t(offset) & 0x1ffu) << 12) | ((mode & 0x3u) << 30)); } void instruction_builder::str_fp(std::uint8_t reg_src, std::uint8_t mode, std::uint8_t reg_addr, std::uint16_t offset) { + check_bits(offset, 12, "Bad str_fp offset value"); do_push(0x3d000000u | (reg_src & REG_MASK) | ((reg_addr & REG_MASK) << 5) | ((offset & 0xfffu) << 10) | ((mode & 0x3u) << 30)); } void instruction_builder::stur_fp(std::uint8_t reg_src, std::uint8_t mode, std::uint8_t reg_addr, std::int16_t offset) { + check_bits(offset, 9, "Bad stur_fp offset value"); do_push(0x3c000000u | (reg_src & REG_MASK) | ((reg_addr & REG_MASK) << 5) | ((std::uint16_t(offset) & 0x1ffu) << 12) | ((mode & 0x3u) << 30)); }