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 "frame_builder.h"
17
18namespace OHOS::ArkCompiler::Toolchain {
19ServerFrameBuilder& ServerFrameBuilder::SetFinal(bool fin)
20{
21    fin_ = fin;
22    return *this;
23}
24
25ServerFrameBuilder& ServerFrameBuilder::SetOpcode(FrameType opcode)
26{
27    opcode_ = opcode;
28    return *this;
29}
30
31ServerFrameBuilder& ServerFrameBuilder::SetPayload(const std::string& payload)
32{
33    payload_ = payload;
34    return *this;
35}
36
37ServerFrameBuilder& ServerFrameBuilder::SetPayload(std::string&& payload)
38{
39    payload_ = std::move(payload);
40    return *this;
41}
42
43ServerFrameBuilder& ServerFrameBuilder::AppendPayload(const std::string& payload)
44{
45    payload_.append(payload);
46    return *this;
47}
48
49std::string ServerFrameBuilder::Build() const
50{
51    std::string message;
52    PushFullHeader(message, 0);
53    PushPayload(message);
54    return message;
55}
56
57void 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
78void 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
92void 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
102void ServerFrameBuilder::PushPayload(std::string& message) const
103{
104    message.append(payload_);
105}
106
107ClientFrameBuilder::ClientFrameBuilder(bool final, FrameType opcode, const uint8_t maskingKey[WebSocketFrame::MASK_LEN])
108    : ServerFrameBuilder(final, opcode)
109{
110    SetMask(maskingKey);
111}
112
113ClientFrameBuilder& 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
121void 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
131void 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
140void 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