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