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 : : #pragma once
18 : :
19 : : #include <algorithm>
20 : : #include <functional>
21 : : #include <limits>
22 : : #include <optional>
23 : : #include <typeinfo>
24 : :
25 : : #include "wirestead/framer/length_prefix_framer.hpp"
26 : : #include "wirestead/framer/line_framer.hpp"
27 : : #include "wirestead/framer/packet_framer.hpp"
28 : : #include "wirestead/wrapper/context.hpp"
29 : : #include "wirestead/wrapper/receive_budget.hpp"
30 : :
31 : : namespace wirestead::wrapper::detail {
32 : : struct ReceiveOverflow {
33 : : ReceiveOverflowReason reason;
34 : : };
35 : 3740 : inline std::shared_ptr<ReceiveCharge> reserve_receive(const std::shared_ptr<ReceiveScope>& scope, size_t bytes) {
36 : 3740 : auto charge = scope->reserve(bytes);
37 [ + + ]: 3740 : if (!charge) throw ReceiveOverflow{ReceiveOverflowReason::ByteLimit};
38 : 3722 : return charge;
39 : 18 : }
40 : :
41 : : // Reservation precedes payload allocation and outlives the payload, including
42 : : // a batch moved out of the queue into a user callback.
43 : : struct ReceiveMessage {
44 : : std::shared_ptr<ReceiveCharge> charge;
45 : : MessageContext context;
46 : 1348 : ReceiveMessage(std::shared_ptr<ReceiveCharge> c, ClientId id, memory::ConstByteSpan bytes)
47 : 4044 : : charge(std::move(c)), context(id, memory::SafeDataBuffer(bytes)) {}
48 : 3231 : ReceiveMessage(ReceiveMessage&&) noexcept = default;
49 : : ReceiveMessage& operator=(ReceiveMessage&&) noexcept = default;
50 : : };
51 : : inline thread_local ReceiveMessage* prepared_message = nullptr;
52 : : class PreparedMessageGuard {
53 : : public:
54 : 1235 : explicit PreparedMessageGuard(ReceiveMessage& message) : previous_(prepared_message) { prepared_message = &message; }
55 : 1235 : ~PreparedMessageGuard() { prepared_message = previous_; }
56 : :
57 : : private:
58 : : ReceiveMessage* previous_;
59 : : };
60 : : // Consume staged ownership at the wrapper message boundary, before any user
61 : : // callback can reenter another receive path. Raw input retention never adopts
62 : : // another dispatch's storage, even when it happens to borrow the same address.
63 : 1227 : inline std::optional<ReceiveMessage> take_prepared_message(ClientId id, memory::ConstByteSpan bytes) {
64 [ + + + - : 2453 : if (prepared_message && prepared_message->charge && prepared_message->context.client_id() == id &&
+ - ]
65 [ + - + - : 3680 : prepared_message->context.safe_data().data() == bytes.data() &&
+ + ]
66 : 1226 : prepared_message->context.safe_data().size() == bytes.size()) {
67 : 1226 : return std::move(*prepared_message);
68 : : }
69 : 1 : return std::nullopt;
70 : : }
71 : 1366 : inline ReceiveMessage retain_received(const std::shared_ptr<ReceiveScope>& scope, ClientId id,
72 : : memory::ConstByteSpan bytes) {
73 : : // Conservative record allowance also covers staging and vector growth.
74 : 1366 : constexpr size_t records = 4 * (sizeof(MessageContext) + sizeof(std::shared_ptr<ReceiveCharge>));
75 [ - + ]: 1366 : if (bytes.size() > std::numeric_limits<size_t>::max() - records)
76 : 0 : throw ReceiveOverflow{ReceiveOverflowReason::ByteLimit};
77 : 1366 : return ReceiveMessage(reserve_receive(scope, bytes.size() + records), id, bytes);
78 : : }
79 : :
80 : : class ReceiveBatch {
81 : : public:
82 : 3812 : ReceiveBatch() = default;
83 : 749 : ReceiveBatch(ReceiveBatch&&) noexcept = default;
84 : 51 : ReceiveBatch& operator=(ReceiveBatch&& other) noexcept {
85 [ + - ]: 51 : if (this != &other) {
86 : 51 : clear();
87 : 51 : charges_ = std::move(other.charges_);
88 : 51 : messages_ = std::move(other.messages_);
89 : : }
90 : 51 : return *this;
91 : : }
92 : 1111 : bool empty() const { return messages_.empty(); }
93 : 470 : size_t size() const { return messages_.size(); }
94 : 10029 : void clear() {
95 : 10029 : std::vector<MessageContext>().swap(messages_);
96 : 10029 : std::vector<std::shared_ptr<ReceiveCharge>>().swap(charges_);
97 : 10029 : }
98 : 263 : void emplace_back(ReceiveMessage message) {
99 : 263 : charges_.push_back(message.charge);
100 : : try {
101 : 263 : messages_.push_back(std::move(message.context));
102 : 0 : } catch (...) {
103 : 0 : charges_.pop_back();
104 : 0 : throw;
105 : 0 : }
106 : 263 : }
107 : 110 : operator const std::vector<MessageContext>&() const { return messages_; }
108 : :
109 : : private:
110 : : std::vector<std::shared_ptr<ReceiveCharge>> charges_;
111 : : std::vector<MessageContext> messages_;
112 : : };
113 : :
114 : : // Friend access changes neither existing public virtual tables nor layouts.
115 : : // Exact type checks intentionally exclude custom subclasses with hidden state.
116 : : struct ReceiveFramerAccess {
117 : : template <class F>
118 : 9886 : static F* exact(framer::IFramer* f) {
119 [ + - + - : 9886 : return f && typeid(*f) == typeid(F) ? static_cast<F*>(f) : nullptr;
+ + ]
120 : : }
121 : : template <class Fn>
122 : 9808 : static bool visit(framer::IFramer* f, Fn&& fn) {
123 [ + + ]: 9808 : if (auto* p = exact<framer::LineFramer>(f)) {
124 : 9736 : fn(*p);
125 : 9736 : return true;
126 : : }
127 [ + + ]: 72 : if (auto* p = exact<framer::LengthPrefixFramer>(f)) {
128 : 66 : fn(*p);
129 : 66 : return true;
130 : : }
131 [ + - ]: 6 : if (auto* p = exact<framer::PacketFramer>(f)) {
132 : 6 : fn(*p);
133 : 6 : return true;
134 : : }
135 : 0 : return false;
136 : : }
137 : 2375 : static size_t capacity(framer::IFramer* f) {
138 : 2375 : size_t n = 0;
139 : 4750 : visit(f, [&](auto& p) { n = p.buffer_.capacity(); });
140 : 2375 : return n;
141 : : }
142 : 1206 : static size_t size(framer::IFramer* f) {
143 : 1206 : size_t n = 0;
144 : 2412 : visit(f, [&](auto& p) { n = p.buffer_.size(); });
145 : 1206 : return n;
146 : : }
147 : 1179 : static bool supported(framer::IFramer* f) {
148 : 2358 : return visit(f, [](auto&) {});
149 : : }
150 : 1179 : static std::unique_ptr<framer::IFramer> clone(framer::IFramer* f) {
151 : 1179 : std::unique_ptr<framer::IFramer> result;
152 : 2358 : visit(f, [&](auto& p) {
153 : : using F = std::decay_t<decltype(p)>;
154 : 1179 : result = std::make_unique<F>(p);
155 : 1179 : });
156 : 1179 : return result;
157 : 0 : }
158 : 1195 : static void reserve_exact(framer::IFramer* f, size_t n) {
159 : 2390 : visit(f, [&](auto& p) {
160 [ - + ]: 1195 : if (n <= p.buffer_.capacity()) return;
161 : 1195 : std::vector<uint8_t> replacement;
162 : 1195 : replacement.reserve(n);
163 : 1195 : replacement.insert(replacement.end(), p.buffer_.begin(), p.buffer_.end());
164 : 1195 : p.buffer_.swap(replacement);
165 : 1195 : });
166 : 1195 : }
167 : 1170 : static void commit(framer::IFramer* target, framer::IFramer* source) {
168 : 2340 : visit(target, [&](auto& p) {
169 : : using F = std::decay_t<decltype(p)>;
170 : 1170 : auto& q = *static_cast<F*>(source);
171 : 1170 : p.buffer_.swap(q.buffer_);
172 : : if constexpr (std::is_same_v<F, framer::LineFramer>) {
173 : 1168 : std::swap(p.scanned_idx_, q.scanned_idx_);
174 : 1168 : std::swap(p.discarding_, q.discarding_);
175 : : } else if constexpr (std::is_same_v<F, framer::PacketFramer>) {
176 : 0 : std::swap(p.scanned_idx_, q.scanned_idx_);
177 : 0 : std::swap(p.state_, q.state_);
178 : : }
179 : 1170 : });
180 : 1170 : }
181 : 1168 : static framer::IFramer::MessageCallback callback(framer::IFramer* f) {
182 : 1168 : framer::IFramer::MessageCallback result;
183 : 2336 : visit(f, [&](auto& p) { result = p.on_message_; });
184 : 1168 : return result;
185 : 0 : }
186 : 4640 : static void release_storage(framer::IFramer* f) {
187 [ + + ]: 4640 : if (!f) return;
188 : 336 : f->reset();
189 : 672 : visit(f, [](auto& p) { std::vector<uint8_t>().swap(p.buffer_); });
190 : : }
191 : : };
192 : : struct ReceiveState {
193 : 2032 : explicit ReceiveState(std::shared_ptr<ReceiveScope> s) : scope(std::move(s)) {}
194 : : std::shared_ptr<ReceiveScope> scope;
195 : : std::shared_ptr<ReceiveCharge> frame_storage;
196 : 4640 : void reset(framer::IFramer* f) {
197 : 4640 : ReceiveFramerAccess::release_storage(f);
198 : 4640 : frame_storage.reset();
199 : 4640 : }
200 : : };
201 : :
202 : : struct PreparedReceive {
203 : : std::optional<ReceiveMessage> raw;
204 : : // MSVC debug vectors allocate an iterator proxy even when empty. Construct
205 : : // staging only for a built-in framer so raw scalar receive stays allocation-free.
206 : : std::optional<std::vector<ReceiveMessage>> messages;
207 : : std::shared_ptr<framer::IFramer> framer;
208 : : memory::ConstByteSpan input;
209 : : bool builtin = false;
210 : 3346 : void deliver() {
211 [ + + ]: 3370 : if (!framer) return;
212 [ - + ]: 1168 : if (!builtin) {
213 : 0 : framer->push_bytes(input);
214 : 0 : return;
215 : : }
216 : 1168 : auto callback = ReceiveFramerAccess::callback(framer.get());
217 [ + + ]: 1168 : if (!callback) return;
218 [ + + ]: 2375 : for (auto& message : *messages) {
219 : 1231 : auto bytes = message.context.safe_data().as_span();
220 : 1231 : PreparedMessageGuard guard(message);
221 : 1231 : callback(bytes);
222 : 1231 : }
223 : 1168 : }
224 : : };
225 : :
226 : : // Stage a whole input before exposing callbacks. A rejected UDP datagram cannot
227 : : // partially alter an existing built-in framer or leak some decoded messages.
228 : 3384 : inline PreparedReceive prepare_receive(ReceiveState& state, const std::shared_ptr<framer::IFramer>& f, ClientId id,
229 : : memory::ConstByteSpan input, bool raw_batch) {
230 : 3384 : PreparedReceive result;
231 : 3384 : result.framer = f;
232 : 3384 : result.input = input;
233 [ + + + + ]: 3384 : if (raw_batch) result.raw.emplace(retain_received(state.scope, id, input));
234 [ + + + - : 3367 : if (!f || !ReceiveFramerAccess::supported(f.get())) return result;
- + + + ]
235 : 1179 : result.builtin = true;
236 : 1179 : result.messages.emplace();
237 : 1179 : const auto limit = state.scope->limits().max_frame_bytes;
238 : 1179 : const auto capacity = ReceiveFramerAccess::capacity(f.get());
239 [ - + ]: 1179 : if (capacity > limit) throw ReceiveOverflow{ReceiveOverflowReason::FrameLimit};
240 [ + + - + : 1179 : if (!state.frame_storage && capacity) state.frame_storage = reserve_receive(state.scope, capacity);
- + - - ]
241 : 1179 : auto work_charge = reserve_receive(state.scope, capacity);
242 : 1179 : auto work = ReceiveFramerAccess::clone(f.get());
243 : 2358 : work->on_message(
244 : 2426 : [&](memory::ConstByteSpan message) { result.messages->emplace_back(retain_received(state.scope, id, message)); });
245 : 1179 : size_t offset = 0;
246 [ + + ]: 2373 : while (offset < input.size()) {
247 : 1203 : const auto used = ReceiveFramerAccess::size(work.get());
248 [ + + ]: 1203 : if (used >= limit) throw ReceiveOverflow{ReceiveOverflowReason::FrameLimit};
249 : 1195 : const auto n = std::min({input.size() - offset, limit - used, size_t(4096)});
250 : 1195 : const auto needed = used + n;
251 [ + - + - ]: 1195 : if (needed > ReceiveFramerAccess::capacity(work.get())) {
252 : : // Keep the old allocation charged until reserve_exact has destroyed it.
253 : 1195 : auto next = reserve_receive(state.scope, needed);
254 : 1195 : ReceiveFramerAccess::reserve_exact(work.get(), needed);
255 : 1195 : work_charge = std::move(next);
256 : 1195 : }
257 : 1195 : work->push_bytes(input.subspan(offset, n));
258 : 1194 : offset += n;
259 : : }
260 : 1170 : ReceiveFramerAccess::commit(f.get(), work.get());
261 : 1170 : work.reset(); // destroy old source storage before releasing its reservation
262 : 1170 : state.frame_storage = std::move(work_charge);
263 : 1170 : return result;
264 : 1214 : }
265 : : } // namespace wirestead::wrapper::detail
|
