From f41b4d44f95978dfa97af04af00055dc3fbf7978 Mon Sep 17 00:00:00 2001
From: Alan Wu <XrXr@users.noreply.github.com>
Date: Wed, 1 Dec 2021 14:15:23 -0500
Subject: YJIT: Bounds check every byte in the assembler

Previously, YJIT assumed that basic blocks never consume more than
1 KiB of memory. This assumption does not hold for long Ruby methods
such as the one in the following:

```ruby
eval(<<RUBY)
def set_local_a_lot
  #{'_=0;'*0x40000}
end
RUBY

set_local_a_lot
```

For low `--yjit-exec-mem-size` values, one basic block could exhaust the
entire buffer.

Introduce a new field `codeblock_t::dropped_bytes` that the assembler
sets whenever it runs out of space. Check this field in
gen_single_block() to respond to out of memory situations and other
error conditions. This design avoids making the control flow graph of
existing code generation functions more complex.

Use POSIX shell in misc/test_yjit_asm.sh since bash is expanding
`0%/*/*` differently.

Co-authored-by: Aaron Patterson <tenderlove@ruby-lang.org>
---
 yjit_codegen.c | 18 +++++++-----------
 1 file changed, 7 insertions(+), 11 deletions(-)

(limited to 'yjit_codegen.c')

diff --git a/yjit_codegen.c b/yjit_codegen.c
index 061cd1ead7..8c888fd53a 100644
--- a/yjit_codegen.c
+++ b/yjit_codegen.c
@@ -554,7 +554,7 @@ yjit_entry_prologue(codeblock_t *cb, const rb_iseq_t *iseq)
 
     const uint32_t old_write_pos = cb->write_pos;
 
-    // Align the current write positon to cache line boundaries
+    // Align the current write position to cache line boundaries
     cb_align_pos(cb, 64);
 
     uint8_t *code_ptr = cb_get_ptr(cb, cb->write_pos);
@@ -640,16 +640,6 @@ gen_single_block(blockid_t blockid, const ctx_t *start_ctx, rb_execution_context
 {
     RUBY_ASSERT(cb != NULL);
 
-    // Check if there is enough executable memory.
-    // FIXME: This bound isn't enforced and long blocks can potentially use more.
-    enum { MAX_CODE_PER_BLOCK = 1024 };
-    if (cb->write_pos + MAX_CODE_PER_BLOCK >= cb->mem_size) {
-        return NULL;
-    }
-    if (ocb->write_pos + MAX_CODE_PER_BLOCK >= ocb->mem_size) {
-        return NULL;
-    }
-
     // Allocate the new block
     block_t *block = calloc(1, sizeof(block_t));
     if (!block) {
@@ -778,6 +768,12 @@ gen_single_block(blockid_t blockid, const ctx_t *start_ctx, rb_execution_context
     // doesn't go to the next instruction.
     RUBY_ASSERT(!jit.record_boundary_patch_point);
 
+    // If code for the block doesn't fit, free the block and fail.
+    if (cb->dropped_bytes || ocb->dropped_bytes) {
+        yjit_free_block(block);
+        return NULL;
+    }
+
     if (YJIT_DUMP_MODE >= 2) {
         // Dump list of compiled instrutions
         fprintf(stderr, "Compiled the following for iseq=%p:\n", (void *)iseq);
-- 
cgit v1.2.3