Skip to content

Commit d534db8

Browse files
committed
fix and improve
1 parent 1348545 commit d534db8

File tree

8 files changed

+74
-74
lines changed

8 files changed

+74
-74
lines changed

include/ylt/standalone/cinatra/coro_http_client.hpp

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -436,7 +436,7 @@ class coro_http_client : public std::enable_shared_from_this<coro_http_client> {
436436
}
437437
else {
438438
#endif
439-
std::string encode_header = ws.encode_frame(source, op, true);
439+
auto encode_header = ws.encode_frame(source, op, true);
440440
std::vector<asio::const_buffer> buffers{
441441
asio::buffer(encode_header.data(), encode_header.size()),
442442
asio::buffer(source.data(), source.size())};
@@ -459,7 +459,7 @@ class coro_http_client : public std::enable_shared_from_this<coro_http_client> {
459459
if (cinatra::gzip_codec::deflate(
460460
{result.buf.data(), result.buf.size()}, dest_buf)) {
461461
std::span<char> msg(dest_buf.data(), dest_buf.size());
462-
std::string header = ws.encode_frame(msg, op, result.eof, true);
462+
auto header = ws.encode_frame(msg, op, result.eof, true);
463463
std::vector<asio::const_buffer> buffers{asio::buffer(header),
464464
asio::buffer(dest_buf)};
465465
auto [ec, sz] = co_await async_write(buffers);
@@ -478,7 +478,7 @@ class coro_http_client : public std::enable_shared_from_this<coro_http_client> {
478478
else {
479479
#endif
480480
std::span<char> msg(result.buf.data(), result.buf.size());
481-
std::string encode_header = ws.encode_frame(msg, op, result.eof);
481+
auto encode_header = ws.encode_frame(msg, op, result.eof);
482482
std::vector<asio::const_buffer> buffers{
483483
asio::buffer(encode_header.data(), encode_header.size()),
484484
asio::buffer(msg.data(), msg.size())};
@@ -1953,7 +1953,7 @@ class coro_http_client : public std::enable_shared_from_this<coro_http_client> {
19531953
auto close_str = ws.format_close_payload(close_code::normal,
19541954
reason.data(), reason.size());
19551955
auto span = std::span<char>(close_str);
1956-
std::string encode_header = ws.encode_frame(span, opcode::close, true);
1956+
auto encode_header = ws.encode_frame(span, opcode::close, true);
19571957
std::vector<asio::const_buffer> buffers{asio::buffer(encode_header),
19581958
asio::buffer(reason)};
19591959

include/ylt/standalone/cinatra/coro_http_connection.hpp

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -557,9 +557,9 @@ class coro_http_connection
557557
}
558558

559559
async_simple::coro::Lazy<std::error_code> write_websocket(
560-
std::string_view msg, opcode op = opcode::text) {
560+
std::string_view msg, opcode op = opcode::text, bool eof = true) {
561561
std::vector<asio::const_buffer> buffers;
562-
std::string header;
562+
std::string_view header;
563563
#ifdef CINATRA_ENABLE_GZIP
564564
std::string dest_buf;
565565
if (is_client_ws_compressed_ && msg.size() > 0) {
@@ -568,13 +568,13 @@ class coro_http_connection
568568
co_return std::make_error_code(std::errc::protocol_error);
569569
}
570570

571-
header = ws_.format_header(dest_buf.length(), op, true);
571+
header = ws_.encode_ws_header(dest_buf.length(), op, eof, true, false);
572572
buffers.push_back(asio::buffer(header));
573573
buffers.push_back(asio::buffer(dest_buf));
574574
}
575575
else {
576576
#endif
577-
header = ws_.format_header(msg.length(), op);
577+
header = ws_.encode_ws_header(msg.length(), op, eof, false, false);
578578
buffers.push_back(asio::buffer(header));
579579
buffers.push_back(asio::buffer(msg));
580580
#ifdef CINATRA_ENABLE_GZIP
@@ -666,7 +666,6 @@ class coro_http_connection
666666

667667
std::string close_msg = ws_.format_close_payload(
668668
close_code::normal, close_frame.message, close_frame.length);
669-
auto header = ws_.format_header(close_msg.length(), opcode::close);
670669

671670
co_await write_websocket(close_msg, opcode::close);
672671
close();

include/ylt/standalone/cinatra/coro_http_router.hpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -75,7 +75,7 @@ class coro_http_router {
7575

7676
if (whole_str.find(":") != std::string::npos) {
7777
std::string method_str(method_name);
78-
coro_router_tree_->coro_insert(key, std::move(http_handler),
78+
coro_router_tree_->coro_insert(whole_str, std::move(http_handler),
7979
method_str);
8080
}
8181
else {

include/ylt/standalone/cinatra/coro_radix_tree.hpp

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -300,12 +300,12 @@ class radix_tree {
300300
if (j < m) {
301301
std::shared_ptr<radix_tree_node> child(
302302
std::make_shared<radix_tree_node>(root->path.substr(j)));
303-
child->handler = root->handler;
303+
child->coro_handler = root->coro_handler;
304304
child->indices = root->indices;
305305
child->children = root->children;
306306

307307
root->path = root->path.substr(0, j);
308-
root->handler = {};
308+
root->coro_handler = {};
309309
root->indices = child->path[0];
310310
root->children = {child};
311311
}

include/ylt/standalone/cinatra/http_parser.hpp

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -101,6 +101,9 @@ class http_parser {
101101
}
102102

103103
full_url_ = url_;
104+
if (!queries_.empty()) {
105+
queries_.clear();
106+
}
104107
if (has_query) {
105108
size_t pos = url_.find('?');
106109
parse_query(url_.substr(pos + 1, url_len - pos - 1));

include/ylt/standalone/cinatra/picohttpparser.h

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -884,7 +884,7 @@ static const char *parse_headers(const char *buf, const char *buf_end,
884884
// has connection
885885
has_connection = true;
886886
char ch = *value;
887-
if (ch == 'U') {
887+
if (ch == 'U' || ch == 'u') {
888888
// has upgrade
889889
has_upgrade = true;
890890
}
Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
#pragma once
2+
3+
// Note: Update the version when release a new version.
4+
5+
// CINATRA_VERSION % 100 is the sub-minor version
6+
// CINATRA_VERSION / 100 % 1000 is the minor version
7+
// CINATRA_VERSION / 100000 is the major version
8+
#define CINATRA_VERSION 901 // 0.9.1

include/ylt/standalone/cinatra/websocket.hpp

Lines changed: 51 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -83,7 +83,7 @@ class websocket {
8383
}
8484

8585
if (msg_masked) {
86-
std::memcpy(mask_, inp + pos, 4);
86+
std::memcpy(mask_key_, inp + pos, 4);
8787
}
8888

8989
return left_header_len_ == 0 ? ws_header_status::complete
@@ -95,9 +95,9 @@ class websocket {
9595

9696
ws_frame_type parse_payload(std::span<char> buf) {
9797
// unmask data:
98-
if (*(uint32_t *)mask_ != 0) {
98+
if (*(uint32_t *)mask_key_ != 0) {
9999
for (size_t i = 0; i < payload_length_; i++) {
100-
buf[i] = buf[i] ^ mask_[i % 4];
100+
buf[i] = buf[i] ^ mask_key_[i % 4];
101101
}
102102
}
103103

@@ -121,16 +121,9 @@ class websocket {
121121
return ws_frame_type::WS_BINARY_FRAME;
122122
}
123123

124-
std::string format_header(size_t length, opcode code,
125-
bool is_compressed = false) {
126-
size_t header_length = encode_header(length, code, is_compressed);
127-
return {msg_header_, header_length};
128-
}
129-
130-
std::string encode_frame(std::span<char> &data, opcode op, bool eof,
131-
bool need_compression = false) {
132-
std::string header;
133-
/// Base header.
124+
std::string_view encode_ws_header(size_t size, opcode op, bool eof,
125+
bool need_compression = false,
126+
bool is_client = true) {
134127
frame_header hdr{};
135128
hdr.fin = eof;
136129
hdr.rsv1 = 0;
@@ -140,56 +133,55 @@ class websocket {
140133
hdr.rsv2 = 0;
141134
hdr.rsv3 = 0;
142135
hdr.opcode = static_cast<uint8_t>(op);
143-
hdr.mask = 1;
144-
145-
if (data.empty()) {
146-
int mask = 0;
147-
header.resize(sizeof(frame_header) + sizeof(mask));
148-
std::memcpy(header.data(), &hdr, sizeof(hdr));
149-
std::memcpy(header.data() + sizeof(hdr), &mask, sizeof(mask));
150-
return header;
151-
}
136+
hdr.mask = is_client;
137+
138+
hdr.len = size < 126 ? size : (size < 65536 ? 126 : 127);
152139

153-
hdr.len =
154-
data.size() < 126 ? data.size() : (data.size() < 65536 ? 126 : 127);
155-
156-
uint8_t buffer[sizeof(frame_header)];
157-
std::memcpy(buffer, (uint8_t *)&hdr, sizeof(hdr));
158-
std::string str_hdr_len =
159-
std::string((const char *)buffer, sizeof(frame_header));
160-
header.append(str_hdr_len);
161-
162-
/// The payload length may be larger than 126 bytes.
163-
std::string str_payload_len;
164-
if (data.size() >= 126) {
165-
if (data.size() >= 65536) {
166-
uint64_t len = data.size();
167-
str_payload_len.resize(sizeof(uint64_t));
168-
*((uint64_t *)&str_payload_len[0]) = htobe64(len);
140+
std::memcpy(msg_header_, (char *)&hdr, sizeof(hdr));
141+
142+
size_t len_bytes = 0;
143+
if (size >= 126) {
144+
if (size >= 65536) {
145+
len_bytes = 8;
146+
*((uint64_t *)(msg_header_ + 2)) = htobe64(size);
169147
}
170148
else {
171-
uint16_t len = data.size();
172-
str_payload_len.resize(sizeof(uint16_t));
173-
*((uint16_t *)&str_payload_len[0]) = htons(static_cast<uint16_t>(len));
149+
len_bytes = 2;
150+
*((uint16_t *)(msg_header_ + 2)) = htons(static_cast<uint16_t>(size));
174151
}
175-
header.append(str_payload_len);
176152
}
177153

178-
/// The mask is a 32-bit value.
179-
uint8_t mask[4] = {};
180-
header[1] |= 0x80;
181-
uint32_t random = (uint32_t)rand();
182-
memcpy(mask, &random, 4);
154+
size_t header_len = 6;
155+
156+
if (is_client) {
157+
if (size > 0) {
158+
// generate mask key.
159+
uint32_t random = (uint32_t)rand();
160+
memcpy(mask_key_, &random, 4);
161+
}
162+
163+
std::memcpy(msg_header_ + 2 + len_bytes, mask_key_, 4);
164+
}
165+
else {
166+
header_len = 2;
167+
}
183168

184-
size_t size = header.size();
185-
header.resize(size + 4);
186-
std::memcpy(header.data() + size, mask, 4);
169+
return {msg_header_, header_len + len_bytes};
170+
}
187171

172+
void encode_ws_payload(std::span<char> &data) {
188173
for (int i = 0; i < data.size(); ++i) {
189-
data[i] ^= mask[i % 4];
174+
data[i] ^= mask_key_[i % 4];
190175
}
176+
}
177+
178+
std::string_view encode_frame(std::span<char> &data, opcode op, bool eof,
179+
bool need_compression = false) {
180+
std::string_view ws_header =
181+
encode_ws_header(data.size(), op, eof, need_compression);
182+
encode_ws_payload(data);
191183

192-
return header;
184+
return ws_header;
193185
}
194186

195187
close_frame parse_close_payload(char *src, size_t length) {
@@ -209,14 +201,12 @@ class websocket {
209201

210202
std::string format_close_payload(uint16_t code, char *message,
211203
size_t length) {
212-
if (length == 0) {
213-
return "";
214-
}
215204
std::string close_payload;
216-
if (code) {
217-
close_payload.resize(length + 2);
218-
code = htons(code);
219-
std::memcpy(close_payload.data(), &code, 2);
205+
close_payload.resize(length + 2);
206+
code = htons(code);
207+
std::memcpy(close_payload.data(), &code, 2);
208+
209+
if (length > 0) {
220210
std::memcpy(close_payload.data() + 2, message, length);
221211
}
222212
return close_payload;
@@ -264,11 +254,11 @@ class websocket {
264254
size_t payload_length_ = 0;
265255

266256
size_t left_header_len_ = 0;
267-
uint8_t mask_[4] = {};
257+
uint8_t mask_key_[4] = {};
268258
unsigned char msg_opcode_ = 0;
269259
unsigned char msg_fin_ = 0;
270260

271-
char msg_header_[10];
261+
char msg_header_[14];
272262
ws_head_len len_bytes_ = SHORT_HEADER;
273263
};
274264

0 commit comments

Comments
 (0)