Line data Source code
1 : #include "body_reader.hpp"
2 : #include <algorithm>
3 : #include <cstring>
4 : #include <limits>
5 : #include <boost/asio/this_coro.hpp>
6 :
7 : namespace thinger::http {
8 :
9 : namespace {
10 :
11 : // Whether the coroutine reading was cancelled (e.g. the other side of a `||` ended
12 : // first): its read ended without taking anything, so the body is not incomplete,
13 : // just not read yet
14 319 : thinger::awaitable<bool> cancelled() {
15 : auto state = co_await boost::asio::this_coro::cancellation_state;
16 : co_return state.cancelled() != boost::asio::cancellation_type::none;
17 638 : }
18 :
19 : // The first piece of a Content-Length body read whole (see read_all())
20 : constexpr size_t body_first_piece = 64 * 1024;
21 :
22 : }
23 :
24 90361 : void body_reader::set_framing(bool chunked, std::uint64_t content_length) {
25 90361 : chunked_ = chunked;
26 90361 : remaining_ = chunked ? 0 : content_length;
27 90361 : read_ = 0;
28 90361 : decoder_ = {};
29 90361 : error_ = body_error::none;
30 90361 : }
31 :
32 4643 : void body_reader::set_read_ahead(const uint8_t* data, size_t size) {
33 4643 : if (data && size > 0) {
34 4637 : read_ahead_.assign(data, data + size);
35 4637 : read_ahead_offset_ = 0;
36 : }
37 4643 : }
38 :
39 4160441 : size_t body_reader::buffered() const {
40 4160441 : return read_ahead_.size() > read_ahead_offset_ ? read_ahead_.size() - read_ahead_offset_ : 0;
41 : }
42 :
43 17804 : std::vector<uint8_t> body_reader::take_read_ahead() {
44 35608 : std::vector<uint8_t> data(read_ahead_.begin() + static_cast<std::ptrdiff_t>(read_ahead_.size() - buffered()),
45 35608 : read_ahead_.end());
46 17804 : read_ahead_.clear();
47 17804 : read_ahead_offset_ = 0;
48 17804 : return data;
49 : }
50 :
51 122788 : bool body_reader::has_pending() const {
52 122788 : if (error_ != body_error::none) return false;
53 121285 : return chunked_ ? !decoder_.done() : remaining_ > 0;
54 : }
55 :
56 4134515 : size_t body_reader::received(const io_result& read) {
57 4134515 : const auto& [ec, bytes] = read;
58 4134515 : if (bytes == 0 && ec == read_deadline_expired) [[unlikely]] error_ = body_error::timeout;
59 4134515 : return bytes;
60 : }
61 :
62 4125976 : thinger::awaitable<size_t> body_reader::read_raw(uint8_t* buffer, size_t max_size) {
63 : if (size_t available = buffered()) {
64 : size_t bytes = std::min(available, max_size);
65 : std::memcpy(buffer, read_ahead_.data() + read_ahead_offset_, bytes);
66 : read_ahead_offset_ += bytes;
67 : co_return bytes;
68 : }
69 : if (!source_) co_return 0;
70 : co_return received(co_await source_(buffer, max_size));
71 8251952 : }
72 :
73 3589 : thinger::awaitable<bool> body_reader::fill_read_ahead() {
74 : read_ahead_.resize(framing_read_size);
75 : read_ahead_offset_ = 0;
76 : size_t bytes = source_ ? received(co_await source_(read_ahead_.data(), read_ahead_.size())) : 0;
77 : read_ahead_.resize(bytes);
78 : co_return bytes > 0;
79 7178 : }
80 :
81 4126412 : thinger::awaitable<size_t> body_reader::read_some_length(uint8_t* buffer, size_t max_size) {
82 : auto to_read = static_cast<size_t>(std::min<std::uint64_t>(max_size, remaining_));
83 : if (to_read == 0) co_return 0;
84 : size_t bytes = co_await read_raw(buffer, to_read);
85 : if (bytes == 0) {
86 : if (error_ == body_error::none && !co_await cancelled()) error_ = body_error::incomplete;
87 : co_return 0;
88 : }
89 : remaining_ -= bytes;
90 : co_return bytes;
91 8252824 : }
92 :
93 9266 : thinger::awaitable<size_t> body_reader::read_some_chunked(uint8_t* buffer, size_t max_size) {
94 : while (!decoder_.done()) {
95 : if (buffered() == 0) {
96 : // Within chunk data: read it straight into the caller buffer
97 : if (size_t pending = decoder_.data_pending()) {
98 : size_t bytes = source_ ? received(co_await source_(buffer, std::min(pending, max_size))) : 0;
99 : if (bytes == 0) {
100 : if (error_ == body_error::none && !co_await cancelled()) error_ = body_error::incomplete;
101 : co_return 0;
102 : }
103 : decoder_.data_consumed(bytes);
104 : co_return bytes;
105 : }
106 : // Framing: decode it from the read-ahead, where anything read past the end
107 : // of the body stays for the next request
108 : if (!co_await fill_read_ahead()) {
109 : if (error_ == body_error::none && !co_await cancelled()) error_ = body_error::incomplete;
110 : co_return 0;
111 : }
112 : }
113 :
114 : size_t produced = 0;
115 : read_ahead_offset_ += decoder_.decode(read_ahead_.data() + read_ahead_offset_, buffered(),
116 : buffer, max_size, produced);
117 : if (decoder_.failed()) {
118 : error_ = body_error::malformed;
119 : co_return 0;
120 : }
121 : if (produced > 0) co_return produced;
122 : }
123 : co_return 0;
124 18532 : }
125 :
126 4135714 : thinger::awaitable<size_t> body_reader::read_some(uint8_t* buffer, size_t max_size) {
127 : if (error_ != body_error::none || max_size == 0) co_return 0;
128 : // A Content-Length body over the limit fails before anything is read
129 : if (!chunked_ && remaining_ > 0 && (remaining_ > limit_ || read_ > limit_ - remaining_)) {
130 : error_ = body_error::too_large;
131 : co_return 0;
132 : }
133 : size_t bytes = chunked_ ? co_await read_some_chunked(buffer, max_size)
134 : : co_await read_some_length(buffer, max_size);
135 : read_ += bytes;
136 : if (read_ > limit_) {
137 : error_ = body_error::too_large;
138 : co_return 0;
139 : }
140 : co_return bytes;
141 8271428 : }
142 :
143 2317 : thinger::awaitable<size_t> body_reader::read(uint8_t* buffer, size_t size) {
144 : size_t total = 0;
145 : while (total < size) {
146 : size_t bytes = co_await read_some(buffer + total, size - total);
147 : if (bytes == 0) break;
148 : total += bytes;
149 : }
150 : co_return total;
151 4634 : }
152 :
153 2338 : thinger::awaitable<bool> body_reader::read_all(std::string& body, size_t max_size) {
154 : if (error_ != body_error::none) co_return false;
155 :
156 : if (!chunked_) {
157 : // Known size: reject it before allocating anything
158 : if (remaining_ > max_size || body.size() > max_size - remaining_) {
159 : error_ = body_error::too_large;
160 : co_return false;
161 : }
162 : // Grown as its bytes arrive, not to its announced size at once: a head announcing a
163 : // large body takes no memory for bytes that never come (as Go's io.ReadAll). The
164 : // first piece holds a small body whole; each next one doubles what arrived
165 : size_t offset = body.size();
166 : size_t end = offset + static_cast<size_t>(remaining_);
167 : while (body.size() < end) {
168 : size_t size = body.size();
169 : size_t piece = std::min(end - size, std::max(body_first_piece, size - offset));
170 : body.resize(size + piece);
171 : size_t bytes = co_await read(reinterpret_cast<uint8_t*>(body.data()) + size, piece);
172 : if (bytes < piece) {
173 : body.resize(size + bytes);
174 : co_return false;
175 : }
176 : }
177 : co_return true;
178 : }
179 :
180 : uint8_t buffer[8192];
181 : while (size_t bytes = co_await read_some(buffer, sizeof(buffer))) {
182 : if (bytes > max_size || body.size() > max_size - bytes) {
183 : error_ = body_error::too_large;
184 : co_return false;
185 : }
186 : body.append(reinterpret_cast<char*>(buffer), bytes);
187 : }
188 : co_return error_ == body_error::none;
189 4676 : }
190 :
191 197 : thinger::awaitable<bool> body_reader::discard(std::uint64_t max_size) {
192 : if (error_ != body_error::none) co_return false;
193 : if (!has_pending()) co_return true;
194 : if (!chunked_ && remaining_ > max_size) co_return false;
195 :
196 : uint8_t buffer[8192];
197 : std::uint64_t discarded = 0;
198 : while (has_pending()) {
199 : size_t bytes = co_await read_some(buffer, sizeof(buffer));
200 : if (bytes == 0) break;
201 : discarded += bytes;
202 : if (discarded > max_size) co_return false;
203 : }
204 : co_return error_ == body_error::none && !has_pending();
205 394 : }
206 :
207 : }
|
