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/packet_framer.hpp"
18 : :
19 : : #include <algorithm>
20 : : #include <cstring>
21 : : #include <iterator>
22 : : #include <limits>
23 : : #include <stdexcept>
24 : :
25 : : namespace wirestead {
26 : : namespace framer {
27 : : namespace {
28 : :
29 : 19 : void append_bytes(std::vector<uint8_t>& buffer, memory::ConstByteSpan data, size_t offset = 0) {
30 [ - + ]: 19 : if (offset >= data.size()) {
31 : 0 : return;
32 : : }
33 : :
34 : 19 : const size_t append_size = data.size() - offset;
35 : 19 : const size_t old_size = buffer.size();
36 : :
37 : : // Check for size_t overflow before resize
38 [ - + ]: 19 : if (append_size > std::numeric_limits<size_t>::max() - old_size) {
39 : 0 : throw std::length_error("append_bytes: size overflow");
40 : : }
41 : :
42 : 19 : buffer.resize(old_size + append_size);
43 : 19 : std::memcpy(buffer.data() + old_size, data.data() + offset, append_size);
44 : : }
45 : :
46 : : } // namespace
47 : :
48 : 32 : PacketFramer::PacketFramer(const std::vector<uint8_t>& start_pattern, const std::vector<uint8_t>& end_pattern,
49 : 32 : size_t max_length)
50 : 32 : : start_pattern_(start_pattern), end_pattern_(end_pattern), max_length_(max_length), state_(State::Sync) {
51 [ + + + + : 32 : if (start_pattern_.empty() && end_pattern_.empty()) {
+ + ]
52 : 1 : throw std::invalid_argument("PacketFramer: start_pattern and end_pattern cannot both be empty.");
53 : : }
54 : 36 : }
55 : :
56 : 30 : void PacketFramer::push_bytes(memory::ConstByteSpan data) {
57 [ + + ]: 30 : if (data.empty()) return;
58 : :
59 : : // Fast Path: Zero-copy processing if buffer is empty
60 [ + + ]: 29 : if (buffer_.empty()) {
61 : 21 : size_t processed_count = 0;
62 [ + + ]: 32 : while (processed_count < data.size()) {
63 : : // Find start pattern
64 : 24 : auto search_start = data.begin() + static_cast<std::ptrdiff_t>(processed_count);
65 : 24 : auto it_start = std::search(search_start, data.end(), start_pattern_.begin(), start_pattern_.end());
66 : :
67 [ + + ]: 24 : if (it_start == data.end()) {
68 : : // Start pattern not found.
69 : : // Keep partial match at the end if applicable.
70 : 7 : size_t remaining = data.size() - processed_count;
71 [ + + ]: 7 : size_t keep_len = (start_pattern_.size() > 1) ? (start_pattern_.size() - 1) : 0;
72 [ + + ]: 7 : if (remaining > keep_len) {
73 : 6 : processed_count = data.size() - keep_len;
74 : : }
75 : 13 : break;
76 : : }
77 : :
78 : : // Found start pattern
79 : 17 : size_t start_idx = static_cast<size_t>(std::distance(data.begin(), it_start));
80 : 17 : size_t search_end_start_idx = start_idx + start_pattern_.size();
81 : :
82 [ + + ]: 17 : if (end_pattern_.empty()) {
83 : : // Assume minimal packet is start pattern only
84 : 3 : size_t packet_len = start_pattern_.size();
85 [ + - ]: 3 : if (on_message_) {
86 : 3 : on_message_(data.subspan(start_idx, packet_len));
87 : : }
88 : 3 : processed_count = start_idx + packet_len;
89 : 3 : continue;
90 : 3 : }
91 : :
92 : : // Find end pattern
93 : 14 : auto search_end_start = data.begin() + static_cast<std::ptrdiff_t>(search_end_start_idx);
94 : 14 : auto it_end = std::search(search_end_start, data.end(), end_pattern_.begin(), end_pattern_.end());
95 : :
96 [ + + ]: 14 : if (it_end == data.end()) {
97 : : // Found start pattern but not end pattern.
98 : : // Buffer everything starting from start_idx.
99 : 6 : processed_count = start_idx;
100 : : // The remaining data starts with start_pattern, so we will transition to Collect state
101 : : // after appending to buffer.
102 : 6 : break;
103 : : }
104 : :
105 : : // Found end pattern
106 : 16 : size_t packet_len = static_cast<size_t>(std::distance(data.begin(), it_end)) + end_pattern_.size() - start_idx;
107 : :
108 [ + + ]: 8 : if (packet_len <= max_length_) {
109 [ + - ]: 7 : if (on_message_) {
110 : 7 : on_message_(data.subspan(start_idx, packet_len));
111 : : }
112 : : }
113 : : // If > max_length, discard by advancing processed_count past it
114 : :
115 : 8 : processed_count = start_idx + packet_len;
116 : : }
117 : :
118 [ + + ]: 21 : if (processed_count < data.size()) {
119 : 11 : append_bytes(buffer_, data, processed_count);
120 : :
121 : : // Update state if we buffered a partial packet starting with start_pattern
122 [ + - + - : 11 : if (state_ == State::Sync && !buffer_.empty()) {
+ - ]
123 [ + + ]: 11 : if (buffer_.size() >= start_pattern_.size()) {
124 [ + - ]: 6 : if (std::equal(start_pattern_.begin(), start_pattern_.end(), buffer_.begin())) {
125 : 6 : state_ = State::Collect;
126 : 6 : scanned_idx_ = start_pattern_.size();
127 : : }
128 : : }
129 : : }
130 : :
131 : : // The end pattern wasn't found within this call. Apply the same
132 : : // max_length_ cap the Collect-state loop below applies on later calls -
133 : : // otherwise a single push_bytes() call with a start marker and a large
134 : : // unterminated payload could grow buffer_ past the documented limit
135 : : // before any call gets a chance to check it.
136 [ + + + + : 11 : if (state_ == State::Collect && buffer_.size() > max_length_) {
+ + ]
137 : 2 : buffer_.clear();
138 : 2 : state_ = State::Sync;
139 : 2 : scanned_idx_ = 0;
140 : : }
141 : : }
142 : 21 : return;
143 : : }
144 : :
145 : 8 : append_bytes(buffer_, data);
146 : :
147 : : while (true) {
148 [ + + ]: 22 : if (state_ == State::Sync) {
149 [ + + ]: 13 : if (start_pattern_.empty()) {
150 : 1 : state_ = State::Collect;
151 : 1 : continue;
152 : : }
153 : :
154 : 12 : auto it = std::search(buffer_.begin(), buffer_.end(), start_pattern_.begin(), start_pattern_.end());
155 [ + + ]: 12 : if (it != buffer_.end()) {
156 : : // Found start pattern.
157 : : // Discard everything before start pattern.
158 [ + + ]: 5 : if (it != buffer_.begin()) {
159 : 1 : buffer_.erase(buffer_.begin(), it);
160 : : }
161 : 5 : state_ = State::Collect;
162 : : // Start scanning for end pattern after the start pattern we just found
163 : 5 : scanned_idx_ = start_pattern_.size();
164 : : // Continue to check for end pattern immediately
165 : : } else {
166 : : // Start pattern not found.
167 : : // Keep partial match at the end.
168 [ + + ]: 7 : if (start_pattern_.size() > 1) {
169 : 6 : size_t keep_len = start_pattern_.size() - 1;
170 [ + + ]: 6 : if (buffer_.size() > keep_len) {
171 : 1 : buffer_.erase(buffer_.begin(), buffer_.end() - static_cast<std::ptrdiff_t>(keep_len));
172 : : }
173 : : } else {
174 : 1 : buffer_.clear();
175 : : }
176 : 7 : break; // Need more data
177 : : }
178 [ + - ]: 9 : } else if (state_ == State::Collect) {
179 [ + + ]: 9 : if (end_pattern_.empty()) {
180 : : // If end pattern is empty, packet ends immediately after start pattern?
181 : : // Assume minimal packet is start pattern only
182 : 1 : size_t packet_len = start_pattern_.size();
183 [ + - ]: 1 : if (on_message_) {
184 : 1 : on_message_(memory::ConstByteSpan(buffer_.data(), packet_len));
185 : : }
186 [ - + ]: 1 : if (buffer_.empty()) return;
187 : :
188 : 1 : buffer_.erase(buffer_.begin(), buffer_.begin() + static_cast<std::ptrdiff_t>(packet_len));
189 : 1 : state_ = State::Sync;
190 : 1 : continue;
191 : 1 : }
192 : :
193 : : // Search for end pattern *after* start pattern
194 : : // Optimization: use scanned_idx_ to avoid re-scanning
195 : 8 : size_t search_offset = std::max(start_pattern_.size(), scanned_idx_);
196 : :
197 : : // Back up slightly to catch split end pattern if we are resuming search
198 [ - + ]: 8 : if (search_offset > start_pattern_.size()) {
199 [ # # ]: 0 : size_t overlap = (end_pattern_.size() > 1) ? (end_pattern_.size() - 1) : 0;
200 [ # # ]: 0 : if (search_offset >= overlap) {
201 : 0 : search_offset -= overlap;
202 : : } else {
203 : 0 : search_offset = 0;
204 : : }
205 : : }
206 : :
207 : : // Safety clamp to ensure we don't search inside start pattern
208 [ - + ]: 8 : if (search_offset < start_pattern_.size()) {
209 : 0 : search_offset = start_pattern_.size();
210 : : }
211 : :
212 [ - + ]: 8 : if (buffer_.size() < search_offset) {
213 : : // Should not happen if Sync worked correctly
214 : 1 : break;
215 : : }
216 : :
217 : 8 : auto search_start = buffer_.begin() + static_cast<std::ptrdiff_t>(search_offset);
218 : 8 : auto it = std::search(search_start, buffer_.end(), end_pattern_.begin(), end_pattern_.end());
219 : :
220 [ + + ]: 8 : if (it != buffer_.end()) {
221 : : // Found end pattern.
222 : 14 : size_t packet_len = static_cast<size_t>(std::distance(buffer_.begin(), it)) + end_pattern_.size();
223 : :
224 [ + + ]: 7 : if (packet_len <= max_length_) {
225 [ + - ]: 6 : if (on_message_) {
226 : 6 : on_message_(memory::ConstByteSpan(buffer_.data(), packet_len));
227 : : }
228 [ - + ]: 6 : if (buffer_.empty()) return;
229 : :
230 : 6 : buffer_.erase(buffer_.begin(), buffer_.begin() + static_cast<std::ptrdiff_t>(packet_len));
231 : 6 : state_ = State::Sync;
232 : 6 : scanned_idx_ = 0;
233 : : } else {
234 : : // Exceeded max length, discard packet
235 : 1 : buffer_.erase(buffer_.begin(), buffer_.begin() + static_cast<std::ptrdiff_t>(packet_len));
236 : 1 : state_ = State::Sync;
237 : 1 : scanned_idx_ = 0;
238 : : }
239 : : } else {
240 : : // End pattern not found.
241 : 1 : scanned_idx_ = buffer_.size();
242 [ - + ]: 1 : if (buffer_.size() > max_length_) {
243 : : // Exceeded limit while collecting. Reset.
244 : 0 : buffer_.clear();
245 : 0 : state_ = State::Sync;
246 : 0 : scanned_idx_ = 0;
247 : : }
248 : 1 : break; // Need more data
249 : : }
250 : : }
251 : 14 : }
252 : : }
253 : :
254 : 36 : void PacketFramer::on_message(MessageCallback cb) { on_message_ = std::move(cb); }
255 : :
256 : 2 : void PacketFramer::reset() {
257 : 2 : buffer_.clear();
258 : 2 : state_ = State::Sync;
259 : 2 : scanned_idx_ = 0;
260 : 2 : }
261 : :
262 : : } // namespace framer
263 : : } // namespace wirestead
|
