Arm64Emitter: Check end of allocated space when emitting code

JitArm64 port of 5b52b3e.

Source/Core/Common/Arm64Emitter.cpp

@@ -71,14 +71,16 @@ std::optional<u8> FPImm8FromFloat(float value)
 }
 }  // Anonymous namespace
 
-void ARM64XEmitter::SetCodePtrUnsafe(u8* ptr)
+void ARM64XEmitter::SetCodePtrUnsafe(u8* ptr, u8* end, bool write_failed)
 {
   m_code = ptr;
+  m_code_end = end;
+  m_write_failed = write_failed;
 }
 
 void ARM64XEmitter::SetCodePtr(u8* ptr, u8* end, bool write_failed)
 {
-  SetCodePtrUnsafe(ptr);
+  SetCodePtrUnsafe(ptr, end, write_failed);
   m_lastCacheFlushEnd = ptr;
 }
 
@@ -92,6 +94,16 @@ u8* ARM64XEmitter::GetWritableCodePtr()
   return m_code;
 }
 
+const u8* ARM64XEmitter::GetCodeEnd() const
+{
+  return m_code_end;
+}
+
+u8* ARM64XEmitter::GetWritableCodeEnd()
+{
+  return m_code_end;
+}
+
 void ARM64XEmitter::ReserveCodeSpace(u32 bytes)
 {
   for (u32 i = 0; i < bytes / 4; i++)
@@ -116,6 +128,13 @@ u8* ARM64XEmitter::AlignCodePage()
 
 void ARM64XEmitter::Write32(u32 value)
 {
+  if (m_code + sizeof(u32) > m_code_end)
+  {
+    m_code = m_code_end;
+    m_write_failed = true;
+    return;
+  }
+
   std::memcpy(m_code, &value, sizeof(u32));
   m_code += sizeof(u32);
 }
@@ -659,6 +678,9 @@ static constexpr u32 MaskImm26(s64 distance)
 // FixupBranch branching
 void ARM64XEmitter::SetJumpTarget(FixupBranch const& branch)
 {
+  if (!branch.ptr)
+    return;
+
   bool Not = false;
   u32 inst = 0;
   s64 distance = (s64)(m_code - branch.ptr);
@@ -709,67 +731,68 @@ void ARM64XEmitter::SetJumpTarget(FixupBranch const& branch)
   std::memcpy(branch.ptr, &inst, sizeof(inst));
 }
 
-FixupBranch ARM64XEmitter::CBZ(ARM64Reg Rt)
+FixupBranch ARM64XEmitter::WriteFixupBranch()
 {
   FixupBranch branch{};
   branch.ptr = m_code;
+  BRK(0);
+
+  // If we couldn't write the full jump instruction, indicate that in the returned FixupBranch by
+  // setting the branch's address to null. This will prevent a later SetJumpTarget() from writing to
+  // invalid memory.
+  if (HasWriteFailed())
+    branch.ptr = nullptr;
+
+  return branch;
+}
+
+FixupBranch ARM64XEmitter::CBZ(ARM64Reg Rt)
+{
+  FixupBranch branch = WriteFixupBranch();
   branch.type = FixupBranch::Type::CBZ;
   branch.reg = Rt;
-  NOP();
   return branch;
 }
 FixupBranch ARM64XEmitter::CBNZ(ARM64Reg Rt)
 {
-  FixupBranch branch{};
-  branch.ptr = m_code;
+  FixupBranch branch = WriteFixupBranch();
   branch.type = FixupBranch::Type::CBNZ;
   branch.reg = Rt;
-  NOP();
   return branch;
 }
 FixupBranch ARM64XEmitter::B(CCFlags cond)
 {
-  FixupBranch branch{};
-  branch.ptr = m_code;
+  FixupBranch branch = WriteFixupBranch();
   branch.type = FixupBranch::Type::BConditional;
   branch.cond = cond;
-  NOP();
   return branch;
 }
 FixupBranch ARM64XEmitter::TBZ(ARM64Reg Rt, u8 bit)
 {
-  FixupBranch branch{};
-  branch.ptr = m_code;
+  FixupBranch branch = WriteFixupBranch();
   branch.type = FixupBranch::Type::TBZ;
   branch.reg = Rt;
   branch.bit = bit;
-  NOP();
   return branch;
 }
 FixupBranch ARM64XEmitter::TBNZ(ARM64Reg Rt, u8 bit)
 {
-  FixupBranch branch{};
-  branch.ptr = m_code;
+  FixupBranch branch = WriteFixupBranch();
   branch.type = FixupBranch::Type::TBNZ;
   branch.reg = Rt;
   branch.bit = bit;
-  NOP();
   return branch;
 }
 FixupBranch ARM64XEmitter::B()
 {
-  FixupBranch branch{};
-  branch.ptr = m_code;
+  FixupBranch branch = WriteFixupBranch();
   branch.type = FixupBranch::Type::B;
-  NOP();
   return branch;
 }
 FixupBranch ARM64XEmitter::BL()
 {
-  FixupBranch branch{};
-  branch.ptr = m_code;
+  FixupBranch branch = WriteFixupBranch();
   branch.type = FixupBranch::Type::BL;
-  NOP();
   return branch;
 }
 
@@ -1945,12 +1968,12 @@ bool ARM64XEmitter::MOVI2R2(ARM64Reg Rd, u64 imm1, u64 imm2)
   MOVI2R(Rd, imm1);
   int size1 = GetCodePtr() - start_pointer;
 
-  SetCodePtrUnsafe(start_pointer);
+  m_code = start_pointer;
 
   MOVI2R(Rd, imm2);
   int size2 = GetCodePtr() - start_pointer;
 
-  SetCodePtrUnsafe(start_pointer);
+  m_code = start_pointer;
 
   bool element = size1 > size2;
 

Source/Core/Common/Arm64Emitter.h

@@ -725,8 +725,18 @@ class ARM64XEmitter
   friend class ARM64FloatEmitter;
 
 private:
-  u8* m_code;
-  u8* m_lastCacheFlushEnd;
+  // Pointer to memory where code will be emitted to.
+  u8* m_code = nullptr;
+
+  // Pointer past the end of the memory region we're allowed to emit to.
+  // Writes that would reach this memory are refused and will set the m_write_failed flag instead.
+  u8* m_code_end = nullptr;
+
+  u8* m_lastCacheFlushEnd = nullptr;
+
+  // Set to true when a write request happens that would write past m_code_end.
+  // Must be cleared with SetCodePtr() afterwards.
+  bool m_write_failed = false;
 
   void AddImmediate(ARM64Reg Rd, ARM64Reg Rn, u64 imm, bool shift, bool negative, bool flags);
   void EncodeCompareBranchInst(u32 op, ARM64Reg Rt, const void* ptr);
@@ -760,6 +770,8 @@ private:
   void EncodeAddressInst(u32 op, ARM64Reg Rd, s32 imm);
   void EncodeLoadStoreUnscaled(u32 size, u32 op, ARM64Reg Rt, ARM64Reg Rn, s32 imm);
 
+  FixupBranch WriteFixupBranch();
+
   template <typename T>
   void MOVI2RImpl(ARM64Reg Rd, T imm);
 
@@ -767,27 +779,30 @@ protected:
   void Write32(u32 value);
 
 public:
-  ARM64XEmitter() : m_code(nullptr), m_lastCacheFlushEnd(nullptr) {}
-  ARM64XEmitter(u8* code_ptr)
+  ARM64XEmitter() = default;
+  ARM64XEmitter(u8* code, u8* code_end)
+      : m_code(code), m_code_end(code_end), m_lastCacheFlushEnd(code)
   {
-    m_code = code_ptr;
-    m_lastCacheFlushEnd = code_ptr;
   }
 
   virtual ~ARM64XEmitter() {}
 
-  // 'end' and 'write_failed' are unused in the ARM code emitter at the moment.
-  // They're just here for interface compatibility with the x64 code emitter.
   void SetCodePtr(u8* ptr, u8* end, bool write_failed = false);
 
-  void SetCodePtrUnsafe(u8* ptr);
+  void SetCodePtrUnsafe(u8* ptr, u8* end, bool write_failed = false);
+  const u8* GetCodePtr() const;
+  u8* GetWritableCodePtr();
+  const u8* GetCodeEnd() const;
+  u8* GetWritableCodeEnd();
   void ReserveCodeSpace(u32 bytes);
   u8* AlignCode16();
   u8* AlignCodePage();
-  const u8* GetCodePtr() const;
   void FlushIcache();
   void FlushIcacheSection(u8* start, u8* end);
-  u8* GetWritableCodePtr();
+
+  // Should be checked after a block of code has been generated to see if the code has been
+  // successfully written to memory. Do not call the generated code when this returns true!
+  bool HasWriteFailed() const { return m_write_failed; }
 
   // FixupBranch branching
   void SetJumpTarget(FixupBranch const& branch);