Stack budget

CHANGELOG.md

@@ -6,6 +6,10 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/) a
 
 ## [Unreleased]
 
+### Added
+
+- `pm_options_stack_budget_set` was added to the C API. It sets the number of bytes of C stack the parser may use before reporting that the source is nested too deeply, instead of overflowing the stack. `SIZE_MAX` means no budget, which is the default. The nesting depth limit alone assumes a stack of roughly 8MB, so embedders that parse on a smaller stack, such as a thread's or a fiber's, should pass a budget.
+
 ### Changed
 
 - `Prism::Translation::Parser` now returns a `nil` AST whenever an error is reported and errors are not fatal, matching the parser gem.

include/prism/compiler/stack.h

@@ -0,0 +1,54 @@
+/**
+ * @file compiler/stack.h
+ */
+#ifndef PRISM_COMPILER_STACK_H
+#define PRISM_COMPILER_STACK_H
+
+#include "prism/compiler/inline.h"
+
+#include <stdint.h>
+
+#if defined(__has_builtin)
+    #if __has_builtin(__builtin_stack_address)
+        #define PRISM_HAS_BUILTIN_STACK_ADDRESS
+    #endif
+#endif
+
+#if defined(_MSC_VER)
+    #include <intrin.h>
+#endif
+
+/**
+ * Returns the current position of the C stack, for measuring how much of it has
+ * been consumed since an earlier call. The stack pointer is read directly where
+ * the compiler or architecture allows it. Unlike __builtin_frame_address, that
+ * does not force the calling function to set up a frame pointer, which costs a
+ * register on targets that omit frame pointers by default (x86-64 at -O2), and
+ * unlike the address of a local it does not add a slot to the frame.
+ */
+static PRISM_INLINE uintptr_t
+pm_stack_address(void) {
+#if defined(PRISM_HAS_BUILTIN_STACK_ADDRESS)
+    return (uintptr_t) __builtin_stack_address();
+#elif defined(_MSC_VER) && defined(_M_ARM64)
+    return (uintptr_t) __getReg(31);
+#elif defined(_MSC_VER)
+    return (uintptr_t) _AddressOfReturnAddress();
+#elif defined(__GNUC__) && defined(__x86_64__)
+    uintptr_t sp;
+    __asm__ ("{movq %%rsp, %0|mov %0, rsp}" : "=r" (sp));
+    return sp;
+#elif defined(__GNUC__) && defined(__aarch64__)
+    /* The register is 64 bits even where pointers are 32 (arm64_32). */
+    uint64_t sp;
+    __asm__ ("mov %0, sp" : "=r" (sp));
+    return (uintptr_t) sp;
+#elif defined(__GNUC__)
+    return (uintptr_t) __builtin_frame_address(0);
+#else
+    volatile char marker = 0;
+    return (uintptr_t) &marker;
+#endif
+}
+
+#endif

include/prism/internal/options.h

