Alexei Frolov | f9ae189 | 2021-04-01 18:24:27 -0700 | [diff] [blame] | 1 | // Copyright 2021 The Pigweed Authors |
Alexei Frolov | bbf164c | 2019-12-16 12:51:59 -0800 | [diff] [blame] | 2 | // |
| 3 | // Licensed under the Apache License, Version 2.0 (the "License"); you may not |
| 4 | // use this file except in compliance with the License. You may obtain a copy of |
| 5 | // the License at |
| 6 | // |
| 7 | // https://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, WITHOUT |
| 11 | // WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the |
| 12 | // License for the specific language governing permissions and limitations under |
| 13 | // the License. |
| 14 | |
| 15 | #include "pw_protobuf/encoder.h" |
| 16 | |
Alexei Frolov | f9ae189 | 2021-04-01 18:24:27 -0700 | [diff] [blame] | 17 | #include <limits> |
| 18 | |
Alexei Frolov | bbf164c | 2019-12-16 12:51:59 -0800 | [diff] [blame] | 19 | namespace pw::protobuf { |
| 20 | |
| 21 | Status Encoder::WriteUint64(uint32_t field_number, uint64_t value) { |
Alexei Frolov | 311c620 | 2020-01-06 11:16:20 -0800 | [diff] [blame] | 22 | std::byte* original_cursor = cursor_; |
Alexei Frolov | bbf164c | 2019-12-16 12:51:59 -0800 | [diff] [blame] | 23 | WriteFieldKey(field_number, WireType::kVarint); |
Alexei Frolov | f9ae189 | 2021-04-01 18:24:27 -0700 | [diff] [blame] | 24 | WriteVarint(value); |
| 25 | return IncreaseParentSize(cursor_ - original_cursor); |
Alexei Frolov | bbf164c | 2019-12-16 12:51:59 -0800 | [diff] [blame] | 26 | } |
| 27 | |
| 28 | // Encodes a base-128 varint to the buffer. |
Rob Mohr | 85cf08f | 2019-12-20 11:54:08 -0800 | [diff] [blame] | 29 | Status Encoder::WriteVarint(uint64_t value) { |
Alexei Frolov | bbf164c | 2019-12-16 12:51:59 -0800 | [diff] [blame] | 30 | if (!encode_status_.ok()) { |
| 31 | return encode_status_; |
| 32 | } |
| 33 | |
Wyatt Hepler | e2cbadf | 2020-06-22 11:21:45 -0700 | [diff] [blame] | 34 | std::span varint_buf = buffer_.last(RemainingSize()); |
Alexei Frolov | bbf164c | 2019-12-16 12:51:59 -0800 | [diff] [blame] | 35 | if (varint_buf.empty()) { |
Wyatt Hepler | d78f7c6 | 2020-09-28 14:27:32 -0700 | [diff] [blame] | 36 | encode_status_ = Status::ResourceExhausted(); |
Alexei Frolov | bbf164c | 2019-12-16 12:51:59 -0800 | [diff] [blame] | 37 | return encode_status_; |
| 38 | } |
| 39 | |
Rob Mohr | 85cf08f | 2019-12-20 11:54:08 -0800 | [diff] [blame] | 40 | size_t written = pw::varint::EncodeLittleEndianBase128(value, varint_buf); |
Alexei Frolov | bbf164c | 2019-12-16 12:51:59 -0800 | [diff] [blame] | 41 | if (written == 0) { |
Wyatt Hepler | d78f7c6 | 2020-09-28 14:27:32 -0700 | [diff] [blame] | 42 | encode_status_ = Status::ResourceExhausted(); |
Alexei Frolov | bbf164c | 2019-12-16 12:51:59 -0800 | [diff] [blame] | 43 | return encode_status_; |
| 44 | } |
| 45 | |
| 46 | cursor_ += written; |
Wyatt Hepler | 1b3da3a | 2021-01-07 13:26:57 -0800 | [diff] [blame] | 47 | return OkStatus(); |
Alexei Frolov | bbf164c | 2019-12-16 12:51:59 -0800 | [diff] [blame] | 48 | } |
| 49 | |
| 50 | Status Encoder::WriteRawBytes(const std::byte* ptr, size_t size) { |
| 51 | if (!encode_status_.ok()) { |
| 52 | return encode_status_; |
| 53 | } |
| 54 | |
| 55 | if (size > RemainingSize()) { |
Wyatt Hepler | d78f7c6 | 2020-09-28 14:27:32 -0700 | [diff] [blame] | 56 | encode_status_ = Status::ResourceExhausted(); |
Alexei Frolov | bbf164c | 2019-12-16 12:51:59 -0800 | [diff] [blame] | 57 | return encode_status_; |
| 58 | } |
| 59 | |
Alexei Frolov | 33a1e8f | 2020-05-26 08:39:32 -0700 | [diff] [blame] | 60 | // Memmove the value into place as it's possible that it shares the encode |
| 61 | // buffer on a memory-constrained system. |
| 62 | std::memmove(cursor_, ptr, size); |
| 63 | |
Alexei Frolov | bbf164c | 2019-12-16 12:51:59 -0800 | [diff] [blame] | 64 | cursor_ += size; |
Wyatt Hepler | 1b3da3a | 2021-01-07 13:26:57 -0800 | [diff] [blame] | 65 | return OkStatus(); |
Alexei Frolov | bbf164c | 2019-12-16 12:51:59 -0800 | [diff] [blame] | 66 | } |
| 67 | |
| 68 | Status Encoder::Push(uint32_t field_number) { |
| 69 | if (!encode_status_.ok()) { |
| 70 | return encode_status_; |
| 71 | } |
| 72 | |
| 73 | if (blob_count_ == blob_locations_.size() || depth_ == blob_stack_.size()) { |
Wyatt Hepler | d78f7c6 | 2020-09-28 14:27:32 -0700 | [diff] [blame] | 74 | encode_status_ = Status::ResourceExhausted(); |
Alexei Frolov | bbf164c | 2019-12-16 12:51:59 -0800 | [diff] [blame] | 75 | return encode_status_; |
| 76 | } |
| 77 | |
| 78 | // Write the key for the nested field. |
Alexei Frolov | 311c620 | 2020-01-06 11:16:20 -0800 | [diff] [blame] | 79 | std::byte* original_cursor = cursor_; |
Alexei Frolov | bbf164c | 2019-12-16 12:51:59 -0800 | [diff] [blame] | 80 | if (Status status = WriteFieldKey(field_number, WireType::kDelimited); |
| 81 | !status.ok()) { |
| 82 | encode_status_ = status; |
| 83 | return status; |
| 84 | } |
| 85 | |
| 86 | if (sizeof(SizeType) > RemainingSize()) { |
| 87 | // Rollback if there isn't enough space. |
| 88 | cursor_ = original_cursor; |
Wyatt Hepler | d78f7c6 | 2020-09-28 14:27:32 -0700 | [diff] [blame] | 89 | encode_status_ = Status::ResourceExhausted(); |
Alexei Frolov | bbf164c | 2019-12-16 12:51:59 -0800 | [diff] [blame] | 90 | return encode_status_; |
| 91 | } |
| 92 | |
| 93 | // Update parent size with the written key. |
Alexei Frolov | f9ae189 | 2021-04-01 18:24:27 -0700 | [diff] [blame] | 94 | PW_TRY(IncreaseParentSize(cursor_ - original_cursor)); |
Alexei Frolov | bbf164c | 2019-12-16 12:51:59 -0800 | [diff] [blame] | 95 | |
| 96 | union { |
Alexei Frolov | 311c620 | 2020-01-06 11:16:20 -0800 | [diff] [blame] | 97 | std::byte* cursor; |
Alexei Frolov | bbf164c | 2019-12-16 12:51:59 -0800 | [diff] [blame] | 98 | SizeType* size_cursor; |
| 99 | }; |
| 100 | |
| 101 | // Create a size entry for the new blob and append it to both the nesting |
| 102 | // stack and location list. |
| 103 | cursor = cursor_; |
| 104 | *size_cursor = 0; |
| 105 | blob_locations_[blob_count_++] = size_cursor; |
| 106 | blob_stack_[depth_++] = size_cursor; |
| 107 | |
| 108 | cursor_ += sizeof(*size_cursor); |
Wyatt Hepler | 1b3da3a | 2021-01-07 13:26:57 -0800 | [diff] [blame] | 109 | return OkStatus(); |
Alexei Frolov | bbf164c | 2019-12-16 12:51:59 -0800 | [diff] [blame] | 110 | } |
| 111 | |
| 112 | Status Encoder::Pop() { |
| 113 | if (!encode_status_.ok()) { |
| 114 | return encode_status_; |
| 115 | } |
| 116 | |
| 117 | if (depth_ == 0) { |
Wyatt Hepler | d78f7c6 | 2020-09-28 14:27:32 -0700 | [diff] [blame] | 118 | encode_status_ = Status::FailedPrecondition(); |
Alexei Frolov | bbf164c | 2019-12-16 12:51:59 -0800 | [diff] [blame] | 119 | return encode_status_; |
| 120 | } |
| 121 | |
| 122 | // Update the parent's size with how much total space the child will take |
| 123 | // after its size field is varint encoded. |
| 124 | SizeType child_size = *blob_stack_[--depth_]; |
Alexei Frolov | f9ae189 | 2021-04-01 18:24:27 -0700 | [diff] [blame] | 125 | PW_TRY(IncreaseParentSize(child_size + VarintSizeBytes(child_size))); |
Alexei Frolov | bbf164c | 2019-12-16 12:51:59 -0800 | [diff] [blame] | 126 | |
Paul Mathieu | a762d17 | 2020-12-01 16:52:09 -0800 | [diff] [blame] | 127 | // Encode the child |
| 128 | if (Status status = EncodeFrom(blob_count_ - 1).status(); !status.ok()) { |
| 129 | encode_status_ = status; |
| 130 | return encode_status_; |
| 131 | } |
| 132 | blob_count_--; |
| 133 | |
Wyatt Hepler | 1b3da3a | 2021-01-07 13:26:57 -0800 | [diff] [blame] | 134 | return OkStatus(); |
Alexei Frolov | bbf164c | 2019-12-16 12:51:59 -0800 | [diff] [blame] | 135 | } |
| 136 | |
Paul Mathieu | a762d17 | 2020-12-01 16:52:09 -0800 | [diff] [blame] | 137 | Result<ConstByteSpan> Encoder::Encode() { return EncodeFrom(0); } |
| 138 | |
| 139 | Result<ConstByteSpan> Encoder::EncodeFrom(size_t blob) { |
Alexei Frolov | bbf164c | 2019-12-16 12:51:59 -0800 | [diff] [blame] | 140 | if (!encode_status_.ok()) { |
Alexei Frolov | bbf164c | 2019-12-16 12:51:59 -0800 | [diff] [blame] | 141 | return encode_status_; |
| 142 | } |
| 143 | |
Paul Mathieu | a762d17 | 2020-12-01 16:52:09 -0800 | [diff] [blame] | 144 | if (blob >= blob_count_) { |
Alexei Frolov | bbf164c | 2019-12-16 12:51:59 -0800 | [diff] [blame] | 145 | // If there are no nested blobs, the buffer already contains a valid proto. |
Alexei Frolov | 13c7c4f | 2020-10-02 09:41:01 -0700 | [diff] [blame] | 146 | return Result<ConstByteSpan>(buffer_.first(EncodedSize())); |
Alexei Frolov | bbf164c | 2019-12-16 12:51:59 -0800 | [diff] [blame] | 147 | } |
| 148 | |
| 149 | union { |
Alexei Frolov | 311c620 | 2020-01-06 11:16:20 -0800 | [diff] [blame] | 150 | std::byte* read_cursor; |
Alexei Frolov | bbf164c | 2019-12-16 12:51:59 -0800 | [diff] [blame] | 151 | SizeType* size_cursor; |
| 152 | }; |
| 153 | |
| 154 | // Starting from the first blob, encode each size field as a varint and |
| 155 | // shift all subsequent data downwards. |
Alexei Frolov | bbf164c | 2019-12-16 12:51:59 -0800 | [diff] [blame] | 156 | size_cursor = blob_locations_[blob]; |
Alexei Frolov | 311c620 | 2020-01-06 11:16:20 -0800 | [diff] [blame] | 157 | std::byte* write_cursor = read_cursor; |
Alexei Frolov | bbf164c | 2019-12-16 12:51:59 -0800 | [diff] [blame] | 158 | |
| 159 | while (read_cursor < cursor_) { |
| 160 | SizeType nested_size = *size_cursor; |
| 161 | |
Wyatt Hepler | e2cbadf | 2020-06-22 11:21:45 -0700 | [diff] [blame] | 162 | std::span<std::byte> varint_buf(write_cursor, sizeof(*size_cursor)); |
Alexei Frolov | bbf164c | 2019-12-16 12:51:59 -0800 | [diff] [blame] | 163 | size_t varint_size = |
| 164 | pw::varint::EncodeLittleEndianBase128(nested_size, varint_buf); |
| 165 | |
| 166 | // Place the write cursor after the encoded varint and the read cursor at |
| 167 | // the location of the next proto field. |
| 168 | write_cursor += varint_size; |
| 169 | read_cursor += varint_buf.size(); |
| 170 | |
| 171 | size_t to_copy; |
| 172 | |
| 173 | if (blob == blob_count_ - 1) { |
| 174 | to_copy = cursor_ - read_cursor; |
| 175 | } else { |
Alexei Frolov | 311c620 | 2020-01-06 11:16:20 -0800 | [diff] [blame] | 176 | std::byte* end = reinterpret_cast<std::byte*>(blob_locations_[blob + 1]); |
Alexei Frolov | bbf164c | 2019-12-16 12:51:59 -0800 | [diff] [blame] | 177 | to_copy = end - read_cursor; |
| 178 | } |
| 179 | |
Alexei Frolov | 33a1e8f | 2020-05-26 08:39:32 -0700 | [diff] [blame] | 180 | std::memmove(write_cursor, read_cursor, to_copy); |
Alexei Frolov | bbf164c | 2019-12-16 12:51:59 -0800 | [diff] [blame] | 181 | write_cursor += to_copy; |
| 182 | read_cursor += to_copy; |
| 183 | |
| 184 | ++blob; |
| 185 | } |
| 186 | |
| 187 | // Point the cursor to the end of the encoded proto. |
| 188 | cursor_ = write_cursor; |
Alexei Frolov | 13c7c4f | 2020-10-02 09:41:01 -0700 | [diff] [blame] | 189 | return Result<ConstByteSpan>(buffer_.first(EncodedSize())); |
Alexei Frolov | bbf164c | 2019-12-16 12:51:59 -0800 | [diff] [blame] | 190 | } |
| 191 | |
Alexei Frolov | f9ae189 | 2021-04-01 18:24:27 -0700 | [diff] [blame] | 192 | Status Encoder::IncreaseParentSize(size_t size_bytes) { |
| 193 | if (!encode_status_.ok()) { |
| 194 | return encode_status_; |
| 195 | } |
| 196 | |
| 197 | if (depth_ == 0) { |
| 198 | return OkStatus(); |
| 199 | } |
| 200 | |
| 201 | size_t current_size = *blob_stack_[depth_ - 1]; |
| 202 | |
| 203 | constexpr size_t max_size = |
| 204 | std::min(varint::MaxValueInBytes(sizeof(SizeType)), |
| 205 | static_cast<uint64_t>(std::numeric_limits<uint32_t>::max())); |
| 206 | |
| 207 | if (size_bytes > max_size || current_size > max_size - size_bytes) { |
| 208 | encode_status_ = Status::OutOfRange(); |
| 209 | return encode_status_; |
| 210 | } |
| 211 | |
| 212 | *blob_stack_[depth_ - 1] = current_size + size_bytes; |
| 213 | return OkStatus(); |
| 214 | } |
| 215 | |
Alexei Frolov | bbf164c | 2019-12-16 12:51:59 -0800 | [diff] [blame] | 216 | } // namespace pw::protobuf |