diff options
Diffstat (limited to 'src/mongo/util/str_escape.cpp')
| -rw-r--r-- | src/mongo/util/str_escape.cpp | 120 |
1 files changed, 29 insertions, 91 deletions
diff --git a/src/mongo/util/str_escape.cpp b/src/mongo/util/str_escape.cpp index c42a916bab0..d191fb92252 100644 --- a/src/mongo/util/str_escape.cpp +++ b/src/mongo/util/str_escape.cpp @@ -37,82 +37,45 @@ namespace mongo::str { namespace { constexpr char kHexChar[] = "0123456789abcdef"; -// Appends the bytes in the range [begin, end) to the output buffer, -// which can either be a fmt::memory_buffer, or a std::string. -template <typename Buffer, typename Iterator> -void appendBuffer(Buffer& buffer, Iterator begin, Iterator end) { - buffer.append(begin, end); -} - // 'singleHandler' Function to write a valid single byte UTF-8 sequence with desired escaping. // 'invalidByteHandler' Function to write a byte of invalid UTF-8 encoding // 'twoEscaper' Function to write a valid two byte UTF-8 sequence with desired escaping, for C1 // control codes. -// 'maxLength' Max length to write into output buffer; A value of std::string::npos means unbounded. -// An escape sequence will not be written if appending the entire sequence will exceed this limit. -// 'wouldWrite' Output to contain the total bytes that would have been written to the buffer if no -// size limit is in place. -// // All these functions take a function object as their first parameter to perform the // writing of any escaped data. This function expects the number of handled bytes as its first // parameter and the corresponding escaped string as the second. They are templates to they can be // inlined. -template <typename Buffer, - typename SingleByteHandler, - typename InvalidByteHandler, - typename TwoByteEscaper> -void escape(Buffer& buffer, +template <typename SingleByteHandler, typename InvalidByteHandler, typename TwoByteEscaper> +void escape(fmt::memory_buffer& buffer, StringData str, SingleByteHandler singleHandler, InvalidByteHandler invalidByteHandler, - TwoByteEscaper twoEscaper, - size_t maxLength, - size_t* wouldWrite) { - // The range [inFirst, it) contains input that does not need to be escaped and that has not been + TwoByteEscaper twoEscaper) { + // The range [begin, it) contains input that does not need to be escaped and that has not been // written to output yet. - // The range [it, inLast) contains remaining input to scan. 'inFirst' is pointing to the - // beginning of the input that has not yet been written to 'escaped'. 'it' is pointing to the - // beginning of the unicode code point we're currently processing in the while-loop below. - // 'inLast' is the end of the input sequence. - auto inFirst = str.begin(); - auto inLast = str.end(); - auto it = inFirst; - size_t cap = maxLength; - size_t total = 0; + // The range [it end) contains remaining input to scan 'begin' is pointing to the beginning of + // the input that has not yet been written to 'escaped'. + // 'it' is pointing to the beginning of the unicode code point we're currently processing in the + // while-loop below. 'end' is the end of the input sequence. + auto begin = str.begin(); + auto it = str.begin(); + auto end = str.end(); // Writes an escaped sequence to output after flushing pending input that does not need to be // escaped. 'it' is assumed to be at the beginning of the input sequence represented by the // escaped data. // 'numHandled' the number of bytes of unescaped data being written escaped in 'escapeSequence' auto flushAndWrite = [&](size_t numHandled, StringData escapeSequence) { - // Appends the range [wFirst, wLast) to the output if the result is within the max length. - // 'canTruncate' controls the behavior if appending the entire range would exceed the limit. - // If true, this appends input up to the length limit. Otherwise, none is appended. - auto boundedWrite = [&](auto wFirst, auto wLast, bool canTruncate) { - size_t len = std::distance(wFirst, wLast); - total += len; - if (maxLength != std::string::npos) { - if (len > cap) { - if (!canTruncate) { - cap = 0; - } - len = cap; - } - cap -= len; - } - appendBuffer(buffer, wFirst, wFirst + len); - }; - // Flush range of unmodified input - boundedWrite(inFirst, it, true); - inFirst = it + numHandled; + buffer.append(begin, it); + begin = it + numHandled; // Write escaped data - boundedWrite(escapeSequence.begin(), escapeSequence.end(), false); + buffer.append(escapeSequence.rawData(), escapeSequence.rawData() + escapeSequence.size()); }; auto isValidCodePoint = [&](auto pos, int len) { - return std::distance(pos, inLast) >= len && + return std::distance(pos, end) >= len && std::all_of(pos + 1, pos + len, [](uint8_t c) { return (c >> 6) == 0b10; }); }; @@ -135,7 +98,7 @@ void escape(Buffer& buffer, auto writeInvalid = [&](uint8_t c) { invalidByteHandler(flushAndWrite, c); }; - while (it != inLast) { + while (it != end) { uint8_t c = *it; bool bit7 = (c >> 7) & 1; if (MONGO_likely(!bit7)) { @@ -193,15 +156,10 @@ void escape(Buffer& buffer, } } // Write last block - flushAndWrite(0, {}); - if (wouldWrite) { - *wouldWrite = total; - } + buffer.append(begin, it); } } // namespace - -template <typename Buffer> -void escapeForTextCommon(Buffer& buffer, StringData str, size_t maxLength, size_t* wouldWrite) { +void escapeForText(fmt::memory_buffer& buffer, StringData str) { auto singleByteHandler = [](const auto& writer, uint8_t unescaped) { switch (unescaped) { case '\0': @@ -329,26 +287,16 @@ void escapeForTextCommon(Buffer& buffer, StringData str, size_t maxLength, size_ str, std::move(singleByteHandler), std::move(invalidByteHandler), - std::move(twoByteEscaper), - maxLength, - wouldWrite); + std::move(twoByteEscaper)); } -void escapeForText(fmt::memory_buffer& buffer, - StringData str, - size_t maxLength, - size_t* wouldWrite) { - escapeForTextCommon(buffer, str, maxLength, wouldWrite); +std::string escapeForText(StringData str) { + fmt::memory_buffer buffer; + escapeForText(buffer, str); + return fmt::to_string(buffer); } -std::string escapeForText(StringData str, size_t maxLength, size_t* wouldWrite) { - std::string buffer; - escapeForTextCommon(buffer, str, maxLength, wouldWrite); - return buffer; -} - -template <typename Buffer> -void escapeForJSONCommon(Buffer& buffer, StringData str, size_t maxLength, size_t* wouldWrite) { +void escapeForJSON(fmt::memory_buffer& buffer, StringData str) { auto singleByteHandler = [](const auto& writer, uint8_t unescaped) { switch (unescaped) { case '\0': @@ -479,21 +427,11 @@ void escapeForJSONCommon(Buffer& buffer, StringData str, size_t maxLength, size_ str, std::move(singleByteHandler), std::move(invalidByteHandler), - std::move(twoByteEscaper), - maxLength, - wouldWrite); + std::move(twoByteEscaper)); } - -void escapeForJSON(fmt::memory_buffer& buffer, - StringData str, - size_t maxLength, - size_t* wouldWrite) { - escapeForJSONCommon(buffer, str, maxLength, wouldWrite); -} - -std::string escapeForJSON(StringData str, size_t maxLength, size_t* wouldWrite) { - std::string buffer; - escapeForJSONCommon(buffer, str, maxLength, wouldWrite); - return buffer; +std::string escapeForJSON(StringData str) { + fmt::memory_buffer buffer; + escapeForJSON(buffer, str); + return fmt::to_string(buffer); } } // namespace mongo::str |