Add bounds checks for aarch64 instruction builder immediate arguments

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;

libs/jit/source/arch/aarch64/instruction_builder.cpp

@@ -1,8 +1,40 @@
 #include <pslang/jit/arch/aarch64/instruction_builder.hpp>
 
+#include <format>
+
 namespace pslang::jit::aarch64
 {
 
+	template <typename T>
+	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<T>)
+			{
+				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));
 	}