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

@@ -1,10 +1,13 @@
 #include "testing.h"
 
 #include "../tools/server/server-common.h"
+#include "../tools/server/server-decision.h"
 
 #include <cstdlib>
+#include <regex>
 #include <string>
 #include <thread>
+#include <utility>
 #include <vector>
 
 // 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<std::pair<std::string, std::string>> 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();
 }

tools/server/server-decision.cpp

@@ -5,7 +5,6 @@
 #include <algorithm>
 #include <cctype>
 #include <cmath>
-#include <regex>
 #include <stdexcept>
 
 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<std::string>(), re_special, "<\xC2\xA6$1\xC2\xA6>");
+        return server_decision_escape_special_tokens(val.get<std::string>());
     }
     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;
     }

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_task> server_decision_group_tasks(std::vector<server_task> && 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);