1/*
2 * Copyright (c) Huawei Technologies Co., Ltd. 2023. All rights reserved.
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#include "base_message.h"
16
17#include "securec.h"
18#include "varint_encode.h"
19
20namespace OHOS {
21namespace Developtools {
22namespace Profiler {
23namespace ProtoEncoder {
24bool BaseMessage::AllocateSubMessage()
25{
26    if (subMessageStack_ == nullptr) {
27        subMessage_ = new (std::nothrow) BaseMessage(writeCtx_);
28    } else {
29        subMessage_ = subMessageStack_->Get();
30    }
31
32    if (subMessage_ == nullptr) {
33        return false;
34    }
35    subMessage_->Reset(writeCtx_, subMessageStack_);
36    return true;
37}
38
39void BaseMessage::FinishSubMessage()
40{
41    int32_t subSize = subMessage_->Finish();
42    if (subMessageStack_ == nullptr) {
43        delete subMessage_;
44    } else {
45        subMessageStack_->Release();
46    }
47    subMessage_ = nullptr;
48    if (subSize < 0) {
49        Drop();
50        return;
51    }
52
53    uint8_t* fieldMemory = nullptr;
54    // backfill length
55    writeCtx_->seek(writeCtx_, backfillOffset_);
56    if (!writeCtx_->getMemory(writeCtx_, SIZE_RESERVED_LEN,
57                              &fieldMemory, &backfillOffset_)) {
58        Drop();
59        return;
60    }
61
62    if (subSize == 0) {
63        // reduce the size
64        *fieldMemory = 0;
65        writeCtx_->seek(writeCtx_, backfillOffset_ + 1);
66        size_++;
67        return;
68    }
69
70    // varint(Length)
71    EncodeVarintPadding(fieldMemory, subSize, SIZE_RESERVED_LEN);
72    size_ += SIZE_RESERVED_LEN + subSize;
73    // seek to tail
74    writeCtx_->seek(writeCtx_, backfillOffset_ + SIZE_RESERVED_LEN + subSize);
75}
76
77void BaseMessage::AddBytes(uint32_t fieldId, const void* data, uint32_t dataSize)
78{
79    if (subMessage_ != nullptr) {
80        FinishSubMessage();
81    }
82    if (!isWriting_) {
83        return;
84    }
85
86    uint8_t* fieldMemory = nullptr;
87    uint32_t fieldOffset = 0;
88    if (!writeCtx_->getMemory(writeCtx_, VARINT_ENCODE_MAX_SIZE + SIZE_RESERVED_LEN + dataSize,
89                              &fieldMemory, &fieldOffset)) {
90        Drop();
91        return;
92    }
93
94    uint32_t fieldSize = EncodeVarint(fieldMemory, (fieldId << FIELDID_SHIFT) | ProtoMessageType::LEN);
95    fieldSize += EncodeVarint(fieldMemory + fieldSize, dataSize);
96    if (dataSize != 0) {
97        if (memcpy_s(fieldMemory + fieldSize, dataSize, data, dataSize) != EOK) {
98            Drop();
99            return;
100        }
101    }
102    fieldSize += dataSize;
103    size_ += static_cast<int32_t>(fieldSize);
104    // seek to tail
105    writeCtx_->seek(writeCtx_, fieldOffset + fieldSize);
106}
107
108RandomWriteCtx* BaseMessage::StartAddBytes(uint32_t fieldId)
109{
110    if (subMessage_ != nullptr) {
111        FinishSubMessage();
112    }
113    if (!isWriting_) {
114        // finished or dropped
115        return nullptr;
116    }
117
118    uint8_t* fieldMemory = nullptr;
119    // max field size = varint(fieldId + type) + varint(len)
120    if (!writeCtx_->getMemory(writeCtx_, VARINT_ENCODE_MAX_SIZE + SIZE_RESERVED_LEN,
121                              &fieldMemory, &backfillOffset_)) {
122        Drop();
123        return nullptr;
124    }
125
126    uint32_t tagSize = EncodeVarint(fieldMemory, (fieldId << FIELDID_SHIFT) | ProtoMessageType::LEN);
127    backfillOffset_ += tagSize;
128    size_ += static_cast<int32_t>(tagSize);
129    // reserve length space
130    writeCtx_->seek(writeCtx_, backfillOffset_ + SIZE_RESERVED_LEN);
131    return writeCtx_;
132}
133
134void BaseMessage::FinishAddBytes(int32_t size)
135{
136    if (size < 0) {
137        Drop();
138        return;
139    }
140
141    uint8_t* fieldMemory = nullptr;
142    // backfill length
143    writeCtx_->seek(writeCtx_, backfillOffset_);
144    if (!writeCtx_->getMemory(writeCtx_, SIZE_RESERVED_LEN, &fieldMemory, &backfillOffset_)) {
145        Drop();
146        return;
147    }
148
149    if (size == 0) {
150        // reduce the size
151        *fieldMemory = 0;
152        writeCtx_->seek(writeCtx_, backfillOffset_ + 1);
153        size_++;
154        return;
155    }
156
157    // varint(Length)
158    EncodeVarintPadding(fieldMemory, size, SIZE_RESERVED_LEN);
159    size_ += SIZE_RESERVED_LEN + size;
160    // seek to tail
161    writeCtx_->seek(writeCtx_, backfillOffset_ + SIZE_RESERVED_LEN + size);
162}
163
164void BaseMessage::AddBytesByCallBack(uint32_t fieldId, GetDataCallback getData)
165{
166    if (!getData) {
167        return;
168    }
169
170    if (subMessage_ != nullptr) {
171        FinishSubMessage();
172    }
173    if (!isWriting_) {
174        return;
175    }
176
177    uint8_t* fieldMemory = nullptr;
178    uint32_t fieldOffset = 0;
179    // max field size = varint(fieldId + type) + varint(len)
180    if (!writeCtx_->getMemory(writeCtx_, VARINT_ENCODE_MAX_SIZE + SIZE_RESERVED_LEN,
181                              &fieldMemory, &fieldOffset)) {
182        Drop();
183        return;
184    }
185
186    uint32_t tagSize = EncodeVarint(fieldMemory, (fieldId << FIELDID_SHIFT) | ProtoMessageType::LEN);
187    fieldOffset += tagSize;
188    size_ += static_cast<int32_t>(tagSize);
189    // reserve length space
190    writeCtx_->seek(writeCtx_, fieldOffset + SIZE_RESERVED_LEN);
191
192    int32_t dataSize = getData(writeCtx_);
193    if (dataSize < 0) {
194        Drop();
195        return;
196    }
197
198    // backfill length
199    writeCtx_->seek(writeCtx_, fieldOffset);
200    if (!writeCtx_->getMemory(writeCtx_, SIZE_RESERVED_LEN, &fieldMemory, &fieldOffset)) {
201        Drop();
202        return;
203    }
204
205    if (dataSize == 0) {
206        // reduce the size
207        *fieldMemory = 0;
208        writeCtx_->seek(writeCtx_, fieldOffset + 1);
209        size_++;
210        return;
211    }
212
213    EncodeVarintPadding(fieldMemory, dataSize, SIZE_RESERVED_LEN);
214    size_ += SIZE_RESERVED_LEN + dataSize;
215    // seek to tail
216    writeCtx_->seek(writeCtx_, fieldOffset + SIZE_RESERVED_LEN + dataSize);
217}
218} // namespace ProtoEncoder
219} // namespace Profiler
220} // namespace Developtools
221} // namespace OHOS
222