1 /*
2 * Copyright (c) 2022 Huawei Device Co., Ltd.
3 * Licensed under the Apache License, Version 2.0 (the "License");
4 * you may not use this file except in compliance with the License.
5 * You may obtain a copy of the License at
6 *
7 * http://www.apache.org/licenses/LICENSE-2.0
8 *
9 * Unless required by applicable law or agreed to in writing, software
10 * distributed under the License is distributed on an "AS IS" BASIS,
11 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 * See the License for the specific language governing permissions and
13 * limitations under the License.
14 */
15
16 #include "websocket/frame_builder.h"
17
18 namespace OHOS::ArkCompiler::Toolchain {
SetFinal(bool fin)19 ServerFrameBuilder& ServerFrameBuilder::SetFinal(bool fin)
20 {
21 fin_ = fin;
22 return *this;
23 }
24
SetOpcode(FrameType opcode)25 ServerFrameBuilder& ServerFrameBuilder::SetOpcode(FrameType opcode)
26 {
27 opcode_ = opcode;
28 return *this;
29 }
30
SetPayload(const std::string & payload)31 ServerFrameBuilder& ServerFrameBuilder::SetPayload(const std::string& payload)
32 {
33 payload_ = payload;
34 return *this;
35 }
36
SetPayload(std::string && payload)37 ServerFrameBuilder& ServerFrameBuilder::SetPayload(std::string&& payload)
38 {
39 payload_ = std::move(payload);
40 return *this;
41 }
42
AppendPayload(const std::string & payload)43 ServerFrameBuilder& ServerFrameBuilder::AppendPayload(const std::string& payload)
44 {
45 payload_.append(payload);
46 return *this;
47 }
48
Build() const49 std::string ServerFrameBuilder::Build() const
50 {
51 std::string message;
52 PushFullHeader(message, 0);
53 PushPayload(message);
54 return message;
55 }
56
PushFullHeader(std::string & message,size_t additionalReservedMem) const57 void ServerFrameBuilder::PushFullHeader(std::string& message, size_t additionalReservedMem) const
58 {
59 auto headerBytes = WebSocketFrame::HEADER_LEN;
60 auto payloadBytes = payload_.size();
61 uint8_t payloadLenField = 0;
62
63 if (payloadBytes <= WebSocketFrame::ONE_BYTE_LENTH_ENC_LIMIT) {
64 payloadLenField = static_cast<uint8_t>(payloadBytes);
65 } else if (payloadBytes < WebSocketFrame::TWO_BYTES_LENGTH_LIMIT) {
66 payloadLenField = WebSocketFrame::TWO_BYTES_LENTH_ENC;
67 headerBytes += WebSocketFrame::TWO_BYTES_LENTH;
68 } else {
69 payloadLenField = WebSocketFrame::EIGHT_BYTES_LENTH_ENC;
70 headerBytes += WebSocketFrame::EIGHT_BYTES_LENTH;
71 }
72
73 message.reserve(headerBytes + payloadBytes + additionalReservedMem);
74 PushHeader(message, payloadLenField);
75 PushPayloadLength(message, payloadLenField);
76 }
77
PushHeader(std::string & message,uint8_t payloadLenField) const78 void ServerFrameBuilder::PushHeader(std::string& message, uint8_t payloadLenField) const
79 {
80 uint8_t byte = EnumToNumber(opcode_);
81 if (fin_) {
82 byte |= 0x80;
83 }
84 message.push_back(byte);
85
86 // A server MUST NOT mask any frames that it sends to the client,
87 // hence mask bit must be set to zero (see https://www.rfc-editor.org/rfc/rfc6455#section-5.1)
88 byte = payloadLenField & 0x7f;
89 message.push_back(byte);
90 }
91
PushPayloadLength(std::string & message,uint8_t payloadLenField) const92 void ServerFrameBuilder::PushPayloadLength(std::string& message, uint8_t payloadLenField) const
93 {
94 uint64_t payloadLen = payload_.size();
95 if (payloadLenField == WebSocketFrame::TWO_BYTES_LENTH_ENC) {
96 PushNumberPerByte(message, static_cast<uint16_t>(payloadLen));
97 } else if (payloadLenField == WebSocketFrame::EIGHT_BYTES_LENTH_ENC) {
98 PushNumberPerByte(message, payloadLen);
99 }
100 }
101
PushPayload(std::string & message) const102 void ServerFrameBuilder::PushPayload(std::string& message) const
103 {
104 message.append(payload_);
105 }
106
ClientFrameBuilder(bool final,FrameType opcode,const uint8_t maskingKey[WebSocketFrame::MASK_LEN])107 ClientFrameBuilder::ClientFrameBuilder(bool final, FrameType opcode, const uint8_t maskingKey[WebSocketFrame::MASK_LEN])
108 : ServerFrameBuilder(final, opcode)
109 {
110 SetMask(maskingKey);
111 }
112
SetMask(const uint8_t maskingKey[WebSocketFrame::MASK_LEN])113 ClientFrameBuilder& ClientFrameBuilder::SetMask(const uint8_t maskingKey[WebSocketFrame::MASK_LEN])
114 {
115 for (size_t i = 0; i < WebSocketFrame::MASK_LEN; ++i) {
116 maskingKey_[i] = maskingKey[i];
117 }
118 return *this;
119 }
120
PushFullHeader(std::string & message,size_t additionalReservedMem) const121 void ClientFrameBuilder::PushFullHeader(std::string& message, size_t additionalReservedMem) const
122 {
123 // reserve additional 4 bytes for mask
124 ServerFrameBuilder::PushFullHeader(message, additionalReservedMem + WebSocketFrame::MASK_LEN);
125 // If the data is being sent by the client, the frame(s) MUST be masked
126 // (see https://www.rfc-editor.org/rfc/rfc6455#section-6.1)
127 message[1] |= 0x80;
128 PushMask(message);
129 }
130
PushPayload(std::string & message) const131 void ClientFrameBuilder::PushPayload(std::string& message) const
132 {
133 // push masked payload
134 for (size_t i = 0, end = payload_.size(); i < end; ++i) {
135 char c = payload_[i] ^ maskingKey_[i % WebSocketFrame::MASK_LEN];
136 message.push_back(c);
137 }
138 }
139
PushMask(std::string & message) const140 void ClientFrameBuilder::PushMask(std::string& message) const
141 {
142 for (size_t i = 0; i < WebSocketFrame::MASK_LEN; ++i) {
143 message.push_back(static_cast<char>(maskingKey_[i]));
144 }
145 }
146 } // OHOS::ArkCompiler::Toolchain
147