From 04f7226a2c12c7f382c545a6db5a034266aa53cb Mon Sep 17 00:00:00 2001 From: Yuri Khrustalev Date: Tue, 13 Oct 2026 19:17:11 -0400 Subject: [PATCH] server : escape the special tokens without std::regex (#30243) The d1-omni and kev prompts escape "<|name|>" in the request text with std::regex_replace. libstdc++ recurses for each char of the name, so a state of "<|" and 30000 letters overflows the stack of the HTTP thread and llama-server dies with SIGSEGV. server_decision_escape_special_tokens() writes the same text with a loop. test-server-component compares it with the regex on every string of up to 6 chars of "<|>a-". --- tests/test-server-component.cpp | 66 ++++++++++++++++++++++++++++++++ tools/server/server-decision.cpp | 37 +++++++++++++++--- tools/server/server-decision.h | 3 ++ 3 files changed, 100 insertions(+), 6 deletions(-) diff --git a/tests/test-server-component.cpp b/tests/test-server-component.cpp index d2d103a45..aa353e190 100644 --- a/tests/test-server-component.cpp +++ b/tests/test-server-component.cpp @@ -1,10 +1,13 @@ #include "testing.h" #include "../tools/server/server-common.h" +#include "../tools/server/server-decision.h" #include +#include #include #include +#include #include // unit tests for server internal components, no model required @@ -370,6 +373,68 @@ static void test_server_pipe(testing & t) { }); } +// +// server_decision_escape_special_tokens +// + +// "<|name|>" with U+00A6 in place of "|" +static std::string escaped(const std::string & name) { + return "<\xC2\xA6" + name + "\xC2\xA6>"; +} + +static void test_decision_escape(testing & t) { + t.test("cases", [](testing & t) { + const std::vector> cases = { + { "", "" }, + { "no tokens", "no tokens" }, + { "<|im_start|>user", escaped("im_start") + "user" }, + { "<|a_1|><|B|>", escaped("a_1") + escaped("B") }, + { "<|<|a|>", "<|" + escaped("a") }, + { "<|a|>|>", escaped("a") + "|>" }, + { "<||>", "<||>" }, + { "<|a b|>", "<|a b|>" }, + { "<|a|b|>", "<|a|b|>" }, + { "<|caf\xC3\xA9|>", "<|caf\xC3\xA9|>" }, + { "<|abc", "<|abc" }, + }; + for (const auto & [text, expected] : cases) { + t.assert_equal(text, expected, server_decision_escape_special_tokens(text)); + } + }); + + // every string of up to 6 chars from "<|>a-" gives the text of the regex + t.test("same as std::regex", [](testing & t) { + static const std::regex re_special("<\\|([A-Za-z0-9_]+)\\|>"); + const std::string chars = "<|>a-"; + size_t n_strings = 0; + size_t n_diff = 0; + for (size_t len = 0; len <= 6; len++) { + size_t n = 1; + for (size_t k = 0; k < len; k++) { + n *= chars.size(); + } + for (size_t idx = 0; idx < n; idx++) { + std::string text; + for (size_t k = 0, v = idx; k < len; k++, v /= chars.size()) { + text += chars[v % chars.size()]; + } + if (server_decision_escape_special_tokens(text) != std::regex_replace(text, re_special, "<\xC2\xA6$1\xC2\xA6>")) { + n_diff++; + } + n_strings++; + } + } + t.assert_equal("strings that differ, of " + std::to_string(n_strings), (size_t) 0, n_diff); + }); + + // std::regex overflowed the stack on these + t.test("long input", [](testing & t) { + const std::string name(1000000, 'a'); + t.assert_true("not closed", server_decision_escape_special_tokens("<|" + name) == "<|" + name); + t.assert_true("closed", server_decision_escape_special_tokens("<|" + name + "|>") == escaped(name)); + }); +} + int main(int argc, char ** argv) { testing t; @@ -388,6 +453,7 @@ int main(int argc, char ** argv) { t.test("json_helpers", test_json_helpers); t.test("lora_helpers", test_lora_helpers); t.test("server_pipe", test_server_pipe); + t.test("decision_escape", test_decision_escape); return t.summary(); } diff --git a/tools/server/server-decision.cpp b/tools/server/server-decision.cpp index bd2c9c4aa..b3c19e1f0 100644 --- a/tools/server/server-decision.cpp +++ b/tools/server/server-decision.cpp @@ -5,7 +5,6 @@ #include #include #include -#include #include static const char * decision_question_type_name(server_decision_question_type type) { @@ -449,17 +448,43 @@ static std::string decision_kev_render(const json & val, int indent = 0) { return val.dump(); } +// same result as std::regex_replace with <\|([A-Za-z0-9_]+)\|> +// std::regex is not used: libstdc++ recurses for each char, a long input overflows the stack +std::string server_decision_escape_special_tokens(const std::string & text) { + const auto is_name = [](char c) { + return (c >= 'A' && c <= 'Z') || (c >= 'a' && c <= 'z') || (c >= '0' && c <= '9') || c == '_'; + }; + std::string out; + size_t n_done = 0; // the chars of text that are in out + for (size_t i = text.find("<|"); i != std::string::npos; i = text.find("<|", std::max(i + 1, n_done))) { + size_t j = i + 2; + while (j < text.size() && is_name(text[j])) { + j++; + } + if (j > i + 2 && text.compare(j, 2, "|>") == 0) { + out.append(text, n_done, i - n_done); + out += "<\xC2\xA6"; + out.append(text, i + 2, j - i - 2); + out += "\xC2\xA6>"; + n_done = j + 2; + } + } + if (n_done == 0) { + return text; + } + out.append(text, n_done, std::string::npos); + return out; +} + // kev text input: special tokens written in the text must not be parsed as such static std::string decision_kev_text(const json & val) { - static const std::regex re_special("<\\|([A-Za-z0-9_]+)\\|>"); - return std::regex_replace(decision_kev_render(val), re_special, "<\xC2\xA6$1\xC2\xA6>"); + return server_decision_escape_special_tokens(decision_kev_render(val)); } // lfm2-d1-omni: special tokens written in the input must not be parsed as such, in keys too (d1-omni prompt.py: escape) static json decision_d1omni_escape(const json & val) { - static const std::regex re_special("<\\|([A-Za-z0-9_]+)\\|>"); if (val.is_string()) { - return std::regex_replace(val.get(), re_special, "<\xC2\xA6$1\xC2\xA6>"); + return server_decision_escape_special_tokens(val.get()); } if (val.is_array()) { json out = json::array(); @@ -471,7 +496,7 @@ static json decision_d1omni_escape(const json & val) { if (val.is_object()) { json out = json::object(); for (const auto & [key, item] : val.items()) { - out[std::regex_replace(key, re_special, "<\xC2\xA6$1\xC2\xA6>")] = decision_d1omni_escape(item); + out[server_decision_escape_special_tokens(key)] = decision_d1omni_escape(item); } return out; } diff --git a/tools/server/server-decision.h b/tools/server/server-decision.h index 2753daa82..bfe9fd0ee 100644 --- a/tools/server/server-decision.h +++ b/tools/server/server-decision.h @@ -161,3 +161,6 @@ private: // each group is one parent and its children, it takes at most n_slots slots // note: the order of the tasks is preserved std::vector server_decision_group_tasks(std::vector && tasks, size_t n_slots); + +// write "<|name|>" with U+00A6 in place of "|", so text from the request is not parsed as special tokens +std::string server_decision_escape_special_tokens(const std::string & text);