From 48c9f08b7b25766324282b7605177e2f2be755bb Mon Sep 17 00:00:00 2001 From: Xuan-Son Nguyen Date: Sun, 25 Oct 2026 23:53:08 +0100 Subject: [PATCH] common, grammar: add recursive depth limit, make it a requirement (#31153) * common, grammar: add recursive depth limit, make it a requirement * catch more recursive functions * common : tidy up peg parser changes * grammar : reject candidates in a loop, without a depth limit The depth limit in llama_grammar_reject_candidates_for_stack counted the code points of a vocab token, so any token of 256 chars or more was always rejected under a grammar, such as the 512-space token of GLM-4.7-Flash or the long space tokens of Laguna-S-2.1. The rejection now walks the code points with its own stack of levels instead of recursing: there is no limit and no stack risk, whatever the length of a token from the GGUF, and the masks match master. * grammar : track rejected candidates as indices into the input At a given code point a candidate only differs from its input by the position of its code point pointer, so the levels now carry indices into the input candidates instead of copies. The first level needs no special case and the rejects map back to the input directly. Grammar sampling gets faster than master on GLM-4.7, Laguna, Qwen3.6 and gemma-4 with the json, list and spaces grammars, same masks. * nit : shorten the comment of llama_grammar_reject_step * common : add depth_guard for schema walks * grammar : move GRAMMAR_RECURSIVE_LIMIT into the .cpp * common : reduce recursion limit for peg parser * improve --------- Co-authored-by: Alde Rojas Co-authored-by: Pascal --- common/common.cpp | 84 +++++------ common/common.h | 2 + common/jinja/parser.cpp | 24 +++- common/jinja/runtime.cpp | 30 +++- common/jinja/runtime.h | 9 +- common/jinja/utils.h | 2 + common/jinja/value.cpp | 19 +-- common/jinja/value.h | 2 +- common/json-schema-to-grammar.cpp | 60 ++++++-- common/json-schema-to-grammar.h | 1 + common/json-schema.cpp | 50 ++++--- common/peg-parser.cpp | 42 ++++-- common/peg-parser.h | 6 + skills/code-review/SKILL.md | 1 + src/llama-grammar.cpp | 194 +++++++++++++++++--------- src/llama-vocab.cpp | 79 +++++------ tests/peg-parser/test-json-parser.cpp | 19 +++ tests/test-grammar-integration.cpp | 31 ++++ tests/test-jinja.cpp | 44 ++++++ tests/test-json-schema-to-grammar.cpp | 17 +++ tests/test-server-component.cpp | 27 ++++ 21 files changed, 531 insertions(+), 212 deletions(-) diff --git a/common/common.cpp b/common/common.cpp index 3237c0fa7..a4b4feaee 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -742,50 +742,56 @@ static inline bool glob_class_match(const char c, const char * pattern, const ch } // simple glob: * matches non-/ chars, ** matches anything including /, [] matches character class -static inline bool glob_match(const char * pattern, const char * str) { - if (*pattern == '\0') { - return *str == '\0'; - } - if (pattern[0] == '*' && pattern[1] == '*') { - const char * p = pattern + 2; - if (glob_match(p, str)) return true; - if (*str != '\0') return glob_match(pattern, str + 1); +// only wildcards recurse, too many wildcards: no match +static bool glob_match(const char * pattern, const char * str, int depth = 0) { + if (depth >= COMMON_RECURSIVE_LIMIT) { return false; } - if (*pattern == '*') { - const char * p = pattern + 1; - for (; *str != '\0' && *str != '/'; str++) { - if (glob_match(p, str)) return true; - } - return glob_match(p, str); - } - if (*pattern == '?' && *str != '\0' && *str != '/') { - return glob_match(pattern + 1, str + 1); - } - if (*pattern == '[') { - const char * class_end = pattern + 1; - // If first character after '[' is ']' or '-', treat it as literal - if (*class_end == ']' || *class_end == '-') { - class_end++; - } - while (*class_end != '\0' && *class_end != ']') { - class_end++; - } - if (*class_end == ']') { - if (*str == '\0') return false; - bool matched = glob_class_match(*str, pattern + 1, class_end); - return matched && glob_match(class_end + 1, str + 1); - } else { - if (*str == '[') { - return glob_match(pattern + 1, str + 1); + while (*pattern != '\0') { + if (pattern[0] == '*' && pattern[1] == '*') { + const char * p = pattern + 2; + for (;; str++) { + if (glob_match(p, str, depth + 1)) return true; + if (*str == '\0') return false; } - return false; } + if (*pattern == '*') { + const char * p = pattern + 1; + for (; *str != '\0' && *str != '/'; str++) { + if (glob_match(p, str, depth + 1)) return true; + } + return glob_match(p, str, depth + 1); + } + if (*pattern == '?' && *str != '\0' && *str != '/') { + pattern++; + str++; + continue; + } + if (*pattern == '[') { + const char * class_end = pattern + 1; + // If first character after '[' is ']' or '-', treat it as literal + if (*class_end == ']' || *class_end == '-') { + class_end++; + } + while (*class_end != '\0' && *class_end != ']') { + class_end++; + } + if (*class_end == ']') { + if (*str == '\0' || !glob_class_match(*str, pattern + 1, class_end)) return false; + pattern = class_end + 1; + str++; + continue; + } + if (*str != '[') return false; + pattern++; + str++; + continue; + } + if (*pattern != *str) return false; + pattern++; + str++; } - if (*pattern == *str) { - return glob_match(pattern + 1, str + 1); - } - return false; + return *str == '\0'; } bool glob_match(const std::string & pattern, const std::string & str) { diff --git a/common/common.h b/common/common.h index 70d6a671f..8e3fe2c89 100644 --- a/common/common.h +++ b/common/common.h @@ -31,6 +31,8 @@ #define DIRECTORY_SEPARATOR '/' #endif // _WIN32 +#define COMMON_RECURSIVE_LIMIT 128 + #define COM_DBG(fmt, ...) LOG_DBG("cmn %12.*s: " fmt, 12, __func__, __VA_ARGS__) #define COM_TRC(fmt, ...) LOG_TRC("cmn %12.*s: " fmt, 12, __func__, __VA_ARGS__) #define COM_INF(fmt, ...) LOG_INF("cmn %12.*s: " fmt, 12, __func__, __VA_ARGS__) diff --git a/common/jinja/parser.cpp b/common/jinja/parser.cpp index 67dcafe17..7e99a1ac8 100644 --- a/common/jinja/parser.cpp +++ b/common/jinja/parser.cpp @@ -21,6 +21,17 @@ static bool is_type(const statement_ptr & ptr) { class parser { const std::vector & tokens; size_t current = 0; + int depth = 0; + + struct depth_guard { + parser & p; + depth_guard(parser & p) : p(p) { + if (++p.depth > JINJA_RECURSIVE_LIMIT) { + throw parser_exception("Parser Error: template is nested more than " + std::to_string(JINJA_RECURSIVE_LIMIT) + " levels deep", p.source, p.peek().pos); + } + } + ~depth_guard() { p.depth--; } + }; std::string source; // for error reporting @@ -94,6 +105,7 @@ private: } statement_ptr parse_any() { + depth_guard guard(*this); size_t start_pos = current; switch (peek().t) { case token::comment: @@ -254,6 +266,7 @@ private: size_t pos0 = current; ++current; // consume {% ++current; // consume 'elif' + depth_guard guard(*this); alternate.push_back(parse_if_statement(pos0)); // nested If } else if (is_statement({"else"})) { ++current; // consume {% @@ -330,6 +343,7 @@ private: } statement_ptr parse_if_expression() { + depth_guard guard(*this); auto a = parse_logical_or_expression(); if (is_identifier("if")) { // Ternary expression @@ -373,6 +387,7 @@ private: statement_ptr parse_logical_negation_expression() { // Try parse unary operators if (is_identifier("not")) { + depth_guard guard(*this); size_t start_pos = current; auto op = next(); return mk_stmt(start_pos, op, parse_logical_negation_expression()); @@ -458,6 +473,7 @@ private: statement_ptr parse_unary_expression() { if (is(token::unary_operator)) { + depth_guard guard(*this); size_t start_pos = current; auto op = next(); return mk_stmt(start_pos, op, parse_unary_expression()); @@ -477,9 +493,11 @@ private: size_t start_pos = current; auto expr = mk_stmt(start_pos, std::move(callee), parse_args()); auto member = parse_member_expression(std::move(expr)); // foo.x().y - return is(token::open_paren) - ? parse_call_expression(std::move(member)) // foo.x()() - : std::move(member); + if (!is(token::open_paren)) { + return member; + } + depth_guard guard(*this); + return parse_call_expression(std::move(member)); // foo.x()() } statements parse_args() { diff --git a/common/jinja/runtime.cpp b/common/jinja/runtime.cpp index a53acf84c..7ec6606d7 100644 --- a/common/jinja/runtime.cpp +++ b/common/jinja/runtime.cpp @@ -3,6 +3,7 @@ #include "value.h" #include "utils.h" +#include #include #include #include @@ -52,32 +53,39 @@ static void ensure_key_type_allowed(const value & val) { // execute with error handling value statement::execute(context & ctx) const { + std::exception_ptr pending; + std::string message; try { return execute_impl(ctx); } catch (const continue_statement::signal & /* ex */) { - throw; + pending = std::current_exception(); } catch (const break_statement::signal & /* ex */) { - throw; + pending = std::current_exception(); } catch (const rethrown_exception & /* ex */) { - throw; + pending = std::current_exception(); } catch (const not_implemented_exception & /* ex */) { - throw; + pending = std::current_exception(); } catch (const std::exception & e) { const std::string & source = *ctx.src; if (source.empty()) { std::ostringstream oss; oss << "\nError executing " << type() << " at position " << pos << ": " << e.what(); - throw rethrown_exception(oss.str()); + message = oss.str(); } else { std::ostringstream oss; oss << "\n------------\n"; oss << "While executing " << type() << " at " << get_line_col(source, pos) << " in source:\n"; oss << peak_source(source, pos) << "\n"; oss << "Error: " << e.what(); - // throw as another exception to avoid repeated formatting - throw rethrown_exception(oss.str()); + message = oss.str(); } } + // throw outside the catch: MSVC keeps the unwound frames on the stack while a catch block runs + if (pending) { + std::rethrow_exception(pending); + } + // throw as another exception to avoid repeated formatting + throw rethrown_exception(message); } value identifier::execute_impl(context & ctx) const { @@ -747,7 +755,11 @@ value macro_statement::execute_impl(context & ctx) const { std::string name = cast_stmt(this->name)->val; const func_handler func = [this, name](const func_args & args) -> value { + if (args.ctx.macro_depth >= JINJA_RECURSIVE_LIMIT) { + throw std::runtime_error("macro calls are nested more than " + std::to_string(JINJA_RECURSIVE_LIMIT) + " levels deep"); + } context macro_ctx(args.ctx); // new scope for macro execution + macro_ctx.macro_depth = args.ctx.macro_depth + 1; bind_parameters(name, this->args, args, macro_ctx); @@ -778,7 +790,11 @@ value call_statement::execute_impl(context & ctx) const { context caller_ctx(ctx); // new scope for caller execution const func_handler func = [this, caller_ctx = std::move(caller_ctx)](const func_args & args) -> value { + if (args.ctx.macro_depth >= JINJA_RECURSIVE_LIMIT) { + throw std::runtime_error("macro calls are nested more than " + std::to_string(JINJA_RECURSIVE_LIMIT) + " levels deep"); + } context block_ctx(caller_ctx); // new scope for block execution + block_ctx.macro_depth = args.ctx.macro_depth + 1; bind_parameters("caller", this->caller_args, args, block_ctx); diff --git a/common/jinja/runtime.h b/common/jinja/runtime.h index c01902a1a..101e927c4 100644 --- a/common/jinja/runtime.h +++ b/common/jinja/runtime.h @@ -58,6 +58,8 @@ struct context { bool is_get_stats = false; // whether to collect stats + int macro_depth = 0; + visitor_fn visitor; // src is optional, used for error reporting @@ -78,6 +80,7 @@ struct context { src(parent.src), current_time(parent.current_time), is_get_stats(parent.is_get_stats), + macro_depth(parent.macro_depth), // inherit variables (for example, when entering a new scope) env(mk_val(*parent.env)) {} @@ -710,7 +713,7 @@ struct rethrown_exception : public std::exception { ////////////////////// -static void gather_string_parts_recursive(const value & val, value_string & parts) { +static void gather_string_parts_recursive(const value & val, value_string & parts, int depth = 0) { // TODO: probably allow print value_none as "None" string? currently this breaks some templates if (is_val(val)) { const auto & str_val = cast_val(val)->val_str; @@ -718,10 +721,12 @@ static void gather_string_parts_recursive(const value & val, value_string & part } else if (is_val(val) || is_val(val) || is_val(val)) { std::string str_val = val->as_string().str(); parts->val_str.append(str_val); + } else if (is_val(val) && depth >= JINJA_RECURSIVE_LIMIT) { + throw std::runtime_error("value is nested more than " + std::to_string(JINJA_RECURSIVE_LIMIT) + " levels deep"); } else if (is_val(val)) { auto items = cast_val(val)->as_array(); for (const auto & item : items) { - gather_string_parts_recursive(item, parts); + gather_string_parts_recursive(item, parts, depth + 1); } } } diff --git a/common/jinja/utils.h b/common/jinja/utils.h index de6947fc2..da3757a4a 100644 --- a/common/jinja/utils.h +++ b/common/jinja/utils.h @@ -6,6 +6,8 @@ #include #include +#define JINJA_RECURSIVE_LIMIT 128 + namespace jinja { static void string_replace_all(std::string & s, const std::string & search, const std::string & replace) { diff --git a/common/jinja/value.cpp b/common/jinja/value.cpp index e65c20a45..3aceca1e0 100644 --- a/common/jinja/value.cpp +++ b/common/jinja/value.cpp @@ -1441,8 +1441,6 @@ const func_builtins & value_undefined_t::get_builtins() const { // the runtime walks values with recursion, a deeper input overflows the stack -static constexpr int FROM_JSON_MAX_DEPTH = 128; - static value from_json(const common_json & j, bool mark_input, int depth = 0) { if (j.is_null()) { return mk_val(); @@ -1458,8 +1456,8 @@ static value from_json(const common_json & j, bool mark_input, int depth = 0) { str->mark_input(); } return str; - } else if ((j.is_array() || j.is_object()) && depth >= FROM_JSON_MAX_DEPTH) { - throw std::invalid_argument("JSON is nested more than " + std::to_string(FROM_JSON_MAX_DEPTH) + " levels deep"); + } else if ((j.is_array() || j.is_object()) && depth >= JINJA_RECURSIVE_LIMIT) { + throw std::invalid_argument("JSON is nested more than " + std::to_string(JINJA_RECURSIVE_LIMIT) + " levels deep"); } else if (j.is_array()) { auto arr = mk_val(); for (const auto & item : j) { @@ -1664,6 +1662,8 @@ static void value_to_json_internal(std::ostringstream & oss, const value & val, } } oss << "\""; + } else if ((is_val(val) || is_val(val)) && curr_lvl >= JINJA_RECURSIVE_LIMIT) { + throw std::runtime_error("value is nested more than " + std::to_string(JINJA_RECURSIVE_LIMIT) + " levels deep"); } else if (is_val(val)) { const auto & arr = val->as_array(); oss << "["; @@ -1729,17 +1729,20 @@ std::string value_to_string_repr(const value & val) { } // stats utility -void value_t::stats_t::mark_used(value & val, bool deep) { +void value_t::stats_t::mark_used(value & val, bool deep, int depth) { val->stats.used = true; + if (deep && (is_val(val) || is_val(val)) && depth >= JINJA_RECURSIVE_LIMIT) { + throw std::runtime_error("value is nested more than " + std::to_string(JINJA_RECURSIVE_LIMIT) + " levels deep"); + } if (deep) { if (is_val(val)) { for (auto & item : val->val_arr) { - mark_used(item, deep); + mark_used(item, deep, depth + 1); } } else if (is_val(val)) { for (auto & pair : val->val_obj) { - mark_used(pair.first, deep); - mark_used(pair.second, deep); + mark_used(pair.first, deep, depth + 1); + mark_used(pair.second, deep, depth + 1); } } } diff --git a/common/jinja/value.h b/common/jinja/value.h index 142943dfd..9bb1a87bd 100644 --- a/common/jinja/value.h +++ b/common/jinja/value.h @@ -119,7 +119,7 @@ struct value_t { // ops can be builtin calls or operators: "array_access", "object_access" std::set ops; // utility to recursively mark value and its children as used - static void mark_used(value & val, bool deep = false); + static void mark_used(value & val, bool deep = false, int depth = 0); } stats; value_t() = default; diff --git a/common/json-schema-to-grammar.cpp b/common/json-schema-to-grammar.cpp index 3349ae4fe..6f3ecc992 100644 --- a/common/json-schema-to-grammar.cpp +++ b/common/json-schema-to-grammar.cpp @@ -272,15 +272,12 @@ static bool is_reserved_name(const std::string & name) { return RESERVED_NAMES.find(name) != RESERVED_NAMES.end(); } -static std::regex INVALID_RULE_CHARS_RE("[^a-zA-Z0-9-]+"); static std::regex GRAMMAR_LITERAL_ESCAPE_RE("[\r\n\"\\\\]"); static std::regex GRAMMAR_RANGE_LITERAL_ESCAPE_RE("[\r\n\"\\]\\-\\\\]"); static std::unordered_map GRAMMAR_LITERAL_ESCAPES = { {'\r', "\\r"}, {'\n', "\\n"}, {'"', "\\\""}, {'-', "\\-"}, {']', "\\]"}, {'\\', "\\\\"} }; -static const int MAX_PATTERN_DEPTH = 100; - static std::unordered_set NON_LITERAL_SET = {'|', '.', '(', ')', '[', ']', '{', '}', '*', '+', '?', '^', '$'}; static std::unordered_set ESCAPED_IN_REGEXPS_BUT_NOT_IN_LITERALS = {'^', '$', '.', '[', ']', '(', ')', '|', '{', '}', '*', '+', '?'}; @@ -312,6 +309,22 @@ static std::string format_literal(const std::string & literal) { std::string gbnf_format_literal(const std::string & literal) { return format_literal(literal); } +// same as std::regex_replace with [^a-zA-Z0-9-]+ -> "-", no std::regex: libstdc++ recurses per char +std::string gbnf_format_rule_name(const std::string & name) { + std::string out; + bool in_invalid = false; + for (char c : name) { + if ((c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || (c >= '0' && c <= '9') || c == '-') { + out += c; + in_invalid = false; + } else if (!in_invalid) { + out += '-'; + in_invalid = true; + } + } + return out; +} + static size_t gbnf_escape_length(const std::string & pattern, size_t pos) { if (pos + 1 >= pattern.length() || pattern[pos] != '\\') { return 0; @@ -347,6 +360,17 @@ private: std::unordered_set _refs_being_resolved; std::vector _errors; std::vector _warnings; + int _depth = 0; + + struct depth_guard { + int & depth; + depth_guard(int & depth) : depth(depth) { + if (++depth > COMMON_RECURSIVE_LIMIT) { + throw std::invalid_argument("JSON schema is nested more than " + std::to_string(COMMON_RECURSIVE_LIMIT) + " levels deep"); + } + } + ~depth_guard() { depth--; } + }; template static const T & as(const common_chat_schema & node) { @@ -354,7 +378,7 @@ private: } std::string _add_rule(const std::string & name, const std::string & rule) { - std::string esc_name = regex_replace(name, INVALID_RULE_CHARS_RE, "-"); + std::string esc_name = gbnf_format_rule_name(name); if (_rules.find(esc_name) == _rules.end() || _rules[esc_name] == rule) { _rules[esc_name] = rule; return esc_name; @@ -482,7 +506,7 @@ private: } } paren_depth++; - if (paren_depth > MAX_PATTERN_DEPTH) { + if (paren_depth > COMMON_RECURSIVE_LIMIT) { throw unsupported_pattern("pattern nesting too deep"); } seq.emplace_back("(" + to_rule(transform()) + ")", false); @@ -644,7 +668,11 @@ private: std::string char_rule = _add_primitive("char", PRIMITIVE_RULES.at("char")); std::ostringstream out; out << "[\"] ( "; - std::function visit = [&](size_t idx) { + // one level per char of the longest name + std::function visit = [&](size_t idx, int depth) { + if (depth >= COMMON_RECURSIVE_LIMIT) { + throw std::invalid_argument("JSON schema property name is longer than " + std::to_string(COMMON_RECURSIVE_LIMIT) + " chars"); + } const auto & node = trie.nodes[idx]; std::string rejects; auto first = true; @@ -659,7 +687,7 @@ private: out << "[" << c << "]"; if (!trie.nodes[child].children.empty()) { out << " ("; - visit(child); + visit(child, depth + 1); out << ")"; } else { out << " " << char_rule << "+"; @@ -669,7 +697,7 @@ private: out << " | [^\"" << rejects << "] " << char_rule << "*"; } }; - visit(0); + visit(0, 0); out << " )"; if (trie.nodes[0].pattern < 0) { @@ -682,8 +710,7 @@ private: std::string _resolve_ref(const common_chat_schema_ref & schema) { auto it = schema.ref.find('#'); std::string ref_fragment = it != std::string::npos ? schema.ref.substr(it + 1) : schema.ref; - static const std::regex nonalphanumeric_regex(R"([^a-zA-Z0-9-]+)"); - std::string ref_name = "ref" + std::regex_replace(ref_fragment, nonalphanumeric_regex, "-"); + std::string ref_name = "ref" + gbnf_format_rule_name(ref_fragment); if (_rules.find(ref_name) == _rules.end() && _refs_being_resolved.find(schema.ref) == _refs_being_resolved.end()) { if (!schema.target) { _errors.push_back("Unresolved $ref " + schema.ref); @@ -833,10 +860,13 @@ public: std::unordered_set required; std::vector> properties; std::map enum_values; - std::function add_component = [&](const common_chat_schema & comp, bool is_required) { + std::function add_component = [&](const common_chat_schema & comp, bool is_required, int depth) { + if (depth >= COMMON_RECURSIVE_LIMIT) { + throw std::invalid_argument("JSON schema is nested more than " + std::to_string(COMMON_RECURSIVE_LIMIT) + " levels deep"); + } if (comp.kind() == common_chat_schema::KIND_REF) { if (const auto * target = as(comp).target) { - add_component(*target, is_required); + add_component(*target, is_required, depth + 1); } } else if (comp.kind() == common_chat_schema::KIND_OBJECT) { for (const auto & prop : as(comp).properties) { @@ -854,10 +884,10 @@ public: for (const auto & child : schema.children) { if (child->kind() == common_chat_schema::KIND_ANY_OF) { for (const auto & alt : as(*child).children) { - add_component(*alt, false); + add_component(*alt, false, 0); } } else { - add_component(*child, true); + add_component(*child, true, 0); } } if (!enum_values.empty()) { @@ -875,6 +905,8 @@ public: } std::string visit(const common_chat_schema & schema, const std::string & name) { + // the $ref targets are visited too, a $ref chain is not bounded by the builder + depth_guard guard(_depth); std::string rule_name = is_reserved_name(name) ? name + "-" : name.empty() ? "root" : name; std::string sub_name = name + (name.empty() ? "" : "-"); diff --git a/common/json-schema-to-grammar.h b/common/json-schema-to-grammar.h index b928c250b..b94874399 100644 --- a/common/json-schema-to-grammar.h +++ b/common/json-schema-to-grammar.h @@ -19,5 +19,6 @@ struct common_grammar_options { }; std::string gbnf_format_literal(const std::string & literal); +std::string gbnf_format_rule_name(const std::string & name); std::string build_grammar(const std::function & cb, const common_grammar_options & options = {}); diff --git a/common/json-schema.cpp b/common/json-schema.cpp index e1bfa79df..eb57a524a 100644 --- a/common/json-schema.cpp +++ b/common/json-schema.cpp @@ -9,6 +9,18 @@ #include #include +namespace { +struct depth_guard { + int & depth; + depth_guard(int & depth) : depth(depth) { + if (++depth > COMMON_RECURSIVE_LIMIT) { + throw std::invalid_argument("JSON schema is nested more than " + std::to_string(COMMON_RECURSIVE_LIMIT) + " levels deep"); + } + } + ~depth_guard() { depth--; } +}; +} // namespace + class common_chat_schema_builder { const common_json & root_; common_chat_schema_document & doc_; @@ -224,18 +236,11 @@ class common_chat_schema_builder { return node; } - // the builder and the code that walks the built schema recurse once per level, a deeper schema overflows the stack - static constexpr int max_depth = 128; - int depth_ = 0; + int depth_ = 0; common_chat_schema_ptr build_node(const common_json & schema, const std::string & path) { - if (depth_ >= max_depth) { - throw std::invalid_argument("JSON schema is nested more than " + std::to_string(max_depth) + " levels deep"); - } - depth_++; - auto node = build_node_impl(schema, path); - depth_--; - return node; + depth_guard guard(depth_); + return build_node_impl(schema, path); } common_chat_schema_ptr build_node_impl(const common_json & schema, const std::string & path) { @@ -379,7 +384,9 @@ static common_chat_schema::value_type json_type(const common_json & value) { return common_chat_schema::TYPE_OBJECT; } -static common_chat_schema::type_set value_types_impl(const common_chat_schema & s, std::unordered_set & visited) { +// depth counts the $ref, anyOf and allOf hops, a $ref chain is not bounded by the builder +static common_chat_schema::type_set value_types_impl(const common_chat_schema & s, std::unordered_set & visited, int & depth) { + depth_guard guard(depth); switch (s.kind()) { case common_chat_schema::KIND_ANY: return common_chat_schema::type_set::all(); @@ -413,21 +420,21 @@ static common_chat_schema::type_set value_types_impl(const common_chat_schema & // a cycle contributes no type, to be safe return {}; } - auto types = value_types_impl(*target, visited); + auto types = value_types_impl(*target, visited, depth); visited.erase(target); return types; } case common_chat_schema::KIND_ANY_OF: { common_chat_schema::type_set types; for (const auto & child : static_cast(s).children) { - types |= value_types_impl(*child, visited); + types |= value_types_impl(*child, visited, depth); } return types; } case common_chat_schema::KIND_ALL_OF: { auto types = common_chat_schema::type_set::all(); for (const auto & child : static_cast(s).children) { - types &= value_types_impl(*child, visited); + types &= value_types_impl(*child, visited, depth); } return types; } @@ -437,10 +444,12 @@ static common_chat_schema::type_set value_types_impl(const common_chat_schema & common_chat_schema::type_set common_chat_schema::value_types() const { std::unordered_set visited; - return value_types_impl(*this, visited); + int depth = 0; + return value_types_impl(*this, visited, depth); } -static bool may_be_string_impl(const common_chat_schema & s, std::unordered_set & visited) { +static bool may_be_string_impl(const common_chat_schema & s, std::unordered_set & visited, int & depth) { + depth_guard guard(depth); switch (s.kind()) { case common_chat_schema::KIND_STRING: return true; @@ -459,13 +468,13 @@ static bool may_be_string_impl(const common_chat_schema & s, std::unordered_set< if (!target || !visited.insert(target).second) { return false; } - bool result = may_be_string_impl(*target, visited); + bool result = may_be_string_impl(*target, visited, depth); visited.erase(target); return result; } case common_chat_schema::KIND_ANY_OF: for (const auto & child : static_cast(s).children) { - if (may_be_string_impl(*child, visited)) { + if (may_be_string_impl(*child, visited, depth)) { return true; } } @@ -477,7 +486,7 @@ static bool may_be_string_impl(const common_chat_schema & s, std::unordered_set< if (child->kind() == common_chat_schema::KIND_ANY) { continue; } - if (!may_be_string_impl(*child, visited)) { + if (!may_be_string_impl(*child, visited, depth)) { return false; } any_string = true; @@ -491,7 +500,8 @@ static bool may_be_string_impl(const common_chat_schema & s, std::unordered_set< bool common_chat_schema::may_be_string() const { std::unordered_set visited; - return may_be_string_impl(*this, visited); + int depth = 0; + return may_be_string_impl(*this, visited, depth); } const char * common_chat_schema::kind_name(node_kind kind) { diff --git a/common/peg-parser.cpp b/common/peg-parser.cpp index 8033f7e81..167a889e3 100644 --- a/common/peg-parser.cpp +++ b/common/peg-parser.cpp @@ -10,10 +10,12 @@ #include #include #include -#include #include #include +// The parser can run with a lower limit than COMMON_RECURSIVE_LIMIT +static constexpr int RULE_RECURSIVE_LIMIT = 64; + // Trick to catch missing branches template inline constexpr bool is_always_false_v = false; @@ -186,13 +188,17 @@ std::string common_peg_ast_node::sanitized_text() const { } void common_peg_ast_arena::visit(common_peg_ast_id id, const common_peg_ast_visitor & visitor) const { - if (id == COMMON_PEG_INVALID_AST_ID) { - return; - } - const auto & node = get(id); - visitor(node); - for (const auto & child : node.children) { - visit(child, visitor); + // no recursion: the AST is several nodes deep per level of the input + std::vector stack = { id }; + while (!stack.empty()) { + const auto cur = stack.back(); + stack.pop_back(); + if (cur == COMMON_PEG_INVALID_AST_ID) { + continue; + } + const auto & node = get(cur); + visitor(node); + stack.insert(stack.end(), node.children.rbegin(), node.children.rend()); } } @@ -736,8 +742,15 @@ struct parser_executor { } common_peg_parse_result operator()(const common_peg_rule_parser & p) { + // only rules recurse, a recursive rule is entered once per input level + if (ctx.rule_depth[p.child] >= RULE_RECURSIVE_LIMIT) { + throw common_peg_recursion_error("input is nested more than " + std::to_string(RULE_RECURSIVE_LIMIT) + " levels deep"); + } + // Parse the child + ctx.rule_depth[p.child]++; auto result = arena.parse(p.child, ctx, start_pos); + ctx.rule_depth[p.child]--; if (!result.fail()) { std::string_view text; @@ -819,6 +832,10 @@ common_peg_parse_result common_peg_arena::parse(common_peg_parse_context & ctx, if (root_ == COMMON_PEG_INVALID_PARSER_ID) { throw std::runtime_error("No root parser set"); } + + // Initialize rule depth for all parsers + ctx.rule_depth.assign(parsers_.size(), 0); + return parse(root_, ctx, start); } @@ -1069,11 +1086,6 @@ common_peg_parser operator|(const std::string & str, const common_peg_parser & p return operator|(str.c_str(), p); } -static std::string rule_name(const std::string & name) { - static const std::regex invalid_rule_chars_re("[^a-zA-Z0-9-]+"); - return std::regex_replace(name, invalid_rule_chars_re, "-"); -} - common_peg_parser_builder::common_peg_parser_builder() {} common_peg_parser common_peg_parser_builder::sequence(const std::vector & parsers) { @@ -1155,14 +1167,14 @@ common_peg_parser common_peg_parser_builder::schema(const common_peg_parser & p, } common_peg_parser common_peg_parser_builder::rule(const std::string & name, const common_peg_parser & p, bool trigger) { - auto clean_name = rule_name(name); + auto clean_name = gbnf_format_rule_name(name); auto rule_id = arena_.add_parser(common_peg_rule_parser{clean_name, p.id(), trigger}); arena_.add_rule(clean_name, rule_id); return ref(clean_name); } common_peg_parser common_peg_parser_builder::rule(const std::string & name, const std::function & builder_fn, bool trigger) { - auto clean_name = rule_name(name); + auto clean_name = gbnf_format_rule_name(name); if (arena_.has_rule(clean_name)) { return ref(clean_name); } diff --git a/common/peg-parser.h b/common/peg-parser.h index 9255d6eb9..6b348fab0 100644 --- a/common/peg-parser.h +++ b/common/peg-parser.h @@ -6,6 +6,7 @@ #include #include +#include #include #include #include @@ -182,6 +183,10 @@ inline common_peg_parse_flags operator~(common_peg_parse_flags a) { return static_cast(~int(a)); } +struct common_peg_recursion_error : std::invalid_argument { + using std::invalid_argument::invalid_argument; +}; + struct common_peg_parse_context { std::string input; // [h, e, l, l, o, _, w, o, r, l, d] std::vector tokens; // [id, -1, -1, -1, -1, id, -1, -1, -1, -1, -1] @@ -189,6 +194,7 @@ struct common_peg_parse_context { common_peg_ast_arena ast; int parse_depth; + std::vector rule_depth; // index: rule child id common_peg_parse_context(common_peg_parse_flags flags = COMMON_PEG_PARSE_FLAG_NONE) : flags(flags), parse_depth(0) {} diff --git a/skills/code-review/SKILL.md b/skills/code-review/SKILL.md index 2b15d18c4..b282a9164 100644 --- a/skills/code-review/SKILL.md +++ b/skills/code-review/SKILL.md @@ -58,6 +58,7 @@ Mandatory on every review; any finding here is **blocking**. Rule of thumb: GGUF - **Reused/reserved buffers:** recheck bounds after a buffer is shrunk or reused; watch `reserve()` then index-by-assumed-size, and header fields read before their length is checked. - **Server JSON ints:** clamp client-supplied integers (token/discard counts, offsets) to non-negative and an upper bound before they reach index/pointer arithmetic. - **RPC-deserialized fields:** treat every field (type/buffer/data/ne/nb/op_params) as hostile - validate before use. Null/zero buffers skipping validation, attacker data pointers, out-of-range type indices, and negative strides sign-extending past a corner-only assert all give arbitrary read/write. +- **Recursion on untrusted input:** a function that recurses once per level (or per char) of request JSON, model output, a grammar, a template or a regex pattern overflows the stack - the HTTP threads have 512 KB on macOS, so a few hundred levels are enough. Every such function needs its own depth bound, checked in the function itself, not a pre-walk at a caller (the next caller skips it): libcommon uses `COMMON_RECURSIVE_LIMIT` (`common/common.h`), jinja `JINJA_RECURSIVE_LIMIT` (`common/jinja/utils.h`), the grammar `GRAMMAR_RECURSIVE_LIMIT` (`src/llama-grammar.h`). Count the levels of the input, not the frames, and throw `std::invalid_argument` for request input so the server answers 400. Look for recursion the code does not spell out: mutual recursion, recursive `std::function` lambdas, visitor/virtual dispatch, `std::regex` (libstdc++ recurses per char of the input), and per-char helpers that recurse where a loop would do. - **Lifetime/UAF:** flag stored raw pointers to caller/temporary storage, cached pointers to buffers a later free releases, async ops whose source may drop before completion, and structures not invalidated on free/realloc. Null-check conditionally-built or "not required" tensors before dereferencing. ## Approach and design (when a new component/infra is introduced) diff --git a/src/llama-grammar.cpp b/src/llama-grammar.cpp index ffc7cf27c..4e096a368 100644 --- a/src/llama-grammar.cpp +++ b/src/llama-grammar.cpp @@ -8,10 +8,11 @@ #include #include #include +#include #include #define MAX_REPETITION_THRESHOLD 2000 -#define MAX_GRAMMAR_RECURSION_DEPTH 256 +#define GRAMMAR_RECURSIVE_LIMIT 256 // // helpers // @@ -598,8 +599,8 @@ const char * llama_grammar_parser::parse_sequence( n_prev_rules = 1; rule.push_back({LLAMA_GRETYPE_RULE_REF, ref_rule_id}); } else if (*pos == '(') { // grouping - if (depth >= MAX_GRAMMAR_RECURSION_DEPTH) { - throw std::runtime_error("grammar recursion depth limit exceeded (max " + std::to_string(MAX_GRAMMAR_RECURSION_DEPTH) + "), please reduce nesting depth"); + if (depth >= GRAMMAR_RECURSIVE_LIMIT) { + throw std::runtime_error("grammar recursion depth limit exceeded (max " + std::to_string(GRAMMAR_RECURSIVE_LIMIT) + "), please reduce nesting depth"); } // parse nested alternates into synthesized rule pos = parse_space(pos + 1, true); @@ -947,34 +948,19 @@ static void llama_grammar_advance_stack( } } -static llama_grammar_candidates llama_grammar_reject_candidates( - const llama_grammar_rules & rules, - const llama_grammar_stacks & stacks, - const llama_grammar_candidates & candidates) { - GGML_ASSERT(!stacks.empty()); // REVIEW - - if (candidates.empty()) { - return {}; - } - - auto rejects = llama_grammar_reject_candidates_for_stack(rules, stacks.front(), candidates); - - for (size_t i = 1, size = stacks.size(); i < size; ++i) { - rejects = llama_grammar_reject_candidates_for_stack(rules, stacks[i], rejects); - } - - return rejects; -} - static bool llama_grammar_detect_left_recursion( const llama_grammar_rules & rules, size_t rule_index, std::vector * rules_visited, std::vector * rules_in_progress, - std::vector * rules_may_be_empty) { + std::vector * rules_may_be_empty, + uint32_t depth = 0) { if ((*rules_in_progress)[rule_index]) { return true; } + if (depth > GRAMMAR_RECURSIVE_LIMIT) { + throw std::runtime_error("grammar rules are chained more than " + std::to_string(GRAMMAR_RECURSIVE_LIMIT) + " levels deep"); + } (*rules_in_progress)[rule_index] = true; @@ -1000,7 +986,7 @@ static bool llama_grammar_detect_left_recursion( bool recurse_into_nonterminal = true; for (size_t i = 0; i < rule.size(); i++) { if (rule[i].type == LLAMA_GRETYPE_RULE_REF && recurse_into_nonterminal) { - if (llama_grammar_detect_left_recursion(rules, (size_t)rule[i].value, rules_visited, rules_in_progress, rules_may_be_empty)) { + if (llama_grammar_detect_left_recursion(rules, (size_t)rule[i].value, rules_visited, rules_in_progress, rules_may_be_empty, depth + 1)) { return true; } if (!((*rules_may_be_empty)[(size_t)rule[i].value])) { @@ -1064,56 +1050,63 @@ void llama_grammar_accept(struct llama_grammar * grammar, uint32_t chr) { grammar->stacks = std::move(stacks_new); } -llama_grammar_candidates llama_grammar_reject_candidates_for_stack( +// checks the code point at depth of each candidate (an index into the input) against the stack +// returns true when some match: they go to next, and the stacks that follow to next_stacks +static bool llama_grammar_reject_step( const llama_grammar_rules & rules, const llama_grammar_stack & stack, - const llama_grammar_candidates & candidates) { - - llama_grammar_candidates rejects; - rejects.reserve(candidates.size()); + const llama_grammar_candidates & candidates, + size_t depth, + const std::vector & cur, + std::vector & rejects, + std::vector & next, + llama_grammar_stacks & next_stacks) { + rejects.reserve(cur.size()); if (stack.empty()) { - for (const auto & tok : candidates) { - if (*tok.code_points != 0 || tok.partial_utf8.n_remain != 0) { - rejects.push_back(tok); + for (const uint32_t i : cur) { + const auto & tok = candidates[i]; + if (tok.code_points[depth] != 0 || tok.partial_utf8.n_remain != 0) { + rejects.push_back(i); } } - return rejects; + return false; } const llama_grammar_element * stack_pos = stack.back(); // if the top of the stack is a token rule, then we only need to check the token id if (stack_pos->type == LLAMA_GRETYPE_TOKEN || stack_pos->type == LLAMA_GRETYPE_TOKEN_NOT) { - for (const auto & tok : candidates) { - if (*tok.code_points == 0) { + for (const uint32_t i : cur) { + const auto & tok = candidates[i]; + if (tok.code_points[depth] == 0) { // reached the end of a token consumed by char rules, reject iff it ended // in a partial response if (tok.partial_utf8.n_remain != 0) { - rejects.push_back(tok); + rejects.push_back(i); } } else if (!llama_grammar_match_token(stack_pos, tok.id)) { - rejects.push_back(tok); + rejects.push_back(i); } } - return rejects; + return false; } - llama_grammar_candidates next_candidates; - next_candidates.reserve(candidates.size()); + next.reserve(cur.size()); - for (const auto & tok : candidates) { - if (*tok.code_points == 0) { + for (const uint32_t i : cur) { + const auto & tok = candidates[i]; + if (tok.code_points[depth] == 0) { // reached end of full codepoints in token, reject iff it ended in a partial sequence // that cannot satisfy this position in grammar if (tok.partial_utf8.n_remain != 0 && !llama_grammar_match_partial_char(stack_pos, tok.partial_utf8)) { - rejects.push_back(tok); + rejects.push_back(i); } - } else if (llama_grammar_match_char(stack_pos, *tok.code_points).first) { - next_candidates.push_back({ tok.index, tok.code_points + 1, tok.partial_utf8, tok.id }); + } else if (llama_grammar_match_char(stack_pos, tok.code_points[depth]).first) { + next.push_back(i); } else { - rejects.push_back(tok); + rejects.push_back(i); } } @@ -1124,17 +1117,80 @@ llama_grammar_candidates llama_grammar_reject_candidates_for_stack( if (!llama_grammar_is_end_of_sequence(stack_pos_after)) { stack_after.push_back(stack_pos_after); } - llama_grammar_stacks next_stacks; llama_grammar_advance_stack(rules, stack_after, next_stacks); - auto next_rejects = llama_grammar_reject_candidates(rules, next_stacks, next_candidates); - for (const auto & tok : next_rejects) { - rejects.push_back({ tok.index, tok.code_points - 1, tok.partial_utf8, tok.id }); + return true; +} + +// the candidates that no stack accepts, each stack narrows the rejects of the previous one +// no recursion: a level per code point, a token from the GGUF can be long +static llama_grammar_candidates llama_grammar_reject_candidates( + const llama_grammar_rules & rules, + const llama_grammar_stacks & stacks, + const llama_grammar_candidates & candidates) { + GGML_ASSERT(!stacks.empty()); // REVIEW + + // a level checks the candidates that reach its code point against its stacks in turn + struct level { + llama_grammar_stacks stacks; + size_t i; // the stack in progress + std::vector cur; // the candidates rejected by stacks[0, i) + std::vector rejects; // the rejects of stacks[i] at this code point + }; + + std::vector all(candidates.size()); + std::iota(all.begin(), all.end(), 0u); + + std::vector levels; + levels.push_back({ stacks, 0, std::move(all), {} }); + + while (true) { + level & lv = levels.back(); + + if (lv.i == lv.stacks.size() || lv.cur.empty()) { + if (levels.size() == 1) { + break; + } + // the rejects at the next code point belong to stacks[i] of the parent level + level & parent = levels[levels.size() - 2]; + parent.rejects.insert(parent.rejects.end(), lv.cur.begin(), lv.cur.end()); + levels.pop_back(); + parent.cur = std::move(parent.rejects); + parent.rejects = {}; + parent.i++; + continue; + } + + std::vector rejects; + std::vector next; + llama_grammar_stacks next_stacks; + + if (!llama_grammar_reject_step(rules, lv.stacks[lv.i], candidates, levels.size() - 1, lv.cur, rejects, next, next_stacks)) { + lv.cur = std::move(rejects); + lv.i++; + continue; + } + + GGML_ASSERT(!next_stacks.empty()); + lv.rejects = std::move(rejects); + levels.push_back({ std::move(next_stacks), 0, std::move(next), {} }); } + llama_grammar_candidates rejects; + rejects.reserve(levels[0].cur.size()); + for (const uint32_t i : levels[0].cur) { + rejects.push_back(candidates[i]); + } return rejects; } +llama_grammar_candidates llama_grammar_reject_candidates_for_stack( + const llama_grammar_rules & rules, + const llama_grammar_stack & stack, + const llama_grammar_candidates & candidates) { + return llama_grammar_reject_candidates(rules, llama_grammar_stacks{ stack }, candidates); +} + //////////////////// struct llama_grammar * llama_grammar_init_impl( @@ -1169,14 +1225,19 @@ struct llama_grammar * llama_grammar_init_impl( std::vector rules_visited(n_rules); std::vector rules_in_progress(n_rules); std::vector rules_may_be_empty(n_rules); - for (size_t i = 0; i < n_rules; i++) { - if (rules_visited[i]) { - continue; - } - if (llama_grammar_detect_left_recursion(vec_rules, i, &rules_visited, &rules_in_progress, &rules_may_be_empty)) { - LLAMA_LOG_ERROR("unsupported grammar, left recursion detected for nonterminal at index %zu", i); - return nullptr; + try { + for (size_t i = 0; i < n_rules; i++) { + if (rules_visited[i]) { + continue; + } + if (llama_grammar_detect_left_recursion(vec_rules, i, &rules_visited, &rules_in_progress, &rules_may_be_empty)) { + LLAMA_LOG_ERROR("unsupported grammar, left recursion detected for nonterminal at index %zu", i); + return nullptr; + } } + } catch (const std::exception & err) { + LLAMA_LOG_ERROR("unsupported grammar: %s\n", err.what()); + return nullptr; } // loop over alternates of start rule to build initial stacks @@ -1262,14 +1323,19 @@ struct llama_grammar * llama_grammar_init_impl( std::vector rules_visited(n_rules); std::vector rules_in_progress(n_rules); std::vector rules_may_be_empty(n_rules); - for (size_t i = 0; i < n_rules; i++) { - if (rules_visited[i]) { - continue; - } - if (llama_grammar_detect_left_recursion(vec_rules, i, &rules_visited, &rules_in_progress, &rules_may_be_empty)) { - LLAMA_LOG_ERROR("unsupported grammar, left recursion detected for nonterminal at index %zu\n", i); - return nullptr; + try { + for (size_t i = 0; i < n_rules; i++) { + if (rules_visited[i]) { + continue; + } + if (llama_grammar_detect_left_recursion(vec_rules, i, &rules_visited, &rules_in_progress, &rules_may_be_empty)) { + LLAMA_LOG_ERROR("unsupported grammar, left recursion detected for nonterminal at index %zu\n", i); + return nullptr; + } } + } catch (const std::exception & err) { + LLAMA_LOG_ERROR("unsupported grammar: %s\n", err.what()); + return nullptr; } // loop over alternates of start rule to build initial stacks diff --git a/src/llama-vocab.cpp b/src/llama-vocab.cpp index 05334e1fc..65b569762 100644 --- a/src/llama-vocab.cpp +++ b/src/llama-vocab.cpp @@ -29,31 +29,25 @@ struct naive_trie { naive_trie() : has_value(false), value(0) { } + // no recursion: a token from the GGUF can be long void insert(const char * key, size_t len, int32_t value = 0) { - if (len == 0) { - this->has_value = true; - this->value = value; - return; - } - char c = key[0]; - auto res = children.find(c); - if (res != children.end()) { - res->second.insert(key + 1, len - 1, value); - } else { - auto res = children.insert(std::make_pair(c, naive_trie())); - res.first->second.insert(key + 1, len - 1, value); + naive_trie * node = this; + for (size_t i = 0; i < len; i++) { + node = &node->children[key[i]]; } + node->has_value = true; + node->value = value; } std::pair get_longest_prefix(const char * key, size_t len, size_t offset = 0) const { - if (len == 0 || offset == len) { - return std::make_pair(key, offset); + const naive_trie * node = this; + while (offset < len) { + auto res = node->children.find(key[offset]); + if (res == node->children.end()) { + break; + } + node = &res->second; + offset++; } - char c = key[offset]; - auto res = children.find(c); - if (res != children.end()) { - return res->second.get_longest_prefix(key, len, offset + 1); - } - return std::make_pair(key, offset); } const struct naive_trie * traverse(const char c) const { @@ -174,30 +168,37 @@ struct llm_tokenizer_spm_session { } private: + // no recursion: the merges of a long token can be deep void resegment(llm_symbol & symbol, std::vector & output) { - auto text = std::string(symbol.text, symbol.n); - auto token = vocab.text_to_token(text); + std::vector todo = { &symbol }; + while (!todo.empty()) { + const llm_symbol & cur = *todo.back(); + todo.pop_back(); - // Do we need to support is_unused? - if (token != LLAMA_TOKEN_NULL) { - output.push_back(token); - return; - } + auto text = std::string(cur.text, cur.n); + auto token = vocab.text_to_token(text); - const auto p = rev_merge.find(text); - - if (p == rev_merge.end()) { - // output any symbols that did not form tokens as bytes. - output.reserve(output.size() + symbol.n); - for (int j = 0; j < (int)symbol.n; ++j) { - llama_token id = vocab.byte_to_token(symbol.text[j]); - output.push_back(id); + // Do we need to support is_unused? + if (token != LLAMA_TOKEN_NULL) { + output.push_back(token); + continue; } - return; - } - resegment(symbols[p->second.first], output); - resegment(symbols[p->second.second], output); + const auto p = rev_merge.find(text); + + if (p == rev_merge.end()) { + // output any symbols that did not form tokens as bytes. + output.reserve(output.size() + cur.n); + for (int j = 0; j < (int)cur.n; ++j) { + llama_token id = vocab.byte_to_token(cur.text[j]); + output.push_back(id); + } + continue; + } + + todo.push_back(&symbols[p->second.second]); + todo.push_back(&symbols[p->second.first]); + } } void try_add_bigram(int left, int right) { diff --git a/tests/peg-parser/test-json-parser.cpp b/tests/peg-parser/test-json-parser.cpp index ec7c2e668..9752da461 100644 --- a/tests/peg-parser/test-json-parser.cpp +++ b/tests/peg-parser/test-json-parser.cpp @@ -130,4 +130,23 @@ void test_json_parser(testing &t) { t.assert_true("fail", result.fail()); }); }); + + // json-value is entered once per level + t.test("deep JSON", [](testing &t) { + auto json = build_peg_parser([](common_peg_parser_builder & p) { return p.json(); }); + + const auto parse = [&](size_t depth) { + std::string input = std::string(depth, '[') + std::string(depth, ']'); + common_peg_parse_context ctx(input); + try { + return json.parse(ctx).success() ? std::string() : std::string("fail"); + } catch (const common_peg_recursion_error & e) { + return std::string(e.what()); + } + }; + + t.assert_equal("64 levels", std::string(), parse(64)); + t.assert_true("65 levels", parse(65).find("nested more than 64 levels deep") != std::string::npos); + t.assert_true("100000 levels", parse(100000).find("nested more than 64 levels deep") != std::string::npos); + }); } diff --git a/tests/test-grammar-integration.cpp b/tests/test-grammar-integration.cpp index 2af96f8b4..75bfca645 100644 --- a/tests/test-grammar-integration.cpp +++ b/tests/test-grammar-integration.cpp @@ -879,6 +879,18 @@ static void test_failure_left_recursion() { empty ::= "blah" | )"""; assert(test_build_grammar_fails(hardest_str)); + const auto chain = [](int n) { + std::string s = "root ::= r0\n"; + for (int i = 0; i < n; i++) { + s += "r" + std::to_string(i) + " ::= r" + std::to_string(i + 1) + "\n"; + } + return s + "r" + std::to_string(n) + " ::= \"a\"\n"; + }; + llama_grammar * grammar = build_grammar(chain(100)); + assert(grammar != nullptr); + llama_grammar_free_impl(grammar); + assert(test_build_grammar_fails(chain(1000))); + fprintf(stderr, " ✅︎ Passed\n"); } @@ -1476,6 +1488,25 @@ static void test_json_schema() { // R"""({"productId": 1, "productName": "A green door", "price": 12.50, "tags": ["home", "green", "home"]})""", } ); + + // the deepest schema of json_schema_to_grammar must give a valid grammar + std::string deep_schema; + for (int i = 0; i < 127; i++) { + deep_schema += R"({"type": "array", "items": )"; + } + deep_schema += R"({"type": "string"})" + std::string(127, '}'); + test_schema( + "127 nested arrays", + deep_schema, + // Passing strings + { + std::string(127, '[') + R"("a")" + std::string(127, ']'), + }, + // Failing strings + { + std::string(126, '[') + R"("a")" + std::string(126, ']'), + } + ); } int main() { diff --git a/tests/test-jinja.cpp b/tests/test-jinja.cpp index eb05f8780..2a0117284 100644 --- a/tests/test-jinja.cpp +++ b/tests/test-jinja.cpp @@ -17,6 +17,7 @@ using json = common_json; static void test_template(testing & t, const std::string & name, const std::string & tmpl, const json & vars, const std::string & expect); +static void test_template_error(testing & t, const std::string & name, const std::string & tmpl, const std::string & expect); static void test_whitespace_control(testing & t); static void test_conditionals(testing & t); @@ -606,6 +607,17 @@ static void test_expressions(testing & t) { {{"cond", true}}, "yes" ); + + test_template(t, "nested parentheses", + "{{ " + std::string(30, '(') + "1" + std::string(30, ')') + " }}", + json::object(), + "1" + ); + + test_template_error(t, "too deeply nested parentheses", + "{{ " + std::string(1000, '(') + "1" + std::string(1000, ')') + " }}", + "template is nested more than 128 levels deep" + ); } static void test_set_statement(testing & t) { @@ -1029,6 +1041,11 @@ static void test_filters(testing & t) { {{"x", nullptr}}, "None" ); + + test_template_error(t, "tojson of a list that contains itself", + "{% set arr = [] %}{% set _ = arr.append(arr) %}{{ arr|tojson }}", + "value is nested more than 128 levels deep" + ); } static void test_literals(testing & t) { @@ -1291,6 +1308,17 @@ static void test_macros(testing & t) { json::object(), "Dict {\n root1: 1\n root2: {\n nest1: 1\n nest2: {\n nest3: 2\n }\n }\n}" ); + + test_template(t, "macro recursion", + "{% macro f(n) %}{% if n < 100 %}{{ f(n + 1) }}{% else %}{{ n }}{% endif %}{% endmacro %}{{ f(0) }}", + json::object(), + "100" + ); + + test_template_error(t, "endless macro recursion", + "{% macro f(n) %}{{ f(n + 1) }}{% endmacro %}{{ f(0) }}", + "macro calls are nested more than 128 levels deep" + ); } static void test_namespace(testing & t) { @@ -2554,6 +2582,22 @@ static void test_string_parts(testing & t) { } +static void test_template_error(testing & t, const std::string & name, const std::string & tmpl, const std::string & expect) { + t.test(name, [&tmpl, &expect](testing & t) { + std::string error; + try { + jinja::lexer lexer; + jinja::program ast = jinja::parse_from_tokens(lexer.tokenize(tmpl)); + jinja::context ctx(tmpl); + jinja::runtime runtime(ctx); + runtime.gather_string_parts(runtime.execute(ast)); + } catch (const std::exception & e) { + error = e.what(); + } + t.assert_true("expected error: " + expect + ", got: " + error, error.find(expect) != std::string::npos); + }); +} + static void test_template_cpp(testing & t, const std::string & name, const std::string & tmpl, const json & vars, const std::string & expect) { t.test(name, [&tmpl, &vars, &expect](testing & t) { jinja::lexer lexer; diff --git a/tests/test-json-schema-to-grammar.cpp b/tests/test-json-schema-to-grammar.cpp index d5873c19d..c7d971b29 100755 --- a/tests/test-json-schema-to-grammar.cpp +++ b/tests/test-json-schema-to-grammar.cpp @@ -1513,6 +1513,23 @@ static void test_all(const std::string & title, std::function