1 /* 2 * Copyright (c) Facebook, Inc. and its affiliates. 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 17 #pragma once 18 19 #include <cstdint> 20 21 #include <folly/io/Cursor.h> 22 #include <folly/io/IOBuf.h> 23 #include <folly/io/IOBufQueue.h> 24 25 #include <thrift/lib/cpp/util/VarintUtils.h> 26 #include <thrift/lib/cpp2/protocol/nimble/BufferingNimbleEncoder.h> 27 #include <thrift/lib/cpp2/protocol/nimble/NimbleTypes.h> 28 29 namespace apache { 30 namespace thrift { 31 namespace detail { 32 33 class Encoder { 34 public: Encoder()35 Encoder() : fieldAppender_(&fieldData_, 4096) { 36 sizeStream_.setControlOutput(&sizeControl_); 37 sizeStream_.setDataOutput(&sizeData_); 38 contentStream_.setControlOutput(&contentControl_); 39 contentStream_.setDataOutput(&contentData_); 40 } 41 ~Encoder() = default; 42 Encoder(const Encoder&) = delete; 43 Encoder& operator=(const Encoder&) = delete; 44 encodeSizeChunk(std::uint32_t chunk)45 void encodeSizeChunk(std::uint32_t chunk) { sizeStream_.encodeChunk(chunk); } 46 encodeContentChunk(std::uint32_t chunk)47 void encodeContentChunk(std::uint32_t chunk) { 48 contentStream_.encodeChunk(chunk); 49 } 50 51 // Eventually, this should probably have an IOBuf-taking variant, too. encodeBinary(const void * buf,std::size_t size)52 void encodeBinary(const void* buf, std::size_t size) { 53 binaryData_.append(buf, size); 54 } 55 finalize()56 std::unique_ptr<folly::IOBuf> finalize() { 57 sizeStream_.finalize(); 58 contentStream_.finalize(); 59 60 // An enum field (with yet-to-be determined contents). 61 std::uint64_t messageFlags = 0; 62 63 auto fieldData = fieldData_.move(); 64 auto sizeControl = sizeControl_.move(); 65 auto sizeData = sizeData_.move(); 66 auto contentControl = contentControl_.move(); 67 auto contentData = contentData_.move(); 68 auto binaryData = binaryData_.move(); 69 70 auto chainLengthBytes = [](auto& iobufptr) { 71 // IOBufQueue::move() can return null if no data was ever written to it. 72 return iobufptr ? iobufptr->computeChainDataLength() : 0; 73 }; 74 75 auto sizeHeader = folly::IOBuf::create(0); 76 folly::io::Appender writer(sizeHeader.get(), 32); 77 apache::thrift::util::writeVarint(writer, messageFlags); 78 apache::thrift::util::writeVarint(writer, chainLengthBytes(fieldData)); 79 apache::thrift::util::writeVarint(writer, chainLengthBytes(sizeControl)); 80 apache::thrift::util::writeVarint(writer, chainLengthBytes(sizeData)); 81 apache::thrift::util::writeVarint(writer, chainLengthBytes(contentControl)); 82 apache::thrift::util::writeVarint(writer, chainLengthBytes(contentData)); 83 apache::thrift::util::writeVarint(writer, chainLengthBytes(binaryData)); 84 85 auto prependToSizeHeader = [&](auto ioBufUPtr) { 86 if (ioBufUPtr != nullptr) { 87 sizeHeader->prependChain(std::move(ioBufUPtr)); 88 } 89 }; 90 91 prependToSizeHeader(std::move(fieldData)); 92 prependToSizeHeader(std::move(sizeControl)); 93 prependToSizeHeader(std::move(sizeData)); 94 prependToSizeHeader(std::move(contentControl)); 95 prependToSizeHeader(std::move(contentData)); 96 prependToSizeHeader(std::move(binaryData)); 97 98 return sizeHeader; 99 } 100 encodeFieldBytes(nimble::FieldBytes bytes)101 void encodeFieldBytes(nimble::FieldBytes bytes) { 102 fieldAppender_.push(bytes.bytes, bytes.len); 103 } 104 105 private: 106 folly::IOBufQueue fieldData_; 107 folly::io::QueueAppender fieldAppender_; 108 109 folly::IOBufQueue sizeControl_; 110 folly::IOBufQueue sizeData_; 111 // We expect size stream data to be exclusively positives, and so don't 112 // bother zigzagging it. 113 BufferingNimbleEncoder<ChunkRepr::kRaw> sizeStream_; 114 115 folly::IOBufQueue contentControl_; 116 folly::IOBufQueue contentData_; 117 // Content stream data, on the other hand, may contain negatives. 118 BufferingNimbleEncoder<ChunkRepr::kZigzag> contentStream_; 119 120 // String and binary data is encoded as raw bytes. 121 folly::IOBufQueue binaryData_; 122 }; 123 124 } // namespace detail 125 } // namespace thrift 126 } // namespace apache 127