@@ -75,6 +75,15 @@ struct pm_options_t {
      */
     pm_string_t encoding;
 
+    /*
+     * The number of bytes of C stack that the parser may use, measured from
+     * where pm_parse is called. When the recursive descent goes past it, the
+     * parser reports that the input is nested too deeply instead of
+     * overflowing the stack. SIZE_MAX means no budget was given, in which case
+     * only PRISM_DEPTH_MAXIMUM bounds the recursion.
+     */
+    size_t stack_budget;
+
     /* The number of scopes surrounding the code that is being parsed. */
     size_t scopes_count;
 

include/prism/internal/parser.h

@@ -2,6 +2,8 @@
 #define PRISM_INTERNAL_PARSER_H
 
 #include "prism/compiler/accel.h"
+#include "prism/compiler/inline.h"
+#include "prism/compiler/stack.h"
 
 #include "prism/internal/arena.h"
 #include "prism/internal/comments.h"
@@ -879,6 +881,19 @@ struct pm_parser_t {
      */
     pm_node_list_t *current_block_exits;
 
+    /*
+     * The address of the C stack when pm_parse was entered. The stack budget
+     * is measured as the distance from this address.
+     */
+    uintptr_t stack_base;
+
+    /*
+     * The number of bytes of C stack below stack_base that the recursive
+     * descent may use before it reports that the input is nested too deeply.
+     * SIZE_MAX when no budget was given, so the check never fires.
+     */
+    size_t stack_budget;
+
     /* The version of prism that we should use to parse. */
     pm_options_version_t version;
 
@@ -1012,4 +1027,18 @@ void pm_parser_cleanup(pm_parser_t *parser);
  */
 bool pm_keyword_p(const uint8_t *start, const uint8_t *end);
 
+/*
+ * Returns true if parsing has used more C stack since pm_parse was entered than
+ * the stack budget allows. Every recursion that runs inside pm_parse checks
+ * this alongside its own depth limit, so that the budget bounds all of them.
+ * The distance is taken without assuming which direction the stack grows. With
+ * no budget the limit is SIZE_MAX, so this never fires.
+ */
+static PRISM_INLINE bool
+pm_parser_stack_exhausted(const pm_parser_t *parser) {
+    uintptr_t current = pm_stack_address();
+    uintptr_t used = current < parser->stack_base ? parser->stack_base - current : current - parser->stack_base;
+    return used > parser->stack_budget;
+}
+
 #endif

include/prism/options.h

@@ -288,6 +288,24 @@ PRISM_EXPORTED_FUNCTION uint8_t pm_options_raise_error(const pm_options_t *optio
  */
 PRISM_EXPORTED_FUNCTION void pm_options_raise_error_set(pm_options_t *options, uint8_t raise_error) PRISM_NONNULL(1);
 
+/**
+ * Set the stack_budget option on the given options struct. This is the number
+ * of bytes of C stack that the parser may use, measured from where pm_parse is
+ * called. Deeply nested input that would go past it is reported as nested too
+ * deeply instead of overflowing the stack. Callers should pass less than the
+ * stack they actually have left, since the parser checks the budget between
+ * frames rather than inside every one, and anything that walks the tree
+ * afterward recurses on the same stack. A budget of 0 leaves no room for any
+ * expression, so every one is reported as nested too deeply. A budget of
+ * SIZE_MAX (the default) disables the check, leaving PRISM_DEPTH_MAXIMUM as the
+ * only bound, which assumes a stack of roughly 8MB.
+ *
+ * @param options The options struct to set the stack_budget value on.
+ * @param stack_budget The stack_budget value to set, in bytes, or SIZE_MAX for
+ *     no budget.
+ */
+PRISM_EXPORTED_FUNCTION void pm_options_stack_budget_set(pm_options_t *options, size_t stack_budget) PRISM_NONNULL(1);
+
 /**
  * Allocate and zero out the scopes array on the given options struct.
  *

prism.gemspec

@@ -65,6 +65,7 @@ Gem::Specification.new do |spec|
     "include/prism/compiler/lifetimebound.hpp",
     "include/prism/compiler/nodiscard.h",
     "include/prism/compiler/nonnull.h",
+    "include/prism/compiler/stack.h",
     "include/prism/compiler/unused.h",
     "include/prism/internal/allocator.h",
     "include/prism/internal/allocator_debug.h",

src/options.c

@@ -20,6 +20,7 @@ pm_options_new(void) {
     if (options == NULL) PRISM_OOM();
 
     options->line = 1;
+    options->stack_budget = SIZE_MAX;
     return options;
 }
 
@@ -268,6 +269,14 @@ void pm_options_raise_error_set(pm_options_t *options, uint8_t raise_error) {
     options->raise_error = raise_error;
 }
 
+/**
+ * Set the stack_budget option on the given options struct.
+ */
+void
+pm_options_stack_budget_set(pm_options_t *options, size_t stack_budget) {
+    options->stack_budget = stack_budget;
+}
+
 // For some reason, GCC analyzer thinks we're leaking allocated scopes and
 // locals here, even though we definitely aren't. This is a false positive.
 // Ideally we wouldn't need to suppress this.
@@ -384,6 +393,7 @@ pm_options_read_s32(const char *data) {
 void
 pm_options_read(pm_options_t *options, const char *data) {
     options->line = 1; // default
+    options->stack_budget = SIZE_MAX; // default
     if (data == NULL) return;
 
     uint32_t filepath_length = pm_options_read_u32(data);

src/prism.c

@@ -17849,7 +17849,7 @@ static pm_node_t *
 parse_pattern(pm_parser_t *parser, pm_constant_id_set_t *captures, uint8_t flags, pm_diagnostic_id_t diag_id, uint16_t depth) {
     /* Nested patterns recurse back into parse_pattern without passing through
      * parse_expression, so this cycle needs its own depth check. */
-    if (PRISM_UNLIKELY(depth >= PRISM_DEPTH_MAXIMUM)) {
+    if (PRISM_UNLIKELY(depth >= PRISM_DEPTH_MAXIMUM || pm_parser_stack_exhausted(parser))) {
         pm_parser_err_current(parser, PM_ERR_NESTING_TOO_DEEP);
         return UP(pm_error_recovery_node_create(parser, PM_TOKEN_START(parser, &parser->current), PM_TOKEN_LENGTH(&parser->current)));
     }
@@ -22818,7 +22818,7 @@ parse_expression_terminator(pm_parser_t *parser, pm_node_t *node) {
  */
 static pm_node_t *
 parse_expression(pm_parser_t *parser, pm_binding_power_t binding_power, uint8_t flags, pm_diagnostic_id_t diag_id, uint16_t depth) {
-    if (PRISM_UNLIKELY(depth >= PRISM_DEPTH_MAXIMUM)) {
+    if (PRISM_UNLIKELY(depth >= PRISM_DEPTH_MAXIMUM || pm_parser_stack_exhausted(parser))) {
         pm_parser_err_current(parser, PM_ERR_NESTING_TOO_DEEP);
         return UP(pm_error_recovery_node_create(parser, PM_TOKEN_START(parser, &parser->current), PM_TOKEN_LENGTH(&parser->current)));
     }
@@ -23240,6 +23240,8 @@ pm_parser_init(pm_arena_t *arena, pm_parser_t *parser, const uint8_t *source, si
         .pattern_matching_newlines = false,
         .in_keyword_arg = false,
         .current_block_exits = NULL,
+        .stack_base = 0,
+        .stack_budget = SIZE_MAX,
         .semantic_token_seen = false,
         .frozen_string_literal = PM_OPTIONS_FROZEN_STRING_LITERAL_UNSET,
         .warn_mismatched_indentation = true
@@ -23303,6 +23305,9 @@ pm_parser_init(pm_arena_t *arena, pm_parser_t *parser, const uint8_t *source, si
         // attach_comments
         parser->attach_comments = options->attach_comments;
 
+        // stack_budget
+        parser->stack_budget = options->stack_budget;
+
         // scopes option
         parser->parsing_eval = options->scopes_count > 0;
         if (parser->parsing_eval) parser->warn_mismatched_indentation = false;
@@ -23649,6 +23654,7 @@ pm_parse_continuable(pm_parser_t *parser) {
  */
 pm_node_t *
 pm_parse(pm_parser_t *parser) {
+    parser->stack_base = pm_stack_address();
     pm_node_t *node = parse_program(parser);
     pm_parse_continuable(parser);
     if (parser->attach_comments) pm_comments_attach(parser, node);

src/regexp.c

@@ -710,11 +710,11 @@ pm_regexp_parse_hex_escape(pm_regexp_parser_t *parser) {
  * \p/\P (property), \x (hex), and other single-character escapes. Also
  * tracks encoding state for \M-x and \C-\M-x escapes. The meta and control
  * forms can wrap another escape, so this recurses and is bounded by the same
- * depth limit as groups and character sets.
+ * depth limit and stack budget as groups and character sets.
  */
 static bool
 pm_regexp_parse_backslash_escape(pm_regexp_parser_t *parser, uint16_t depth) {
-    if (depth >= PM_REGEXP_PARSE_DEPTH_MAX) {
+    if (depth >= PM_REGEXP_PARSE_DEPTH_MAX || pm_parser_stack_exhausted(parser->parser)) {
         pm_regexp_parse_error(parser, parser->start, parser->end, "parse depth limit over");
         return false;
     }
@@ -846,7 +846,7 @@ pm_regexp_parse_character_set(pm_regexp_parser_t *parser, uint16_t depth) {
  */
 static bool
 pm_regexp_parse_lbracket(pm_regexp_parser_t *parser, uint16_t depth) {
-    if (depth >= PM_REGEXP_PARSE_DEPTH_MAX) {
+    if (depth >= PM_REGEXP_PARSE_DEPTH_MAX || pm_parser_stack_exhausted(parser->parser)) {
         pm_regexp_parse_error(parser, parser->start, parser->end, "parse depth limit over");
         return false;
     }
@@ -1005,6 +1005,16 @@ pm_regexp_options_state(pm_regexp_options_t *options, uint8_t key) {
  */
 static bool
 pm_regexp_parse_group(pm_regexp_parser_t *parser, uint16_t depth) {
+    /* Every level of nested groups passes through here, so this bounds the
+     * group recursion by the stack budget. The check lives here rather than
+     * beside the depth check in pm_regexp_parse_expression because there it
+     * stops that function from being inlined, which adds a frame to every
+     * level of nesting. */
+    if (pm_parser_stack_exhausted(parser->parser)) {
+        pm_regexp_parse_error(parser, parser->start, parser->end, "parse depth limit over");
+        return false;
+    }
+
     const uint8_t *group_start = parser->cursor;
 
     pm_regexp_options_t options;