Branch data Line data Source code
1 : : /*
2 : : * Copyright 2025 Jinwoo Sung
3 : : *
4 : : * Licensed under the Apache License, Version 2.0 (the "License");
5 : : * you may not use this file except in compliance with the License.
6 : : * You may obtain a copy of the License at
7 : : *
8 : : * http://www.apache.org/licenses/LICENSE-2.0
9 : : *
10 : : * Unless required by applicable law or agreed to in writing, software
11 : : * distributed under the License is distributed on an "AS IS" BASIS,
12 : : * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 : : * See the License for the specific language governing permissions and
14 : : * limitations under the License.
15 : : */
16 : :
17 : : #include "wirestead/framer/line_framer.hpp"
18 : :
19 : : #include <algorithm>
20 : : #include <cstring>
21 : : #include <iterator>
22 : : #include <limits>
23 : : #include <stdexcept>
24 : : #include <string_view>
25 : :
26 : : namespace wirestead {
27 : : namespace framer {
28 : : namespace {
29 : :
30 : 6 : void append_bytes(std::vector<uint8_t>& buffer, memory::ConstByteSpan data, size_t offset = 0) {
31 [ - + ]: 6 : if (offset >= data.size()) {
32 : 0 : return;
33 : : }
34 : :
35 : 6 : const size_t append_size = data.size() - offset;
36 : 6 : const size_t old_size = buffer.size();
37 : :
38 : : // Check for size_t overflow before resize
39 [ - + ]: 6 : if (append_size > std::numeric_limits<size_t>::max() - old_size) {
40 : 0 : throw std::length_error("append_bytes: size overflow");
41 : : }
42 : :
43 : 6 : buffer.resize(old_size + append_size);
44 : 6 : std::memcpy(buffer.data() + old_size, data.data() + offset, append_size);
45 : : }
46 : :
47 : : } // namespace
48 : :
49 : 57 : LineFramer::LineFramer(std::string_view delimiter, bool include_delimiter, size_t max_length)
50 : 171 : : delimiter_(delimiter), include_delimiter_(include_delimiter), max_length_(max_length) {
51 [ - + ]: 57 : if (delimiter_.empty()) {
52 : 0 : delimiter_ = "\n";
53 : : }
54 : 57 : }
55 : :
56 : 1032 : void LineFramer::push_bytes(memory::ConstByteSpan data) {
57 [ - + ]: 1032 : if (data.empty()) return;
58 : :
59 : : // Process data in chunks to prevent large memory allocations (DoS protection)
60 : : // We use max(max_length_, 4096) as a reasonable chunk size.
61 : : // This ensures that even if the user sends a huge payload, we only
62 : : // allocate memory incrementally and have a chance to clear the buffer
63 : : // if limits are exceeded.
64 : 1032 : const size_t chunk_limit = std::max(max_length_, size_t(4096));
65 : :
66 : 1032 : size_t offset = 0;
67 [ + + ]: 2064 : while (offset < data.size()) {
68 : 1032 : size_t len = std::min(data.size() - offset, chunk_limit);
69 : 1032 : push_bytes_internal(data.subspan(offset, len));
70 : 1032 : offset += len;
71 : : }
72 : : }
73 : :
74 : 2 : std::optional<size_t> LineFramer::skip_until_delimiter(memory::ConstByteSpan data) const {
75 : 2 : decltype(data.begin()) it;
76 [ + - ]: 2 : if (delimiter_.size() == 1) {
77 : 2 : const void* found = std::memchr(data.data(), static_cast<uint8_t>(delimiter_[0]), data.size());
78 [ + - ]: 4 : it = found ? data.begin() +
79 : 4 : std::distance(static_cast<const uint8_t*>(data.data()), static_cast<const uint8_t*>(found))
80 : 0 : : data.end();
81 : : } else {
82 : 0 : it = std::search(data.begin(), data.end(), delimiter_.begin(), delimiter_.end());
83 : : }
84 [ - + ]: 2 : if (it == data.end()) {
85 : 0 : return std::nullopt;
86 : : }
87 : 4 : return static_cast<size_t>(std::distance(data.begin(), it)) + delimiter_.size();
88 : : }
89 : :
90 : 1032 : void LineFramer::push_bytes_internal(memory::ConstByteSpan data) {
91 [ - + ]: 1032 : if (data.empty()) return;
92 : :
93 [ + + ]: 1032 : if (discarding_) {
94 : : // Resynchronizing after a discarded oversized message: don't buffer or
95 : : // emit anything from this chunk until the delimiter that ends the
96 : : // discarded message is found. A delimiter split across two discard-mode
97 : : // chunks is missed (we don't carry a partial match across calls here),
98 : : // which only delays resync by one more chunk - it doesn't misdeliver
99 : : // data.
100 : 2 : std::optional<size_t> consumed = skip_until_delimiter(data);
101 [ - + ]: 2 : if (!consumed.has_value()) {
102 : 2 : return; // Still haven't found it; stay in discarding mode.
103 : : }
104 : 2 : discarding_ = false;
105 : 2 : data = data.subspan(*consumed, data.size() - *consumed);
106 [ + - ]: 2 : if (data.empty()) return;
107 : : }
108 : :
109 : : // Fast Path: If buffer is empty, process data directly (zero-copy)
110 [ + + ]: 1030 : if (buffer_.empty()) {
111 : 1028 : size_t processed_count = scan_and_process(data, 0);
112 : :
113 : : // If we haven't processed everything, append the remainder to the buffer
114 [ + + ]: 1028 : if (processed_count < data.size()) {
115 : 4 : append_bytes(buffer_, data, processed_count);
116 : 4 : scanned_idx_ = buffer_.size(); // We scanned all of it
117 : :
118 : : // DoS protection for partial message overflow
119 [ + + ]: 4 : if (buffer_.size() > max_length_) {
120 : 2 : buffer_.clear();
121 : 2 : scanned_idx_ = 0;
122 : 2 : discarding_ = true;
123 : : }
124 : : } else {
125 : : // All processed, buffer remains empty
126 : 1024 : scanned_idx_ = 0;
127 : : }
128 : 1028 : return;
129 : : }
130 : :
131 : : // Slow Path: Append new data to buffer and process
132 : 2 : append_bytes(buffer_, data);
133 : :
134 : : // Determine where to start searching to avoid re-scanning
135 : : // We back up by delimiter length - 1 to catch split delimiters
136 : 2 : size_t search_start_idx = scanned_idx_;
137 [ + - ]: 2 : if (search_start_idx >= delimiter_.length()) {
138 : 2 : search_start_idx -= (delimiter_.length() - 1);
139 : : } else {
140 : 0 : search_start_idx = 0;
141 : : }
142 : :
143 : 2 : size_t processed_count = scan_and_process(memory::ConstByteSpan(buffer_), search_start_idx);
144 : :
145 : : // Batch erase all processed data to ensure O(N) erase complexity
146 [ + - ]: 2 : if (processed_count > 0) {
147 : 2 : buffer_.erase(buffer_.begin(), buffer_.begin() + static_cast<std::ptrdiff_t>(processed_count));
148 : : }
149 : :
150 : : // Update scanned_idx_ for the next call.
151 : : // The buffer size is now reduced. We have scanned everything that remains.
152 : 2 : scanned_idx_ = buffer_.size();
153 : :
154 : : // Final check: if the *remaining* partial message in the buffer already exceeds max_length_,
155 : : // we must reset to prevent unbound growth (DoS protection).
156 [ - + ]: 2 : if (buffer_.size() > max_length_) {
157 : 0 : buffer_.clear();
158 : 0 : scanned_idx_ = 0;
159 : 0 : discarding_ = true;
160 : : }
161 : : }
162 : :
163 : 1030 : size_t LineFramer::scan_and_process(memory::ConstByteSpan data, size_t search_start_offset) {
164 : : // Safety clamp
165 [ - + ]: 1030 : if (search_start_offset > data.size()) {
166 : 0 : search_start_offset = data.size();
167 : : }
168 : :
169 : : // processed_count tracks the number of bytes from the start of the buffer
170 : : // that have been either emitted as messages or skipped due to overflow.
171 : 1030 : size_t processed_count = 0;
172 : :
173 : : // Search cursor
174 : 1030 : size_t search_cursor = search_start_offset;
175 : :
176 : : // O(N) scan loop
177 : : while (true) {
178 [ - + ]: 2265 : if (search_cursor > data.size()) break;
179 : :
180 : : // Perform search using iterators derived from current buffer state
181 : 2265 : auto search_begin = data.begin() + static_cast<std::ptrdiff_t>(search_cursor);
182 : 2265 : decltype(data.begin()) it;
183 : :
184 [ + + ]: 2265 : if (delimiter_.size() == 1) {
185 : : // Optimization: Use std::memchr for single-byte delimiter
186 : 2261 : const void* found = std::memchr(data.data() + search_cursor, static_cast<uint8_t>(delimiter_[0]),
187 : 4522 : static_cast<size_t>(std::distance(search_begin, data.end())));
188 [ + + ]: 2261 : if (found) {
189 : 2466 : it = data.begin() + std::distance(static_cast<const uint8_t*>(data.data()), static_cast<const uint8_t*>(found));
190 : : } else {
191 : 1028 : it = data.end();
192 : : }
193 : : } else {
194 : 4 : it = std::search(search_begin, data.end(), delimiter_.begin(), delimiter_.end());
195 : : }
196 : :
197 [ + + ]: 2265 : if (it == data.end()) {
198 : 1030 : break;
199 : : }
200 : :
201 : : // Found a delimiter
202 : 1235 : size_t match_start_idx = static_cast<size_t>(std::distance(data.begin(), it));
203 : 1235 : size_t match_end_idx = match_start_idx + delimiter_.length();
204 : :
205 : : // Calculate message length (from end of previous processed data to end of current delimiter)
206 : 1235 : size_t current_msg_total_len = match_end_idx - processed_count;
207 : 1235 : size_t content_len = current_msg_total_len - delimiter_.length();
208 : :
209 : : // Check max_length against content length (delimiter excluded)
210 [ + + ]: 1235 : if (content_len > max_length_) {
211 : : // Message exceeds limit. Skip it.
212 : : } else {
213 : : // Valid message
214 [ + - ]: 1234 : if (on_message_) {
215 [ + + ]: 1234 : size_t payload_len = include_delimiter_ ? current_msg_total_len : content_len;
216 : 1234 : on_message_(memory::ConstByteSpan(data.data() + processed_count, payload_len));
217 : : }
218 : : }
219 : :
220 : : // Mark these bytes as processed
221 : 1235 : processed_count = match_end_idx;
222 : :
223 : : // Advance search cursor to start strictly after the current delimiter
224 : 1235 : search_cursor = processed_count;
225 : 1235 : }
226 : :
227 : 1030 : return processed_count;
228 : : }
229 : :
230 : 56 : void LineFramer::on_message(MessageCallback cb) { on_message_ = std::move(cb); }
231 : :
232 : 8 : void LineFramer::reset() {
233 : 8 : buffer_.clear();
234 : 8 : scanned_idx_ = 0;
235 : 8 : discarding_ = false;
236 : 8 : }
237 : :
238 : : } // namespace framer
239 : : } // namespace wirestead
|
